From 1d3c100acb6fcb1a519f7655131cb2a643ce6d2f Mon Sep 17 00:00:00 2001 From: Jack Date: Sun, 5 Jul 2026 20:45:35 +0800 Subject: [PATCH] Refactor: pdf parser (#16625) ### Summary PDF parser refactor --- internal/deepdoc/parser/pdf/layout/layout.go | 461 ++++++---- .../deepdoc/parser/pdf/layout/layout_test.go | 435 +++++++++ internal/deepdoc/parser/pdf/parser.go | 381 +++++--- internal/deepdoc/parser/pdf/parser_ocr.go | 63 +- internal/deepdoc/parser/pdf/parser_test.go | 212 ++++- .../deepdoc/parser/pdf/table/group_boxes.go | 371 ++++++++ .../parser/pdf/table/group_boxes_test.go | 350 ++++++++ .../parser/pdf/table/table_construct.go | 850 +++--------------- .../parser/pdf/table/table_construct_test.go | 669 +++++--------- .../deepdoc/parser/pdf/table/table_html.go | 112 +++ .../parser/pdf/table/table_html_test.go | 108 +++ .../deepdoc/parser/pdf/table/table_merge.go | 128 +++ .../parser/pdf/table/table_merge_test.go | 179 ++++ .../deepdoc/parser/pdf/table/table_post.go | 489 ++++++---- .../parser/pdf/table/table_post_test.go | 802 +++++++++++++++++ .../deepdoc/parser/pdf/table/table_spans.go | 90 ++ .../parser/pdf/table/table_spans_test.go | 46 + 17 files changed, 3988 insertions(+), 1758 deletions(-) create mode 100644 internal/deepdoc/parser/pdf/table/group_boxes.go create mode 100644 internal/deepdoc/parser/pdf/table/group_boxes_test.go create mode 100644 internal/deepdoc/parser/pdf/table/table_html.go create mode 100644 internal/deepdoc/parser/pdf/table/table_html_test.go create mode 100644 internal/deepdoc/parser/pdf/table/table_merge.go create mode 100644 internal/deepdoc/parser/pdf/table/table_merge_test.go create mode 100644 internal/deepdoc/parser/pdf/table/table_post_test.go create mode 100644 internal/deepdoc/parser/pdf/table/table_spans.go create mode 100644 internal/deepdoc/parser/pdf/table/table_spans_test.go diff --git a/internal/deepdoc/parser/pdf/layout/layout.go b/internal/deepdoc/parser/pdf/layout/layout.go index 1bf08d614b..142e0cdaaa 100644 --- a/internal/deepdoc/parser/pdf/layout/layout.go +++ b/internal/deepdoc/parser/pdf/layout/layout.go @@ -24,111 +24,139 @@ func AssignColumn(boxes []pdf.TextBox, zoom float64) []pdf.TextBox { return boxes } - pageGroups := make(map[int][]int) - for i, b := range boxes { - pageGroups[b.PageNumber] = append(pageGroups[b.PageNumber], i) - } + pageGroups, sortedPages := groupBoxesByPage(boxes) result := make([]pdf.TextBox, len(boxes)) copy(result, boxes) // Step A: per-page best k using silhouette score. pageCols := make(map[int]int) - for pg, indices := range pageGroups { - n := len(indices) - if n < 2 { - pageCols[pg] = 1 - for _, idx := range indices { - result[idx].ColID = 0 - } - continue - } - - // Extract x0 values and apply indent tolerance (12% of page width). - x0s := make([]float64, n) - minX0 := math.MaxFloat64 - maxX1 := 0.0 - for i, idx := range indices { - x0s[i] = boxes[idx].X0 - if x0s[i] < minX0 { - minX0 = x0s[i] - } - if boxes[idx].X1 > maxX1 { - maxX1 = boxes[idx].X1 - } - } - pageWidth := maxX1 - minX0 - indentTol := pageWidth * 0.12 - - for i := range x0s { - if math.Abs(x0s[i]-minX0) < indentTol { - x0s[i] = minX0 - } - } - - // Try k = 1 .. min(4, n), pick best by silhouette. - maxTry := min(4, n) - if maxTry < 2 { - maxTry = 1 - } - bestK, bestScore := 1, -1.0 - - for k := 1; k <= maxTry; k++ { - labels, _ := util.KMeans1D(x0s, k) - var score float64 - if k > 1 { - score = util.Silhouette1D(x0s, labels) - } - // score = 0 for k=1; score = -1 if silhouette undefined. - if score > bestScore { - bestScore = score - bestK = k - } - } - pageCols[pg] = bestK + for _, pg := range sortedPages { + indices := pageGroups[pg] + determineBestKForPage(boxes, result, indices, pg, pageCols) } // Step B: assign col_id per page using per-page best k. // Labels are remapped by centroid x-order: leftmost column → 0. - for pg, indices := range pageGroups { - if len(indices) == 0 { - continue - } - k := pageCols[pg] - if len(indices) < k { - k = 1 - } - - x0s := make([]float64, len(indices)) - for i, idx := range indices { - x0s[i] = boxes[idx].X0 - } - - labels, centroids := util.KMeans1D(x0s, k) - - // Sort centroids by x position, remap labels left→right. - type clPair struct { - center float64 - label int - } - var pairs []clPair - for lbl, c := range centroids { - pairs = append(pairs, clPair{c, lbl}) - } - sort.Slice(pairs, func(i, j int) bool { return pairs[i].center < pairs[j].center }) - remap := make(map[int]int, k) - for newL, p := range pairs { - remap[p.label] = newL - } - - for i, idx := range indices { - result[idx].ColID = remap[labels[i]] - } + for _, pg := range sortedPages { + indices := pageGroups[pg] + assignColIDsForPage(boxes, result, indices, pg, pageCols) } return result } +// determineBestKForPage finds the best number of clusters (k) for a page using silhouette score +func determineBestKForPage(boxes, result []pdf.TextBox, indices []int, pg int, pageCols map[int]int) { + n := len(indices) + if n < 2 { + pageCols[pg] = 1 + for _, idx := range indices { + result[idx].ColID = 0 + } + return + } + + x0s, minX0, maxX1 := extractX0Values(boxes, indices) + pageWidth := maxX1 - minX0 + indentTol := pageWidth * 0.12 + applyIndentTolerance(x0s, minX0, indentTol) + + bestK, _ := findBestK(x0s, n) + pageCols[pg] = bestK +} + +// extractX0Values extracts x0 coordinates from boxes on a page and finds minX0 and maxX1 +func extractX0Values(boxes []pdf.TextBox, indices []int) (x0s []float64, minX0 float64, maxX1 float64) { + n := len(indices) + x0s = make([]float64, n) + minX0 = math.MaxFloat64 + maxX1 = 0.0 + for i, idx := range indices { + x0s[i] = boxes[idx].X0 + if x0s[i] < minX0 { + minX0 = x0s[i] + } + if boxes[idx].X1 > maxX1 { + maxX1 = boxes[idx].X1 + } + } + return x0s, minX0, maxX1 +} + +// applyIndentTolerance adjusts x0 values that are close to minX0 to improve clustering +func applyIndentTolerance(x0s []float64, minX0, indentTol float64) { + for i := range x0s { + if math.Abs(x0s[i]-minX0) < indentTol { + x0s[i] = minX0 + } + } +} + +// findBestK tries k from 1 to min(4, n) and returns the k with the best silhouette score +func findBestK(x0s []float64, n int) (bestK int, bestScore float64) { + maxTry := min(4, n) + if maxTry < 2 { + maxTry = 1 + } + bestK, bestScore = 1, -1.0 + + for k := 1; k <= maxTry; k++ { + labels, _ := util.KMeans1D(x0s, k) + var score float64 + if k > 1 { + score = util.Silhouette1D(x0s, labels) + } + // score = 0 for k=1; score = -1 if silhouette undefined. + if score > bestScore { + bestScore = score + bestK = k + } + } + return bestK, bestScore +} + +// assignColIDsForPage assigns column IDs to boxes on a page using the best k +func assignColIDsForPage(boxes, result []pdf.TextBox, indices []int, pg int, pageCols map[int]int) { + if len(indices) == 0 { + return + } + k := pageCols[pg] + if len(indices) < k { + k = 1 + } + + x0s := make([]float64, len(indices)) + for i, idx := range indices { + x0s[i] = boxes[idx].X0 + } + + labels, centroids := util.KMeans1D(x0s, k) + remap := remapLabelsByCentroidOrder(centroids) + + for i, idx := range indices { + result[idx].ColID = remap[labels[i]] + } +} + +// remapLabelsByCentroidOrder remaps cluster labels so leftmost column = 0 +func remapLabelsByCentroidOrder(centroids []float64) map[int]int { + type clPair struct { + center float64 + label int + } + var pairs []clPair + for lbl, c := range centroids { + pairs = append(pairs, clPair{c, lbl}) + } + sort.Slice(pairs, func(i, j int) bool { return pairs[i].center < pairs[j].center }) + remap := make(map[int]int, len(centroids)) + for newL, p := range pairs { + remap[p.label] = newL + } + return remap +} + // ---- Text merge (horizontal) ---- // TextMerge horizontally merges adjacent boxes at similar vertical positions. @@ -183,31 +211,14 @@ func NaiveVerticalMerge(boxes []pdf.TextBox, medianHeights map[int]float64, medi if len(boxes) < 2 { return boxes } - // Group by page only — matches Python's _naive_vertical_merge which - // hardcodes col="x" (pdf_parser.py:868), ignoring column assignment. - // Cross-column merges are prevented by the 30% horizontal overlap check. - groups := make(map[int][]int) - for i, b := range boxes { - groups[b.PageNumber] = append(groups[b.PageNumber], i) - } - // Sort page keys for deterministic output order (Python dict preserves - // insertion order since 3.7, Go map iteration is random). - pageKeys := make([]int, 0, len(groups)) - for pg := range groups { - pageKeys = append(pageKeys, pg) - } - sort.Ints(pageKeys) + + // Group boxes by page + pageGroups, sortedPages := groupBoxesByPage(boxes) var result []pdf.TextBox - for _, pg := range pageKeys { - indices := groups[pg] - sort.Slice(indices, func(i, j int) bool { - bi, bj := boxes[indices[i]], boxes[indices[j]] - if bi.Top != bj.Top { - return bi.Top < bj.Top - } - return bi.X0 < bj.X0 - }) + for _, pg := range sortedPages { + // Collect all boxes for this page + indices := pageGroups[pg] bxs := make([]pdf.TextBox, len(indices)) for i, idx := range indices { bxs[i] = boxes[idx] @@ -222,85 +233,9 @@ func NaiveVerticalMerge(boxes []pdf.TextBox, medianHeights map[int]float64, medi mw = 8 // Python fallback: np.median([...]) if chars else 8 (pdf_parser.py:1465) } - // Collect pattern: build output slice, merging into last element when appropriate. - out := make([]pdf.TextBox, 0, len(bxs)) - for i := 0; i < len(bxs); i++ { - b := bxs[i] - // Cross-page suffix (e.g. page number on previous page): skip. - if i > 0 && bxs[i-1].PageNumber < b.PageNumber && pageNumSuffixPattern.MatchString(bxs[i-1].Text) { - continue - } - if strings.TrimSpace(b.Text) == "" { - // Whitespace gap bridge: absorb into prev box if gap/xov pass, - // extending prev.Bottom. This matches Python's while/pop which - // keeps whitespace inline and lets it extend the previous box. - if len(out) > 0 { - prev := &out[len(out)-1] - if b.Top-prev.Bottom <= mh*1.5 && util.OverlapX(prev, &b) >= 0.3 { - // TODO: prev.Bottom = math.Max(prev.Bottom, b.Bottom) — direct assignment - // can shrink a tall merged box when a short whitespace box overlaps. - // Matches Python behavior (also direct assignment). Defer fix until - // pipeline alignment is shipped. See TestNaiveVerticalMerge_BottomShrink. - prev.Bottom = b.Bottom - } - } - continue - } - if len(out) == 0 { - out = append(out, b) - continue - } - prev := &out[len(out)-1] - if prev.LayoutNo != b.LayoutNo || strings.TrimSpace(b.Text) == "" { - slog.Debug("vm reject", "reason", "layout_no", "prevLayout", prev.LayoutNo, "bLayout", b.LayoutNo) - out = append(out, b) - continue - } - gap := b.Top - prev.Bottom - if gap > mh*1.5 { - slog.Debug("vm reject", "reason", "gap", "gap", gap, "threshold", mh*1.5, "mh", mh) - out = append(out, b) - continue - } - ov := util.OverlapX(prev, &b) - if ov < 0.3 { - slog.Debug("vm reject", "reason", "ovX", "ov", ov, "threshold", 0.3) - out = append(out, b) - continue - } - - // Strip text before checking first/last characters (matching Python's - // b["text"].strip()[-1] / b_["text"].strip()[0]). - prevText := strings.TrimSpace(prev.Text) - bText := strings.TrimSpace(b.Text) - - concatting := []bool{ - endsWithOneOf(prevText, ",;:\"，、‘“；：-"), - endsSecondLastOneOf(prevText, ",;:\"，、‘“；："), - startsWithOneOf(bText, "。；？！”）),，、："), - } - anti := []bool{ - endsWithOneOf(prevText, "。？！?"), - isEnglish && endsWithOneOf(prevText, ".!?"), - prev.PageNumber == b.PageNumber && b.Top-prev.Bottom > mh*1.5, - prev.PageNumber < b.PageNumber && math.Abs(prev.X0-b.X0) > mw*4, - } - detach := []bool{prev.X1 < b.X0, prev.X0 > b.X1} - if (slices.Contains(anti, true) && !slices.Contains(concatting, true)) || slices.Contains(detach, true) { - out = append(out, b) - continue - } - - slog.Debug("vm merge", "gap", gap, "ovX", ov, "mh", mh, "prev", prevText[:min(40, len(prevText))], "next", bText[:min(40, len(bText))]) - // Python: (b["text"].rstrip() + " " + b_["text"].lstrip()).strip() - prev.Text = strings.TrimSpace(strings.TrimRight(prevText, " \t") + " " + strings.TrimLeft(bText, " \t")) - // Preserve the taller bottom when merging (prev.Bottom may already - // extend beyond b.Bottom from a previous merge step). - prev.Bottom = math.Max(prev.Bottom, b.Bottom) - prev.X0 = math.Min(prev.X0, b.X0) - prev.X1 = math.Max(prev.X1, b.X1) - } - result = append(result, out...) + // Process boxes for this page + processed := processPageBoxes(bxs, mh, mw, isEnglish) + result = append(result, processed...) } slog.Debug("vm result", "in", len(boxes), "out", len(result)) return result @@ -333,6 +268,148 @@ func FinalReadingOrderMerge(boxes []pdf.TextBox) []pdf.TextBox { var pageNumSuffixPattern = regexp.MustCompile(`[0-9 •一—-]+$`) +// groupBoxesByPage groups text boxes by page, returning a map from page number to index list and sorted page number list +func groupBoxesByPage(boxes []pdf.TextBox) (map[int][]int, []int) { + if len(boxes) == 0 { + return map[int][]int{}, []int{} + } + + pageGroups := make(map[int][]int) + for i, b := range boxes { + pageGroups[b.PageNumber] = append(pageGroups[b.PageNumber], i) + } + + // Sort page numbers + pageKeys := make([]int, 0, len(pageGroups)) + for pg := range pageGroups { + pageKeys = append(pageKeys, pg) + } + sort.Ints(pageKeys) + + return pageGroups, pageKeys +} + +// shouldMergeBoxes determines whether two boxes should be merged +func shouldMergeBoxes(prev, curr *pdf.TextBox, mh, mw float64, isEnglish bool) bool { + // Check layout number + if prev.LayoutNo != curr.LayoutNo { + slog.Debug("vm reject", "reason", "layoutNo", "prevLayout", prev.LayoutNo, "currLayout", curr.LayoutNo) + return false + } + + // Check vertical gap + gap := curr.Top - prev.Bottom + if gap > mh*1.5 { + slog.Debug("vm reject", "reason", "gap", "gap", gap, "threshold", mh*1.5, "mh", mh) + return false + } + + // Check horizontal overlap + ov := util.OverlapX(prev, curr) + if ov < 0.3 { + slog.Debug("vm reject", "reason", "ovX", "ov", ov, "threshold", 0.3) + return false + } + + // Check merge/block conditions + prevText := strings.TrimSpace(prev.Text) + currText := strings.TrimSpace(curr.Text) + + concatting := []bool{ + endsWithOneOf(prevText, ",;:\"，、‘“；：-"), + endsSecondLastOneOf(prevText, ",;:\"，、‘“；："), + startsWithOneOf(currText, "。；？！?\"）),，、："), + } + anti := []bool{ + endsWithOneOf(prevText, "。？！?"), + isEnglish && endsWithOneOf(prevText, ".!?"), + prev.PageNumber < curr.PageNumber && math.Abs(prev.X0-curr.X0) > mw*4, + } + detach := []bool{prev.X1 < curr.X0, prev.X0 > curr.X1} + + if (slices.Contains(anti, true) && !slices.Contains(concatting, true)) || slices.Contains(detach, true) { + return false + } + + return true +} + +// mergeTwoBoxes merges two text boxes +func mergeTwoBoxes(prev, curr pdf.TextBox) pdf.TextBox { + prevText := strings.TrimSpace(prev.Text) + currText := strings.TrimSpace(curr.Text) + + prev.Text = strings.TrimSpace(strings.TrimRight(prevText, " \t") + " " + strings.TrimLeft(currText, " \t")) + prev.Bottom = math.Max(prev.Bottom, curr.Bottom) + prev.X0 = math.Min(prev.X0, curr.X0) + prev.X1 = math.Max(prev.X1, curr.X1) + + prevTrunc, currTrunc := prevText, currText + if r := []rune(prevTrunc); len(r) > 40 { + prevTrunc = string(r[:40]) + } + if r := []rune(currTrunc); len(r) > 40 { + currTrunc = string(r[:40]) + } + slog.Debug("vm merge", "prev", prevTrunc, "curr", currTrunc) + + return prev +} + +// processPageBoxes processes all boxes for a single page +func processPageBoxes(boxes []pdf.TextBox, mh, mw float64, isEnglish bool) []pdf.TextBox { + if len(boxes) == 0 { + return boxes + } + + // Sort by Top, X0 + sortedBoxes := make([]pdf.TextBox, len(boxes)) + copy(sortedBoxes, boxes) + sort.Slice(sortedBoxes, func(i, j int) bool { + if sortedBoxes[i].Top != sortedBoxes[j].Top { + return sortedBoxes[i].Top < sortedBoxes[j].Top + } + return sortedBoxes[i].X0 < sortedBoxes[j].X0 + }) + + out := make([]pdf.TextBox, 0, len(sortedBoxes)) + for i := 0; i < len(sortedBoxes); i++ { + curr := sortedBoxes[i] + + // Skip cross-page suffixes (like previous page number) + if i > 0 && sortedBoxes[i-1].PageNumber < curr.PageNumber && pageNumSuffixPattern.MatchString(sortedBoxes[i-1].Text) { + continue + } + + // Handle empty boxes + if strings.TrimSpace(curr.Text) == "" { + if len(out) > 0 { + prev := &out[len(out)-1] + if curr.Top-prev.Bottom <= mh*1.5 && util.OverlapX(prev, &curr) >= 0.3 { + // TODO: prev.Bottom = math.Max(prev.Bottom, curr.Bottom) — direct assignment might shrink tall merged boxes + // Matches Python behavior (also direct assignment). Defer fix until pipeline alignment release. + prev.Bottom = curr.Bottom + } + } + continue + } + + if len(out) == 0 { + out = append(out, curr) + continue + } + + prev := &out[len(out)-1] + if shouldMergeBoxes(prev, &curr, mh, mw, isEnglish) { + out[len(out)-1] = mergeTwoBoxes(*prev, curr) + } else { + out = append(out, curr) + } + } + + return out +} + // ---- rune-based text helpers (CJK-safe) ---- func lastRune(s string) rune { diff --git a/internal/deepdoc/parser/pdf/layout/layout_test.go b/internal/deepdoc/parser/pdf/layout/layout_test.go index 019bdfad86..c7ed9841e3 100644 --- a/internal/deepdoc/parser/pdf/layout/layout_test.go +++ b/internal/deepdoc/parser/pdf/layout/layout_test.go @@ -7,6 +7,19 @@ import ( "testing" ) +// ---- test helpers ---- + +func newTestTextBox(page int, x0, x1, top, bottom float64, text string) pdf.TextBox { + return pdf.TextBox{ + PageNumber: page, + X0: x0, + X1: x1, + Top: top, + Bottom: bottom, + Text: text, + } +} + func TestAssignColumn(t *testing.T) { boxes := []pdf.TextBox{ {PageNumber: 0, X0: 50, Text: "col0-left"}, @@ -505,3 +518,425 @@ func TestNaiveVerticalMergeNonMerge(t *testing.T) { t.Errorf("expected 2 separate boxes (large gap), got %d", len(result)) } } + +// ---- 重构辅助函数的测试 ---- + +func TestGroupBoxesByPage(t *testing.T) { + boxes := []pdf.TextBox{ + {PageNumber: 1, Text: "page1-box1"}, + {PageNumber: 0, Text: "page0-box1"}, + {PageNumber: 0, Text: "page0-box2"}, + {PageNumber: 2, Text: "page2-box1"}, + {PageNumber: 1, Text: "page1-box2"}, + } + pageGroups, sortedPages := groupBoxesByPage(boxes) + + if len(sortedPages) != 3 { + t.Errorf("expected 3 unique pages, got %d", len(sortedPages)) + } + if sortedPages[0] != 0 || sortedPages[1] != 1 || sortedPages[2] != 2 { + t.Errorf("pages should be sorted [0,1,2], got %v", sortedPages) + } + if len(pageGroups[0]) != 2 { + t.Errorf("page 0 should have 2 boxes, got %d", len(pageGroups[0])) + } + if len(pageGroups[1]) != 2 { + t.Errorf("page 1 should have 2 boxes, got %d", len(pageGroups[1])) + } + if len(pageGroups[2]) != 1 { + t.Errorf("page 2 should have 1 box, got %d", len(pageGroups[2])) + } + if boxes[pageGroups[0][0]].Text != "page0-box1" { + t.Errorf("first page0 box index incorrect") + } +} + +func TestGroupBoxesByPage_Empty(t *testing.T) { + pageGroups, sortedPages := groupBoxesByPage(nil) + if len(pageGroups) != 0 || len(sortedPages) != 0 { + t.Error("empty input should return empty result") + } + + pageGroups, sortedPages = groupBoxesByPage([]pdf.TextBox{}) + if len(pageGroups) != 0 || len(sortedPages) != 0 { + t.Error("empty input should return empty result") + } +} + +func TestGroupBoxesByPage_SinglePage(t *testing.T) { + boxes := []pdf.TextBox{ + {PageNumber: 5, Text: "box1"}, + {PageNumber: 5, Text: "box2"}, + } + pageGroups, sortedPages := groupBoxesByPage(boxes) + if len(sortedPages) != 1 || sortedPages[0] != 5 { + t.Errorf("expected single page 5, got %v", sortedPages) + } + if len(pageGroups[5]) != 2 { + t.Errorf("page 5 should have 2 boxes, got %d", len(pageGroups[5])) + } +} + +func TestShouldMergeBoxes(t *testing.T) { + t.Run("should merge - basic case", func(t *testing.T) { + prev := &pdf.TextBox{ + PageNumber: 0, X0: 50, X1: 250, Top: 100, Bottom: 112, + Text: "前一句", + LayoutNo: "1", + } + curr := &pdf.TextBox{ + PageNumber: 0, X0: 50, X1: 250, Top: 114, Bottom: 126, + Text: "后一句", + LayoutNo: "1", + } + if !shouldMergeBoxes(prev, curr, 12, 200, false) { + t.Error("should merge basic case") + } + }) + + t.Run("should NOT merge - different layoutNo", func(t *testing.T) { + prev := &pdf.TextBox{PageNumber: 0, LayoutNo: "1", Top: 100, Bottom: 112, X0: 50, X1: 250} + curr := &pdf.TextBox{PageNumber: 0, LayoutNo: "2", Top: 114, Bottom: 126, X0: 50, X1: 250} + if shouldMergeBoxes(prev, curr, 12, 200, false) { + t.Error("should not merge different layoutNo") + } + }) + + t.Run("should NOT merge - gap too large", func(t *testing.T) { + prev := &pdf.TextBox{PageNumber: 0, LayoutNo: "1", Top: 100, Bottom: 112, X0: 50, X1: 250} + curr := &pdf.TextBox{PageNumber: 0, LayoutNo: "1", Top: 200, Bottom: 212, X0: 50, X1: 250} + if shouldMergeBoxes(prev, curr, 12, 200, false) { + t.Error("should not merge large gap") + } + }) + + t.Run("should NOT merge - overlap too small", func(t *testing.T) { + prev := &pdf.TextBox{PageNumber: 0, LayoutNo: "1", Top: 100, Bottom: 112, X0: 50, X1: 100} + curr := &pdf.TextBox{PageNumber: 0, LayoutNo: "1", Top: 114, Bottom: 126, X0: 200, X1: 250} + if shouldMergeBoxes(prev, curr, 12, 200, false) { + t.Error("should not merge small overlap") + } + }) + + t.Run("should merge - comma override period anti", func(t *testing.T) { + prev := &pdf.TextBox{ + PageNumber: 0, LayoutNo: "1", Top: 100, Bottom: 112, X0: 50, X1: 250, + Text: "前一句。", + } + curr := &pdf.TextBox{ + PageNumber: 0, LayoutNo: "1", Top: 114, Bottom: 126, X0: 50, X1: 250, + Text: ", 续句", + } + if !shouldMergeBoxes(prev, curr, 12, 200, false) { + t.Error("should merge when comma overrides period anti") + } + }) + + t.Run("should NOT merge - english period anti", func(t *testing.T) { + prev := &pdf.TextBox{ + PageNumber: 0, LayoutNo: "1", Top: 100, Bottom: 112, X0: 50, X1: 250, + Text: "End of sentence.", + } + curr := &pdf.TextBox{ + PageNumber: 0, LayoutNo: "1", Top: 114, Bottom: 126, X0: 50, X1: 250, + Text: "Next sentence", + } + if shouldMergeBoxes(prev, curr, 12, 200, true) { + t.Error("should not merge english period anti") + } + }) +} + +func TestMergeTwoBoxes(t *testing.T) { + prev := pdf.TextBox{ + PageNumber: 0, X0: 50, X1: 200, Top: 100, Bottom: 112, + Text: "第一行", + LayoutNo: "1", + } + curr := pdf.TextBox{ + PageNumber: 0, X0: 60, X1: 250, Top: 114, Bottom: 130, + Text: "第二行", + LayoutNo: "1", + } + + result := mergeTwoBoxes(prev, curr) + + expectedText := "第一行 第二行" + if result.Text != expectedText { + t.Errorf("expected text %q, got %q", expectedText, result.Text) + } + if result.X0 != 50 { + t.Errorf("expected X0 50, got %f", result.X0) + } + if result.X1 != 250 { + t.Errorf("expected X1 250, got %f", result.X1) + } + if result.Bottom != 130 { + t.Errorf("expected Bottom 130, got %f", result.Bottom) + } + if result.LayoutNo != "1" { + t.Errorf("expected LayoutNo preserved") + } +} + +func TestMergeTwoBoxes_TrimWhitespace(t *testing.T) { + prev := pdf.TextBox{Text: " first line "} + curr := pdf.TextBox{Text: " second line "} + + result := mergeTwoBoxes(prev, curr) + + if result.Text != "first line second line" { + t.Errorf("text should be trimmed and joined, got %q", result.Text) + } +} + +func TestProcessPageBoxes(t *testing.T) { + boxes := []pdf.TextBox{ + { + PageNumber: 0, X0: 50, X1: 250, Top: 114, Bottom: 126, + Text: "第二句", + LayoutNo: "1", + }, + { + PageNumber: 0, X0: 50, X1: 250, Top: 100, Bottom: 112, + Text: "第一句", + LayoutNo: "1", + }, + } + + result := processPageBoxes(boxes, 12, 200, false) + + if len(result) != 1 { + t.Errorf("expected 1 merged box, got %d", len(result)) + } + if !strings.Contains(result[0].Text, "第一句") || !strings.Contains(result[0].Text, "第二句") { + t.Errorf("merged text should contain both parts, got %q", result[0].Text) + } +} + +func TestProcessPageBoxes_WhitespaceBox(t *testing.T) { + boxes := []pdf.TextBox{ + { + PageNumber: 0, X0: 50, X1: 250, Top: 100, Bottom: 112, + Text: "第一句", + LayoutNo: "1", + }, + { + PageNumber: 0, X0: 50, X1: 250, Top: 113, Bottom: 115, + Text: " ", + LayoutNo: "1", + }, + { + PageNumber: 0, X0: 50, X1: 250, Top: 116, Bottom: 128, + Text: "第二句", + LayoutNo: "1", + }, + } + + result := processPageBoxes(boxes, 12, 200, false) + + if len(result) != 1 { + t.Errorf("expected 1 merged box, got %d", len(result)) + } +} + +func TestProcessPageBoxes_NoMerge(t *testing.T) { + boxes := []pdf.TextBox{ + { + PageNumber: 0, X0: 50, X1: 250, Top: 100, Bottom: 112, + Text: "第一句。", + LayoutNo: "1", + }, + { + PageNumber: 0, X0: 50, X1: 250, Top: 200, Bottom: 212, + Text: "第二句", + LayoutNo: "1", + }, + } + + result := processPageBoxes(boxes, 12, 200, false) + + if len(result) != 2 { + t.Errorf("expected 2 boxes, got %d", len(result)) + } +} + +// ── Column-assignment helper tests ────────────────────────────────── + +func TestExtractX0Values(t *testing.T) { + boxes := []pdf.TextBox{ + {PageNumber: 0, X0: 50, X1: 200}, + {PageNumber: 0, X0: 30, X1: 100}, + {PageNumber: 0, X0: 80, X1: 300}, + } + x0s, minX0, maxX1 := extractX0Values(boxes, []int{0, 1, 2}) + if len(x0s) != 3 { + t.Fatalf("expected 3 x0s, got %d", len(x0s)) + } + if x0s[0] != 50 || x0s[1] != 30 || x0s[2] != 80 { + t.Errorf("x0s mismatch: %v", x0s) + } + if minX0 != 30 { + t.Errorf("minX0 = %v, want 30", minX0) + } + if maxX1 != 300 { + t.Errorf("maxX1 = %v, want 300", maxX1) + } +} + +func TestApplyIndentTolerance(t *testing.T) { + values := []float64{100, 105, 200, 210} + applyIndentTolerance(values, 100, 10) + if values[0] != 100 || values[1] != 100 { + t.Errorf("close x0s should be adjusted to minX0: %v", values) + } + if values[2] != 200 || values[3] != 210 { + t.Errorf("distant x0s should remain unchanged: %v", values) + } +} + +func TestApplyIndentTolerance_Zero(t *testing.T) { + values := []float64{100, 101, 200} + applyIndentTolerance(values, 100, 0) + if values[1] != 101 { + t.Errorf("zero tolerance: x0s should be unchanged, got %v", values) + } +} + +func TestApplyIndentTolerance_Negative(t *testing.T) { + values := []float64{-100, -95, 0, 50} + applyIndentTolerance(values, -100, 10) + if values[0] != -100 || values[1] != -100 { + t.Errorf("negative x0s close to minX0 should be adjusted: %v", values) + } +} + +func TestFindBestK_SingleCluster(t *testing.T) { + // Note: KMeans1D uses random initialization, so non-identical values + // may occasionally produce k>1. This test verifies the function runs + // without error and returns k>=1 (not a correctness check). + x0s := []float64{100, 99, 101} + bestK, _ := findBestK(x0s, len(x0s)) + if bestK < 1 { + t.Errorf("expected bestK>=1, got %d", bestK) + } +} + +func TestFindBestK_TwoColumns(t *testing.T) { + x0s := []float64{50, 55, 60, 200, 210, 220} + bestK, _ := findBestK(x0s, len(x0s)) + if bestK != 2 { + t.Errorf("two columns: expected bestK=2, got %d", bestK) + } +} + +func TestFindBestK_OneValue(t *testing.T) { + x0s := []float64{100} + bestK, _ := findBestK(x0s, len(x0s)) + if bestK != 1 { + t.Errorf("single value: expected bestK=1, got %d", bestK) + } +} + +func TestFindBestK_Identical(t *testing.T) { + x0s := []float64{100, 100, 100, 100, 100} + bestK, _ := findBestK(x0s, len(x0s)) + if bestK != 1 { + t.Errorf("identical values: expected bestK=1, got %d", bestK) + } +} + +func TestRemapLabelsByCentroidOrder_Ordered(t *testing.T) { + centroids := []float64{50, 200, 400} + remap := remapLabelsByCentroidOrder(centroids) + if remap[0] != 0 || remap[1] != 1 || remap[2] != 2 { + t.Errorf("ordered centroids: expected 0->0,1->1,2->2, got %v", remap) + } +} + +func TestRemapLabelsByCentroidOrder_Unordered(t *testing.T) { + centroids := []float64{200, 50, 400} + remap := remapLabelsByCentroidOrder(centroids) + if remap[0] != 1 || remap[1] != 0 || remap[2] != 2 { + t.Errorf("unordered centroids: expected {0:1,1:0,2:2}, got %v", remap) + } +} + +func TestRemapLabelsByCentroidOrder_Nil(t *testing.T) { + remap := remapLabelsByCentroidOrder(nil) + if len(remap) != 0 { + t.Errorf("nil centroids: expected empty map, got %v", remap) + } +} + +func TestDetermineBestKForPage_SingleBox(t *testing.T) { + boxes := []pdf.TextBox{{PageNumber: 0, X0: 100, X1: 200}} + result := make([]pdf.TextBox, len(boxes)) + copy(result, boxes) + pageCols := make(map[int]int) + determineBestKForPage(boxes, result, []int{0}, 0, pageCols) + if pageCols[0] != 1 { + t.Errorf("single box: expected pageCols[0]=1, got %d", pageCols[0]) + } + if result[0].ColID != 0 { + t.Errorf("single box: expected ColID=0, got %d", result[0].ColID) + } +} + +func TestDetermineBestKForPage_TwoColumns(t *testing.T) { + boxes := []pdf.TextBox{ + {PageNumber: 0, X0: 50, X1: 100}, + {PageNumber: 0, X0: 55, X1: 100}, + {PageNumber: 0, X0: 300, X1: 400}, + {PageNumber: 0, X0: 310, X1: 400}, + } + result := make([]pdf.TextBox, len(boxes)) + copy(result, boxes) + pageCols := make(map[int]int) + determineBestKForPage(boxes, result, []int{0, 1, 2, 3}, 0, pageCols) + if pageCols[0] != 2 { + t.Errorf("two distinct columns: expected pageCols[0]=2, got %d", pageCols[0]) + } +} + +func TestAssignmentHelpers_IndentTolerance(t *testing.T) { + boxes := []pdf.TextBox{ + {PageNumber: 0, X0: 50, X1: 150, Top: 10, Bottom: 30}, + {PageNumber: 0, X0: 205, X1: 350, Top: 10, Bottom: 30}, + {PageNumber: 0, X0: 58, X1: 150, Top: 40, Bottom: 60}, + } + result := make([]pdf.TextBox, len(boxes)) + copy(result, boxes) + pageCols := make(map[int]int) + determineBestKForPage(boxes, result, []int{0, 1, 2}, 0, pageCols) + if pageCols[0] != 2 { + t.Errorf("expected 2 columns after indent tolerance, got %d", pageCols[0]) + } +} + +func TestAssignColIDsForPage_Normal(t *testing.T) { + boxes := []pdf.TextBox{ + {PageNumber: 0, X0: 50, X1: 100}, + {PageNumber: 0, X0: 200, X1: 300}, + } + result := make([]pdf.TextBox, len(boxes)) + copy(result, boxes) + pageCols := map[int]int{0: 2} + assignColIDsForPage(boxes, result, []int{0, 1}, 0, pageCols) + if result[0].ColID != 0 || result[1].ColID != 1 { + t.Errorf("expected ColIDs 0,1 but got %d,%d", result[0].ColID, result[1].ColID) + } +} + +func TestAssignColIDsForPage_KTooLarge(t *testing.T) { + boxes := []pdf.TextBox{ + {PageNumber: 0, X0: 100, X1: 200}, + } + result := make([]pdf.TextBox, len(boxes)) + copy(result, boxes) + pageCols := map[int]int{0: 3} + assignColIDsForPage(boxes, result, []int{0}, 0, pageCols) + if result[0].ColID != 0 { + t.Errorf("expected ColID=0 (k clamped to 1), got %d", result[0].ColID) + } +} diff --git a/internal/deepdoc/parser/pdf/parser.go b/internal/deepdoc/parser/pdf/parser.go index fc3ac96c73..8c4249309f 100644 --- a/internal/deepdoc/parser/pdf/parser.go +++ b/internal/deepdoc/parser/pdf/parser.go @@ -60,51 +60,20 @@ func NewTableBuilderFor(doc pdf.DocAnalyzer) pdf.TableBuilder { func (p *Parser) ParseRaw(ctx context.Context, engine pdf.PDFEngine, docAnalyzer pdf.DocAnalyzer) (*pdf.ParseResult, error) { tb := NewTableBuilderFor(docAnalyzer) - // Normalize page range - pageCount, err := engine.PageCount() + _, fromPage, toPage, err := p.normalizePageRange(engine) if err != nil { return nil, fmt.Errorf("page count: %w", err) } - toPage := p.Config.ToPage - if toPage < 0 || toPage >= pageCount { - toPage = pageCount - 1 - } - fromPage := p.Config.FromPage if toPage < fromPage { return &pdf.ParseResult{PageImages: make(map[int]image.Image)}, nil } totalPages := toPage - fromPage + 1 - batchSize := p.Config.BatchSize - if batchSize <= 0 { - batchSize = 50 - } + batchSize := p.getBatchSize() - // ── Prescan ── - prescanChars := make(map[int][]pdf.TextChar) - prescanMedianH := make(map[int]float64) - prescanMedianW := make(map[int]float64) - for pg := fromPage; pg <= toPage; pg++ { - chars, extractErr := engine.ExtractChars(pg) - if extractErr != nil { - slog.Warn("prescan: ExtractChars failed", "page", pg, "err", extractErr) - chars = nil - } - prescanChars[pg] = chars - prescanMedianH[pg] = util.MedianCharHeight(chars) - prescanMedianW[pg] = util.MedianCharWidth(chars) - } - isEnglish := util.DetectEnglish(prescanChars, totalPages, nil) - scanNoise := util.IsScanNoise(util.FullTextFromChars(prescanChars)) + prescanChars, prescanMedianH, prescanMedianW, isEnglish, scanNoise := p.prescanPages(ctx, engine, fromPage, toPage, totalPages) + outlines := p.extractOutlines(engine) - // ── Outlines ── - outlines, outlineErr := engine.Outlines() - if outlineErr != nil { - slog.Warn("Failed to extract PDF outlines; continuing without them", "err", outlineErr) - outlines = nil - } - - // ── Small document ── if totalPages <= batchSize { result, err := p.processPages(ctx, engine, fromPage, toPage, prescanChars, prescanMedianH, prescanMedianW, isEnglish, scanNoise, @@ -116,9 +85,91 @@ func (p *Parser) ParseRaw(ctx context.Context, engine pdf.PDFEngine, docAnalyzer return result, nil } - // ── Large document: batched ── - slog.Info("batched processing", "pages", totalPages, "batchSize", batchSize) + result, err := p.processLargeDocument(ctx, engine, fromPage, toPage, batchSize, + prescanChars, prescanMedianH, prescanMedianW, isEnglish, scanNoise, + docAnalyzer, tb) + if err != nil { + return nil, err + } + result.Outlines = outlines + return result, nil +} + +// ── ParseRaw helper functions ─────────────────────────────────────────────── + +// normalizePageRange normalizes the page range based on config and actual page count. +func (p *Parser) normalizePageRange(engine pdf.PDFEngine) (pageCount, fromPage, toPage int, err error) { + pageCount, err = engine.PageCount() + if err != nil { + return 0, 0, 0, err + } + fromPage = p.Config.FromPage + if fromPage < 0 { + fromPage = 0 + } else if fromPage >= pageCount { + fromPage = pageCount - 1 + } + toPage = p.Config.ToPage + if toPage < 0 || toPage >= pageCount { + toPage = pageCount - 1 + } + if toPage < fromPage { + toPage = fromPage + } + return pageCount, fromPage, toPage, nil +} + +// getBatchSize returns the batch size, defaulting to 50 if <= 0. +func (p *Parser) getBatchSize() int { + batchSize := p.Config.BatchSize + if batchSize <= 0 { + batchSize = 50 + } + return batchSize +} + +// prescanPages extracts chars from all pages and computes median heights/widths. +func (p *Parser) prescanPages(ctx context.Context, engine pdf.PDFEngine, fromPage, toPage, totalPages int) ( + map[int][]pdf.TextChar, map[int]float64, map[int]float64, bool, bool, +) { + prescanChars := make(map[int][]pdf.TextChar) + prescanMedianH := make(map[int]float64) + prescanMedianW := make(map[int]float64) + + for pg := fromPage; pg <= toPage; pg++ { + chars, extractErr := engine.ExtractChars(pg) + if extractErr != nil { + slog.Warn("prescan: ExtractChars failed", "page", pg, "err", extractErr) + chars = nil + } + prescanChars[pg] = chars + prescanMedianH[pg] = util.MedianCharHeight(chars) + prescanMedianW[pg] = util.MedianCharWidth(chars) + } + + isEnglish := util.DetectEnglish(prescanChars, totalPages, nil) + scanNoise := util.IsScanNoise(util.FullTextFromChars(prescanChars)) + return prescanChars, prescanMedianH, prescanMedianW, isEnglish, scanNoise +} + +// extractOutlines extracts the PDF outlines, returning nil on error. +func (p *Parser) extractOutlines(engine pdf.PDFEngine) []pdf.Outline { + outlines, outlineErr := engine.Outlines() + if outlineErr != nil { + slog.Warn("Failed to extract PDF outlines; continuing without them", "err", outlineErr) + outlines = nil + } + return outlines +} + +// processLargeDocument processes a large document in batches. +func (p *Parser) processLargeDocument(ctx context.Context, engine pdf.PDFEngine, fromPage, toPage, batchSize int, + prescanChars map[int][]pdf.TextChar, prescanMedianH, prescanMedianW map[int]float64, + isEnglish, scanNoise bool, docAnalyzer pdf.DocAnalyzer, tb pdf.TableBuilder, +) (*pdf.ParseResult, error) { + slog.Info("batched processing", "pages", toPage-fromPage+1, "batchSize", batchSize) result := &pdf.ParseResult{PageImages: make(map[int]image.Image)} + for start := fromPage; start <= toPage; start += batchSize { if err := ctx.Err(); err != nil { return nil, fmt.Errorf("cancelled at batch starting page %d: %w", start, err) @@ -141,22 +192,29 @@ func (p *Parser) ParseRaw(ctx context.Context, engine pdf.PDFEngine, docAnalyzer return nil, err } - result.Sections = append(result.Sections, batch.Sections...) - result.Tables = append(result.Tables, batch.Tables...) - for pg, img := range batch.PageImages { - result.PageImages[pg] = img - } - result.Metrics.BoxesInitial += batch.Metrics.BoxesInitial - result.Metrics.BoxesTextMerge += batch.Metrics.BoxesTextMerge - result.Metrics.BoxesVertMerge += batch.Metrics.BoxesVertMerge - result.Metrics.BoxesFinal += batch.Metrics.BoxesFinal - result.Metrics.TablesCount += batch.Metrics.TablesCount + p.mergeBatchResults(result, batch) } - result.Outlines = outlines return result, nil } -// ── Internal pipeline steps ──────────────────────────────────────────────── +// mergeBatchResults merges the batch result into the main result. +func (p *Parser) mergeBatchResults(result, batch *pdf.ParseResult) { + result.Sections = append(result.Sections, batch.Sections...) + result.Tables = append(result.Tables, batch.Tables...) + if result.PageImages == nil { + result.PageImages = make(map[int]image.Image) + } + for pg, img := range batch.PageImages { + result.PageImages[pg] = img + } + result.Metrics.BoxesInitial += batch.Metrics.BoxesInitial + result.Metrics.BoxesTextMerge += batch.Metrics.BoxesTextMerge + result.Metrics.BoxesVertMerge += batch.Metrics.BoxesVertMerge + result.Metrics.BoxesFinal += batch.Metrics.BoxesFinal + result.Metrics.TablesCount += batch.Metrics.TablesCount +} + +// ── extractPages helper functions ─────────────────────────────────────────── func (p *Parser) extractPages(ctx context.Context, engine pdf.PDFEngine, fromPage, toPage int, @@ -165,103 +223,167 @@ func (p *Parser) extractPages(ctx context.Context, engine pdf.PDFEngine, pageImages map[int]image.Image, docAnalyzer pdf.DocAnalyzer, ) ([]pdf.TextBox, map[int][]pdf.TextChar, bool, error) { - var boxes []pdf.TextBox - pageChars := make(map[int][]pdf.TextChar) - ocrUsedAny := false - pageCount := toPage - fromPage + 1 results := make([]pageResult, pageCount) - cap := p.Config.MaxOCRConcurrency - if cap <= 0 { - cap = 1 - } - sem := make(chan struct{}, cap) - var wg sync.WaitGroup + sem, wg := p.setupPageConcurrency() for i := 0; i < pageCount; i++ { pg := fromPage + i chars := prescanChars[pg] if len(chars) > 0 && !util.IsGarbledPage(chars) { - pageImg, renderErr := RenderPageToImage(engine, pg) - if renderErr == nil && pageImg != nil { - pageImages[pg] = pageImg + if err := ctx.Err(); err != nil { + results[i] = pageResult{pg: pg, err: fmt.Errorf("cancelled before sync page %d: %w", pg, err)} + continue } - var ocrBoxes []pdf.TextBox - ocrUsed := false - if !p.Config.SkipOCR && renderErr == nil && pageImg != nil { - ocrBoxes = ocrMergeChars(ctx, pageImg, chars, docAnalyzer, pg) - if ocrBoxes == nil { - ocrBoxes = lyt.CharsToBoxes(chars, pg, p.Config.SortByTop) - } else { - ocrUsed = true - ocrUsedAny = true - } - } else { - ocrBoxes = lyt.CharsToBoxes(chars, pg, p.Config.SortByTop) - } - results[i] = pageResult{pg: pg, ocrBoxes: ocrBoxes, chars: chars, ocrUsed: ocrUsed} + results[i] = p.processPageSync(ctx, engine, pg, chars, pageImages, docAnalyzer) continue } wg.Add(1) go func(i, pg int, chars []pdf.TextChar) { defer wg.Done() - select { - case <-ctx.Done(): - results[i] = pageResult{pg: pg, err: ctx.Err()} - return - case sem <- struct{}{}: - } - defer func() { <-sem }() + results[i] = p.processPageAsync(ctx, engine, pg, chars, sem, docAnalyzer) + }(i, pg, chars) + } + wg.Wait() + return p.collectPageResults(results, pageImages, medianHeights, medianWidths) +} - pageImg, err := RenderPageToImage(engine, pg) - if err != nil { - results[i] = pageResult{pg: pg, err: err} - return - } - if err := ctx.Err(); err != nil { - results[i] = pageResult{pg: pg, err: err} - return - } +// setupPageConcurrency sets up the concurrency primitives for page processing. +func (p *Parser) setupPageConcurrency() (chan struct{}, *sync.WaitGroup) { + maxConc := p.Config.MaxOCRConcurrency + if maxConc <= 0 { + maxConc = 1 + } + return make(chan struct{}, maxConc), &sync.WaitGroup{} +} - var ocrBoxes []pdf.TextBox - ocrUsed := false - if !p.Config.SkipOCR { - label := "scan page" - if len(chars) > 0 { - label = "garbled page" - } - ocrBoxes = ocrDetectAndRecognize(ctx, pageImg, docAnalyzer, pg, label) - if ocrBoxes != nil { - for j := range ocrBoxes { - for _, r := range ocrBoxes[j].Text { - chars = append(chars, pdf.TextChar{Text: string(r), PageNumber: pg}) - break - } +// processPageSync processes a page synchronously (normal pages with non-garbled chars). +func (p *Parser) processPageSync(ctx context.Context, engine pdf.PDFEngine, pg int, chars []pdf.TextChar, + pageImages map[int]image.Image, docAnalyzer pdf.DocAnalyzer, +) pageResult { + pageImg, renderErr := RenderPageToImage(engine, pg) + if renderErr == nil && pageImg != nil { + pageImages[pg] = pageImg + } else if renderErr != nil { + slog.Warn("processPageSync: RenderPageToImage failed", "page", pg, "err", renderErr) + } + + ocrBoxes, updatedChars, ocrUsed := p.processPageBoxes(ctx, pageImg, chars, pg, renderErr, docAnalyzer, false) + return pageResult{pg: pg, ocrBoxes: ocrBoxes, chars: updatedChars, ocrUsed: ocrUsed} +} + +// processPageAsync processes a page asynchronously (garbled pages or scan pages). +func (p *Parser) processPageAsync(ctx context.Context, engine pdf.PDFEngine, pg int, chars []pdf.TextChar, + sem chan struct{}, docAnalyzer pdf.DocAnalyzer, +) pageResult { + select { + case <-ctx.Done(): + return pageResult{pg: pg, err: ctx.Err()} + case sem <- struct{}{}: + } + defer func() { <-sem }() + + pageImg, err := RenderPageToImage(engine, pg) + if err != nil { + return pageResult{pg: pg, err: err} + } + if err := ctx.Err(); err != nil { + return pageResult{pg: pg, err: err} + } + + ocrBoxes, updatedChars, ocrUsed := p.processPageBoxes(ctx, pageImg, chars, pg, nil, docAnalyzer, true) + return pageResult{pg: pg, ocrBoxes: ocrBoxes, chars: updatedChars, ocrUsed: ocrUsed, pageImg: pageImg} +} + +// processPageBoxes processes OCR box extraction for a page, shared between sync and async paths. +// Returns (ocrBoxes, updatedChars, ocrUsed). The updatedChars includes synthetic OCR +// chars appended when OCR detect+recognize succeeds — callers must use the returned +// chars slice, not the original, to get correct median/layout calculations. +func (p *Parser) processPageBoxes(ctx context.Context, pageImg image.Image, chars []pdf.TextChar, pg int, + renderErr error, docAnalyzer pdf.DocAnalyzer, isAsync bool, +) ([]pdf.TextBox, []pdf.TextChar, bool) { + var ocrBoxes []pdf.TextBox + ocrUsed := false + + if !p.Config.SkipOCR { + if isAsync { + label := "scan page" + if len(chars) > 0 { + label = "garbled page" + } + ocrBoxes = ocrDetectAndRecognize(ctx, pageImg, docAnalyzer, pg, label) + if ocrBoxes != nil { + for j := range ocrBoxes { + for _, r := range ocrBoxes[j].Text { + chars = append(chars, pdf.TextChar{Text: string(r), PageNumber: pg}) + break } - ocrUsed = true } + ocrUsed = true } - if !ocrUsed && len(chars) > 0 && !p.Config.SkipOCR { + if !ocrUsed && len(chars) > 0 { ocrBoxes = ocrMergeChars(ctx, pageImg, chars, docAnalyzer, pg) if ocrBoxes != nil { ocrUsed = true } } - if !ocrUsed { - if len(chars) > 0 { - ocrBoxes = lyt.CharsToBoxes(chars, pg, p.Config.SortByTop) + } else { + if renderErr == nil && pageImg != nil { + ocrBoxes = ocrMergeChars(ctx, pageImg, chars, docAnalyzer, pg) + if ocrBoxes != nil { + ocrUsed = true } } - results[i] = pageResult{pg: pg, ocrBoxes: ocrBoxes, chars: chars, ocrUsed: ocrUsed, pageImg: pageImg} - }(i, pg, chars) + } } - wg.Wait() - return mergePageResults(results, boxes, pageImages, pageChars, ocrUsedAny, medianHeights, medianWidths) + + if !ocrUsed && len(chars) > 0 { + if ocrBoxes == nil { + ocrBoxes = lyt.CharsToBoxes(chars, pg, p.Config.SortByTop) + } + } + + return ocrBoxes, chars, ocrUsed } +// collectPageResults collects and merges the per-page results. +func (p *Parser) collectPageResults(results []pageResult, pageImages map[int]image.Image, + medianHeights, medianWidths map[int]float64, +) ([]pdf.TextBox, map[int][]pdf.TextChar, bool, error) { + var boxes []pdf.TextBox + pageChars := make(map[int][]pdf.TextChar) + ocrUsedAny := false + var errs []error + + for _, r := range results { + if r.err != nil { + slog.Warn("page OCR failed", "page", r.pg, "err", r.err) + errs = append(errs, fmt.Errorf("page %d: %w", r.pg, r.err)) + continue + } + if r.ocrUsed { + boxes = append(boxes, r.ocrBoxes...) + ocrUsedAny = true + } else if len(r.ocrBoxes) > 0 { + boxes = append(boxes, r.ocrBoxes...) + } + if r.pageImg != nil { + pageImages[r.pg] = r.pageImg + } + pageChars[r.pg] = r.chars + if r.ocrUsed { + medianHeights[r.pg] = util.MedianCharHeight(r.chars) + medianWidths[r.pg] = util.MedianCharWidth(r.chars) + } + } + return boxes, pageChars, ocrUsedAny, errors.Join(errs...) +} + +// ── Internal pipeline steps ──────────────────────────────────────────────── + func (p *Parser) retryScanNoise(ctx context.Context, engine pdf.PDFEngine, fromPage, toPage int, pageImages map[int]image.Image, @@ -494,32 +616,3 @@ func matchTableImage(sec *pdf.Section, tableImgByRegion map[string]string) (stri return "", false } -// mergePageResults collects per-page OCR results into the final output. -func mergePageResults(results []pageResult, boxes []pdf.TextBox, pageImages map[int]image.Image, - pageChars map[int][]pdf.TextChar, ocrUsedAny bool, - medianHeights, medianWidths map[int]float64, -) ([]pdf.TextBox, map[int][]pdf.TextChar, bool, error) { - var errs []error - for _, r := range results { - if r.err != nil { - slog.Warn("page OCR failed", "page", r.pg, "err", r.err) - errs = append(errs, fmt.Errorf("page %d: %w", r.pg, r.err)) - continue - } - if r.ocrUsed { - boxes = append(boxes, r.ocrBoxes...) - ocrUsedAny = true - } else if len(r.ocrBoxes) > 0 { - boxes = append(boxes, r.ocrBoxes...) - } - if r.pageImg != nil { - pageImages[r.pg] = r.pageImg - } - pageChars[r.pg] = r.chars - if r.ocrUsed { - medianHeights[r.pg] = util.MedianCharHeight(r.chars) - medianWidths[r.pg] = util.MedianCharWidth(r.chars) - } - } - return boxes, pageChars, ocrUsedAny, errors.Join(errs...) -} diff --git a/internal/deepdoc/parser/pdf/parser_ocr.go b/internal/deepdoc/parser/pdf/parser_ocr.go index fb803aedff..c8ed5c8972 100644 --- a/internal/deepdoc/parser/pdf/parser_ocr.go +++ b/internal/deepdoc/parser/pdf/parser_ocr.go @@ -22,11 +22,11 @@ func ocrDetectAndRecognize(ctx context.Context, pageImg image.Image, doc pdf.Doc } var result []pdf.TextBox - for _, box := range boxes { - x0 := int(math.Min(box.X0, math.Min(box.X1, math.Min(box.X2, box.X3)))) - y0 := int(math.Min(box.Y0, math.Min(box.Y1, math.Min(box.Y2, box.Y3)))) - x1 := int(math.Max(box.X0, math.Max(box.X1, math.Max(box.X2, box.X3)))) - y1 := int(math.Max(box.Y0, math.Max(box.Y1, math.Max(box.Y2, box.Y3)))) + for _, b := range boxes { + x0 := int(math.Min(b.X0, math.Min(b.X1, math.Min(b.X2, b.X3)))) + y0 := int(math.Min(b.Y0, math.Min(b.Y1, math.Min(b.Y2, b.Y3)))) + x1 := int(math.Max(b.X0, math.Max(b.X1, math.Max(b.X2, b.X3)))) + y1 := int(math.Max(b.Y0, math.Max(b.Y1, math.Max(b.Y2, b.Y3)))) if x0 >= x1 || y0 >= y1 { continue } @@ -41,7 +41,7 @@ func ocrDetectAndRecognize(ctx context.Context, pageImg image.Image, doc pdf.Doc result = append(result, pdf.TextBox{ X0: float64(x0), X1: float64(x1), Top: float64(y0), Bottom: float64(y1), - Text: t.Text, + Text: t.Text, PageNumber: pageNum, }) } @@ -60,20 +60,26 @@ type ocrDetectBox struct { } func ocrMergeChars(ctx context.Context, pageImg image.Image, chars []pdf.TextChar, doc pdf.DocAnalyzer, pageNum int) []pdf.TextBox { + boxes, scale, err := detectBoxes(ctx, pageImg, doc, pageNum) + if err != nil || len(boxes) == 0 { + return nil + } + boxChars := matchCharsToBoxes(boxes, chars) + return buildTextBoxes(ctx, pageImg, boxes, boxChars, doc, scale, pageNum) +} + +func detectBoxes(ctx context.Context, pageImg image.Image, doc pdf.DocAnalyzer, pageNum int) ([]ocrDetectBox, float64, error) { ocrDetectBoxes, err := doc.OCRDetect(ctx, pageImg) if err != nil || len(ocrDetectBoxes) == 0 { - return nil + return nil, 0, err } slog.Debug("ocrMergeChars detect", "page", pageNum, "boxes", len(ocrDetectBoxes)) - // Detect boxes are in pixel space (216 DPI). Scale to PDF space (72 DPI) - // so coordinates match embedded chars. scale := pdf.DlaScale // 3.0 imgBounds := pageImg.Bounds() imgW := float64(imgBounds.Dx()) / scale imgH := float64(imgBounds.Dy()) / scale - // Step 1: match embedded chars to detect boxes (Python __ocr char matching). boxes := make([]ocrDetectBox, 0, len(ocrDetectBoxes)) for _, b := range ocrDetectBoxes { x0 := min(b.X0, b.X1, b.X2, b.X3) / scale @@ -100,8 +106,6 @@ func ocrMergeChars(ctx context.Context, pageImg image.Image, chars []pdf.TextCha }, x0: x0, y0: y0, x1: x1, y1: y1}) } - // Sort detect boxes top-down (fuzzy Y-group), matching Python's - // Recognizer.sort_Y_firstly with threshold = median box height / 3. if len(boxes) > 1 { boxHeights := make([]float64, len(boxes)) for i := range boxes { @@ -109,20 +113,21 @@ func ocrMergeChars(ctx context.Context, pageImg image.Image, chars []pdf.TextCha } sort.Float64s(boxHeights) threshold := boxHeights[len(boxHeights)/2] / 3 - sort.Slice(boxes, func(a, b int) bool { - if math.Abs(boxes[a].y0-boxes[b].y0) < threshold { - return boxes[a].x0 < boxes[b].x0 + sort.Slice(boxes, func(i, j int) bool { + if math.Abs(boxes[i].y0-boxes[j].y0) < threshold { + return boxes[i].x0 < boxes[j].x0 } - return boxes[a].y0 < boxes[b].y0 + return boxes[i].y0 < boxes[j].y0 }) } + return boxes, scale, nil +} - // Step 2: match each char to the best overlapping detect box - // (char perspective), matching Python's find_overlapped. +func matchCharsToBoxes(boxes []ocrDetectBox, chars []pdf.TextChar) [][]pdf.TextChar { boxChars := make([][]pdf.TextChar, len(boxes)) for _, c := range chars { bestIdx := -1 - bestOverlap := 1e-6 // Python: thr=1e-6 + bestOverlap := 1e-6 for i := range boxes { overlap := charBoxOverlapRatio(c, boxes[i].x0, boxes[i].x1, boxes[i].y0, boxes[i].y1) if overlap >= bestOverlap { @@ -133,8 +138,6 @@ func ocrMergeChars(ctx context.Context, pageImg image.Image, chars []pdf.TextCha if bestIdx < 0 { continue } - // Height gating, matching Python: skip when height differs >70%, - // except space chars which are always kept. ch := c.Bottom - c.Top if ch <= 0 { ch = 1 @@ -145,27 +148,26 @@ func ocrMergeChars(ctx context.Context, pageImg image.Image, chars []pdf.TextCha } boxChars[bestIdx] = append(boxChars[bestIdx], c) } - - return buildTextBoxes(ctx, pageImg, boxes, boxChars, doc, scale, pageNum) + return boxChars } -// sortYFirstly sorts chars by Y (fuzzy group by threshold), then by X. +// sortCharsYFirstly sorts chars by Y (fuzzy group by threshold), then by X. // Matching Python Recognizer.sort_Y_firstly in recognizer.py:26-33: // // If two chars have Y diff < threshold → same line → sort by X. // Otherwise → sort by Y. func sortCharsYFirstly(chars []pdf.TextChar, threshold float64) { - sort.Slice(chars, func(a, b int) bool { - diff := chars[a].Top - chars[b].Top + sort.Slice(chars, func(i, j int) bool { + diff := chars[i].Top - chars[j].Top if math.Abs(diff) < threshold { - return chars[a].X0 < chars[b].X0 + return chars[i].X0 < chars[j].X0 } return diff < 0 }) } -// charBoxOverlapRatio computes the overlap ratio between a char and a box, -// from the char's perspective. Returns overlap_area / char_area. +// charBoxOverlapRatio computes overlap ratio between a char and a box, +// from char perspective. Returns overlap_area / char_area. // Matching Python's Recognizer.overlapped_area(char, box, ratio=True). func charBoxOverlapRatio(c pdf.TextChar, x0, x1, y0, y1 float64) float64 { cw := c.X1 - c.X0 @@ -216,8 +218,7 @@ func ocrTableCells(ctx context.Context, cells []pdf.TSRCell, tableImg image.Imag } } -// buildTextBoxes assembles detect box text from embedded chars and fills -// empty boxes via batch OCR. +// buildTextBoxes assembles detect box text from embedded chars and fills empty boxes via batch OCR. func buildTextBoxes(ctx context.Context, pageImg image.Image, boxes []ocrDetectBox, boxChars [][]pdf.TextChar, doc pdf.DocAnalyzer, scale float64, pageNum int, ) []pdf.TextBox { diff --git a/internal/deepdoc/parser/pdf/parser_test.go b/internal/deepdoc/parser/pdf/parser_test.go index db37502918..de44ab0b1d 100644 --- a/internal/deepdoc/parser/pdf/parser_test.go +++ b/internal/deepdoc/parser/pdf/parser_test.go @@ -14,6 +14,24 @@ import ( util "ragflow/internal/deepdoc/parser/pdf/util" ) +// ---- test helpers ---- + +func newTestParser() *Parser { + return &Parser{Config: pdf.DefaultParserConfig()} +} + +func newMockDocAnalyzer(healthy bool, boxes []pdf.OCRBox, texts []pdf.OCRText) *MockDocAnalyzer { + return &MockDocAnalyzer{ + Healthy: healthy, + OCRBoxes: boxes, + OCRTexts: texts, + } +} + +func newSimpleMockDocAnalyzer() *MockDocAnalyzer { + return &MockDocAnalyzer{Healthy: true} +} + // ── OCR fallback ────────────────────────────────────────────────────── func TestOCR_Fallback(t *testing.T) { @@ -34,7 +52,7 @@ func TestOCR_Fallback(t *testing.T) { t.Run("detect + recognize success", func(t *testing.T) { mock := &MockDocAnalyzer{ - Healthy: true, + Healthy: true, OCRBoxes: []pdf.OCRBox{{X0: 10, Y0: 20, X1: 90, Y1: 20, X2: 90, Y2: 40, X3: 10, Y3: 40}}, OCRTexts: []pdf.OCRText{{Text: "Hello", Confidence: 0.9}}, } @@ -49,7 +67,7 @@ func TestOCR_Fallback(t *testing.T) { t.Run("detect boxes but rec returns empty text", func(t *testing.T) { mock := &MockDocAnalyzer{ - Healthy: true, + Healthy: true, OCRBoxes: []pdf.OCRBox{{X0: 10, Y0: 20, X1: 90, Y1: 20, X2: 90, Y2: 40, X3: 10, Y3: 40}}, OCRTexts: []pdf.OCRText{{Text: "", Confidence: 0.1}}, } @@ -97,7 +115,7 @@ func TestOCR_ScanPage(t *testing.T) { t.Run("detect success but rec returns empty", func(t *testing.T) { mock := &MockDocAnalyzer{ - Healthy: true, + Healthy: true, OCRBoxes: []pdf.OCRBox{{X0: 10, Y0: 20, X1: 90, Y1: 20, X2: 90, Y2: 40, X3: 10, Y3: 40}}, OCRTexts: []pdf.OCRText{}, } @@ -196,7 +214,7 @@ func garbledSample() []pdf.TextChar { return chars } -// ── OCR fallback integration through Parse ───────────────────────────── +// ── OCR fallback integration through Parse ────────────────────────────── func TestOCR_FallbackIntegration(t *testing.T) { // ocrFallback logic is tested via TestOCR_fallback. @@ -314,7 +332,7 @@ func TestOCR_Fallback_PUAGarbled(t *testing.T) { } dummyImg := image.NewRGBA(image.Rect(0, 0, 100, 100)) mock := &MockDocAnalyzer{ - Healthy: true, + Healthy: true, OCRBoxes: []pdf.OCRBox{{X0: 10, Y0: 20, X1: 90, Y1: 20, X2: 90, Y2: 40, X3: 10, Y3: 40}}, OCRTexts: []pdf.OCRText{{Text: "PUA OCR text", Confidence: 0.9}}, } @@ -324,7 +342,7 @@ func TestOCR_Fallback_PUAGarbled(t *testing.T) { } } -// ── ocrMergeChars ───────────────────────────────────────────────────── +// ── ocrMergeChars ────────────────────────────────────────────────────── func TestOCR_MergeChars(t *testing.T) { dummyImg := image.NewRGBA(image.Rect(0, 0, 600, 600)) @@ -346,7 +364,7 @@ func TestOCR_MergeChars(t *testing.T) { t.Run("detect boxes — all overlap with chars (chars used, Python-aligned)", func(t *testing.T) { mock := &MockDocAnalyzer{ - Healthy: true, + Healthy: true, OCRBoxes: []pdf.OCRBox{{X0: 15, Y0: 15, X1: 150, Y1: 15, X2: 150, Y2: 150, X3: 15, Y3: 150}}, OCRTexts: []pdf.OCRText{{Text: "Hello OCR", Confidence: 0.9}}, } @@ -363,7 +381,7 @@ func TestOCR_MergeChars(t *testing.T) { t.Run("detect boxes — none overlap with chars", func(t *testing.T) { mock := &MockDocAnalyzer{ - Healthy: true, + Healthy: true, OCRBoxes: []pdf.OCRBox{{X0: 240, Y0: 240, X1: 270, Y1: 240, X2: 270, Y2: 270, X3: 240, Y3: 270}}, OCRTexts: []pdf.OCRText{{Text: "OCR", Confidence: 0.9}}, } @@ -379,7 +397,7 @@ func TestOCR_MergeChars(t *testing.T) { t.Run("detect box — no chars and OCR returns empty", func(t *testing.T) { mock := &MockDocAnalyzer{ - Healthy: true, + Healthy: true, OCRBoxes: []pdf.OCRBox{{X0: 240, Y0: 240, X1: 270, Y1: 240, X2: 270, Y2: 270, X3: 240, Y3: 270}}, OCRTexts: []pdf.OCRText{}, } @@ -495,7 +513,7 @@ func TestOCR_MergeChars(t *testing.T) { // Chars at (10,10-25) → within the box region. Char text "do" is // used (Python-aligned: embedded chars are more precise than OCR). mock := &MockDocAnalyzer{ - Healthy: true, + Healthy: true, OCRBoxes: []pdf.OCRBox{{X0: 30, Y0: 30, X1: 90, Y1: 30, X2: 90, Y2: 90, X3: 30, Y3: 90}}, OCRTexts: []pdf.OCRText{{Text: "docker commit infiniflow", Confidence: 0.95}}, } @@ -546,7 +564,7 @@ func TestTableSectionCaptionInHTML(t *testing.T) { figures := pdf.CollectFigures(sections) sections = tbl.MergeCaptions(sections, figures) - if !strings.HasPrefix(sections[0].Text, "表1: 交通工具等级") { + if !strings.HasPrefix(sections[0].Text, "表1: 交通工具等级 maxR { + maxR = b.R + } + } + if maxR <= 0 { + return GroupBoxesByYX(boxes) + } + // Sort by R index first (Python: sort_R_firstly), then Y, then X. + sort.Slice(boxes, func(i, j int) bool { + if boxes[i].R != boxes[j].R { + return boxes[i].R < boxes[j].R + } + if boxes[i].Top != boxes[j].Top { + return boxes[i].Top < boxes[j].Top + } + return boxes[i].X0 < boxes[j].X0 + }) + + // Compress R indices: Python's sort_R_firstly grouping. + rowMap, compressed := compressRowIndices(boxes) + + // Collect boxes per row. + cmap, _ := collectBoxesPerRow(boxes, rowMap) + + // Compress C indices per row. + cCompressed, cMaxCol := compressColIndices(boxes, rowMap, compressed) + + // Build grid. + return buildGrid(cmap, cCompressed, cMaxCol, compressed) +} + +// GroupBoxesByYX groups boxes into a cell grid by Y/X coordinates, +// matching Python's construct_table which uses sort_R_firstly and +// sort_C_firstly when R/C annotations are absent. Falls back from +// GroupBoxesByRC when boxes lack R/C annotations. +func GroupBoxesByYX(boxes []pdf.TextBox) [][]pdf.TSRCell { + if len(boxes) == 0 { + return nil + } + // Sort by (page, top, x0) — same as Python sort_R_firstly with R=-1. + sort.Slice(boxes, func(i, j int) bool { + if boxes[i].PageNumber != boxes[j].PageNumber { + return boxes[i].PageNumber < boxes[j].PageNumber + } + if boxes[i].Top != boxes[j].Top { + return boxes[i].Top < boxes[j].Top + } + return boxes[i].X0 < boxes[j].X0 + }) + + // Group into rows by Y proximity (Python's row grouping). + type rowGroup struct { + boxes []pdf.TextBox + top, btm float64 + } + var rowGroups []rowGroup + rowGroups = append(rowGroups, rowGroup{ + boxes: []pdf.TextBox{boxes[0]}, + top: boxes[0].Top, + btm: boxes[0].Bottom, + }) + for i := 1; i < len(boxes); i++ { + prev := &rowGroups[len(rowGroups)-1] + // Python: same row if top < prev.btm (Y overlaps) and same page. + if boxes[i].PageNumber == prev.boxes[0].PageNumber && boxes[i].Top < prev.btm { + prev.boxes = append(prev.boxes, boxes[i]) + if boxes[i].Top < prev.top { + prev.top = boxes[i].Top + } + if boxes[i].Bottom > prev.btm { + prev.btm = boxes[i].Bottom + } + } else { + rowGroups = append(rowGroups, rowGroup{ + boxes: []pdf.TextBox{boxes[i]}, + top: boxes[i].Top, + btm: boxes[i].Bottom, + }) + } + } + + // Within each row, group into columns by X proximity. + rows := make([][]pdf.TSRCell, len(rowGroups)) + for ri, rg := range rowGroups { + // Sort by X0. + sort.Slice(rg.boxes, func(i, j int) bool { + return rg.boxes[i].X0 < rg.boxes[j].X0 + }) + // Group by X overlap. + var cols []struct { + boxes []pdf.TextBox + x1 float64 + } + cols = append(cols, struct { + boxes []pdf.TextBox + x1 float64 + }{ + boxes: []pdf.TextBox{rg.boxes[0]}, + x1: rg.boxes[0].X1, + }) + for i := 1; i < len(rg.boxes); i++ { + prev := &cols[len(cols)-1] + if rg.boxes[i].X0 < prev.x1 { + prev.boxes = append(prev.boxes, rg.boxes[i]) + if rg.boxes[i].X1 > prev.x1 { + prev.x1 = rg.boxes[i].X1 + } + } else { + cols = append(cols, struct { + boxes []pdf.TextBox + x1 float64 + }{ + boxes: []pdf.TextBox{rg.boxes[i]}, + x1: rg.boxes[i].X1, + }) + } + } + rows[ri] = make([]pdf.TSRCell, len(cols)) + for ci, col := range cols { + var sb strings.Builder + for _, b := range col.boxes { + t := strings.TrimSpace(b.Text) + if t == "" { + continue + } + if sb.Len() > 0 { + sb.WriteByte(' ') + } + sb.WriteString(t) + } + rows[ri][ci].Text = sb.String() + } + } + return rows +} + +// cellPosFromBox returns the position coordinates and label for a cell +// derived from a text box. Header cells use HLeft/HRight/HTop/HBott +// for spanning-aware positions; regular cells use the box's own bounds. +func cellPosFromBox(b pdf.TextBox) (x0, y0, x1, y1 float64, label string) { + x0, y0, x1, y1 = b.X0, b.Top, b.X1, b.Bottom + if b.H > 0 { + label = "table header" + if b.HLeft != 0 || b.HRight != 0 { + if b.HLeft != 0 { + x0 = b.HLeft + } + if b.HRight != 0 { + x1 = b.HRight + } + } + if b.HTop != 0 { + y0 = b.HTop + } + if b.HBott != 0 { + y1 = b.HBott + } + } else if b.SP > 0 { + label = "table spanning cell" + } + return +} + +// cellLabelFromBox returns the TSR label for a box based on H/SP annotations. +// Used when merging multiple boxes into one cell — preserves the spanning label. +func cellLabelFromBox(b pdf.TextBox) string { + if b.H > 0 { + return "table header" + } + if b.SP > 0 { + return "table spanning cell" + } + return "" +} + +// compressRowIndices compresses R values into contiguous row indices. +// Returns rowMap (original R → compressed index) and the maximum compressed index. +// Boxes must already be sorted by R, Y, X. +func compressRowIndices(boxes []pdf.TextBox) (map[int]int, int) { + rowMap := make(map[int]int) // original R → compressed row index + compressed := 0 + rowMap[boxes[0].R] = 0 + lastR := boxes[0].R + for i := 1; i < len(boxes); i++ { + if boxes[i].R != lastR { + compressed++ + rowMap[boxes[i].R] = compressed + lastR = boxes[i].R + } else { + rowMap[boxes[i].R] = compressed + } + } + return rowMap, compressed +} + +// collectBoxesPerRow collects boxes into row groups, merging boxes in the same cell. +// Returns cmap (row → col → entry) and maxCols (max column index per row). +func collectBoxesPerRow(boxes []pdf.TextBox, rowMap map[int]int) (map[int]map[int]*rb, map[int]int) { + cmap := make(map[int]map[int]*rb) // row → col → entry + maxCols := make(map[int]int) + for _, b := range boxes { + t := strings.TrimSpace(b.Text) + // Keep boxes with SP/H annotations even if text is empty — + // their coordinates are needed for colspan/rowspan calculation. + if t == "" && b.H <= 0 && b.SP <= 0 { + continue + } + r := rowMap[b.R] + c := b.C + if cmap[r] == nil { + cmap[r] = make(map[int]*rb) + } + x0, y0, x1, y1, label := cellPosFromBox(b) + if v, ok := cmap[r][c]; ok { + if t != "" { + v.txt += " " + t + } + // Merge spanning coordinates (use widest extent). + if b.H > 0 || b.SP > 0 { + v.label = cellLabelFromBox(b) + if v.x0 > x0 { + v.x0 = x0 + } + if v.y0 > y0 { + v.y0 = y0 + } + if v.x1 < x1 { + v.x1 = x1 + } + if v.y1 < y1 { + v.y1 = y1 + } + } + } else { + cmap[r][c] = &rb{r, c, t, x0, y0, x1, y1, label} + } + if c > maxCols[r] { + maxCols[r] = c + } + } + return cmap, maxCols +} + +// rowBox is a helper for compressColIndices. +type rowBox struct { + c, idx int + x0, x1 float64 + txt string +} + +// compressColIndices compresses column indices per row based on X0 ordering and overlap. +// Returns cCompressed (row → original C → compressed C) and cMaxCol (max compressed C per row). +func compressColIndices(boxes []pdf.TextBox, rowMap map[int]int, compressed int) (map[int]map[int]int, map[int]int) { + cCompressed := make(map[int]map[int]int) // row → (original C → compressed col) + cMaxCol := make(map[int]int) + for ri := 0; ri <= compressed; ri++ { + // Collect all boxes in this row, sorted by X0. + var rowBoxes []rowBox + for i, b := range boxes { + if rowMap[b.R] == ri && (strings.TrimSpace(b.Text) != "" || b.H > 0 || b.SP > 0) { + rowBoxes = append(rowBoxes, rowBox{c: b.C, idx: i, x0: b.X0, x1: b.X1, txt: b.Text}) + } + } + sort.Slice(rowBoxes, func(i, j int) bool { return rowBoxes[i].x0 < rowBoxes[j].x0 }) + // Assign compressed column by X-order (disjoint X → new col). + cMap := make(map[int]int) // original C → compressed col + right := 0.0 + nCols := 0 + for _, rb := range rowBoxes { + if len(cMap) == 0 || rb.x0 >= right { + cMap[rb.c] = nCols + nCols++ + right = rb.x1 + } else { + // Overlapping X → merge into last column. + cMap[rb.c] = nCols - 1 + if rb.x1 > right { + right = rb.x1 + } + } + } + cCompressed[ri] = cMap + cMaxCol[ri] = nCols - 1 + } + return cCompressed, cMaxCol +} + +// buildGrid builds the final cell grid from the collected and compressed data. +func buildGrid(cmap map[int]map[int]*rb, cCompressed map[int]map[int]int, cMaxCol map[int]int, compressed int) [][]pdf.TSRCell { + rows := make([][]pdf.TSRCell, compressed+1) + for ri := 0; ri <= compressed; ri++ { + maxC := cMaxCol[ri] + rows[ri] = make([]pdf.TSRCell, maxC+1) + for ci, v := range cmap[ri] { + cci := cCompressed[ri][ci] + if cci <= maxC { + if rows[ri][cci].Text == "" { + rows[ri][cci].Text = v.txt + rows[ri][cci].X0 = v.x0 + rows[ri][cci].Y0 = v.y0 + rows[ri][cci].X1 = v.x1 + rows[ri][cci].Y1 = v.y1 + rows[ri][cci].Label = v.label + } else { + // Multiple originals map to same compressed cell — merge deterministically. + if v.txt != "" { + rows[ri][cci].Text += " " + v.txt + } + if v.x0 < rows[ri][cci].X0 { + rows[ri][cci].X0 = v.x0 + } + if v.y0 < rows[ri][cci].Y0 { + rows[ri][cci].Y0 = v.y0 + } + if v.x1 > rows[ri][cci].X1 { + rows[ri][cci].X1 = v.x1 + } + if v.y1 > rows[ri][cci].Y1 { + rows[ri][cci].Y1 = v.y1 + } + if rows[ri][cci].Label == "" && v.label != "" { + rows[ri][cci].Label = v.label + } + } + } + } + } + return rows +} + +func BoxesHaveAnnotations(boxes []pdf.TextBox) bool { + maxR, maxC := 0, 0 + for _, b := range boxes { + if b.R > maxR { + maxR = b.R + } + if b.C > maxC { + maxC = b.C + } + } + // True if at least 2 rows or 2 cols (R/C are 0-based, so maxR>0 means ≥2 rows). + return maxR > 0 || maxC > 0 +} diff --git a/internal/deepdoc/parser/pdf/table/group_boxes_test.go b/internal/deepdoc/parser/pdf/table/group_boxes_test.go new file mode 100644 index 0000000000..6ff203e967 --- /dev/null +++ b/internal/deepdoc/parser/pdf/table/group_boxes_test.go @@ -0,0 +1,350 @@ + +package table + +import ( + "sort" + "strings" + "testing" + + pdf "ragflow/internal/deepdoc/parser/pdf/type" +) + +func TestGroupBoxesByRC_RDiffSplitsRows(t *testing.T) { + // 6 boxes with 6 different R values → 6 rows (Python R-first splitting). + boxes := []pdf.TextBox{ + {X0: 10, X1: 90, Top: 0, Bottom: 30, Text: "A", R: 0, C: 0}, + {X0: 110, X1: 190, Top: 0, Bottom: 30, Text: "B", R: 1, C: 1}, + {X0: 210, X1: 290, Top: 0, Bottom: 30, Text: "C", R: 2, C: 2}, + {X0: 10, X1: 90, Top: 35, Bottom: 65, Text: "D", R: 3, C: 0}, + {X0: 110, X1: 190, Top: 35, Bottom: 65, Text: "E", R: 4, C: 1}, + {X0: 210, X1: 290, Top: 35, Bottom: 65, Text: "F", R: 5, C: 2}, + } + rows := GroupBoxesByRC(boxes) + // R=0,1,2,3,4,5 → 6 rows (Python: R differs → new row). + if len(rows) != 6 { + t.Fatalf("expected 6 rows (R differs → split), got %d", len(rows)) + } +} + +func TestGroupBoxesByRC_MergesCloseCols(t *testing.T) { + // R=0 has C=0,1. R=1 has C=0,1. C compression → 2 cols each. + boxes := []pdf.TextBox{ + {X0: 10, X1: 90, Top: 0, Bottom: 30, Text: "A", R: 0, C: 0}, + {X0: 110, X1: 190, Top: 0, Bottom: 30, Text: "B", R: 0, C: 1}, + {X0: 10, X1: 90, Top: 35, Bottom: 65, Text: "C", R: 1, C: 0}, + {X0: 110, X1: 190, Top: 35, Bottom: 65, Text: "D", R: 1, C: 1}, + } + rows := GroupBoxesByRC(boxes) + if len(rows) != 2 { + t.Fatalf("expected 2 rows (R diff), got %d", len(rows)) + } + if len(rows[0]) != 2 || len(rows[1]) != 2 { + t.Errorf("expected 2 cols/row, got %d/%d", len(rows[0]), len(rows[1])) + } + if rows[0][0].Text != "A" || rows[0][1].Text != "B" { + t.Errorf("row0 wrong: %q %q", rows[0][0].Text, rows[0][1].Text) + } + if rows[1][0].Text != "C" || rows[1][1].Text != "D" { + t.Errorf("row1 wrong: %q %q", rows[1][0].Text, rows[1][1].Text) + } +} + +func TestGroupBoxesByRC_RDiffSplitsRow(t *testing.T) { + // R=0 and R=1 at same Y (overlapping) → two separate rows in the grid. + boxes := []pdf.TextBox{ + {X0: 10, X1: 90, Top: 0, Bottom: 30, Text: "A", R: 0, C: 0}, + {X0: 110, X1: 190, Top: 0, Bottom: 30, Text: "B", R: 1, C: 1}, + {X0: 10, X1: 90, Top: 35, Bottom: 65, Text: "C", R: 2, C: 0}, + {X0: 110, X1: 190, Top: 35, Bottom: 65, Text: "D", R: 3, C: 1}, + } + rows := GroupBoxesByRC(boxes) + // R=0,1,2,3 → 4 different R values → 4 rows (Python: R differs → new row). + if len(rows) != 4 { + t.Fatalf("expected 4 rows (R differs → split), got %d", len(rows)) + } + if rows[0][0].Text != "A" || rows[1][0].Text != "B" { + t.Errorf("row0/1 wrong: A=%q B=%q", rows[0][0].Text, rows[1][0].Text) + } +} + +func TestFillCellTextFromBoxes_RCOnly(t *testing.T) { + cells := []pdf.TSRCell{ + {X0: 0, Y0: 0, X1: 100, Y1: 50, Label: "table"}, + {X0: 90, Y0: 0, X1: 200, Y1: 50, Label: "table"}, + } + // This box straddles cell 0 (X=0-100) and cell 1 (X=90-200). + // Spatial overlap: both match. R/C: should go to cell R=0, C=0 only. + boxes := []pdf.TextBox{ + {X0: 80, X1: 120, Top: 0, Bottom: 50, Text: "TEXT", LayoutType: "table", R: 0, C: 0}, + } + rows := GroupTSRCellsToRows(cells) + for _, b := range boxes { + t := strings.TrimSpace(b.Text) + if t == "" { + continue + } + if b.R >= 0 && b.R < len(rows) && b.C >= 0 && b.C < len(rows[b.R]) { + rows[b.R][b.C].Text = t + } + } + // Cell 0 should have text, cell 1 should NOT. + if rows[0][0].Text != "TEXT" { + t.Errorf("cell[0][0] = %q, want %q", rows[0][0].Text, "TEXT") + } + if rows[0][1].Text != "" { + t.Errorf("cell[0][1] = %q, should be empty (spatial overlap leak)", rows[0][1].Text) + } +} + +func TestGroupBoxesByRC_FallbackToYXWhenNoAnnotations(t *testing.T) { + // When all boxes have R=-1 (Python's case: regex didn't match "table" label), + // groupBoxesByRC should fall back to YX coordinate grouping. + boxes := []pdf.TextBox{ + {X0: 10, X1: 90, Top: 0, Bottom: 30, Text: "A", R: -1, C: -1}, + {X0: 110, X1: 190, Top: 0, Bottom: 30, Text: "B", R: -1, C: -1}, + {X0: 10, X1: 90, Top: 35, Bottom: 65, Text: "C", R: -1, C: -1}, + {X0: 110, X1: 190, Top: 35, Bottom: 65, Text: "D", R: -1, C: -1}, + } + rows := GroupBoxesByRC(boxes) + // R=-1 for all → maxR = -1 → grid would be 0 rows. Must fall back to YX. + if len(rows) == 0 { + t.Fatal("groupBoxesByRC returned 0 rows when R=-1 — no YX fallback") + } + if len(rows) != 2 { + t.Errorf("expected 2 rows (Y-split), got %d", len(rows)) + } +} + +func TestGroupBoxesByRC_ColspanMissing(t *testing.T) { + // Box with SP annotation spanning 2 columns (HLeft→HRight covers cols 0-1). + boxes := []pdf.TextBox{ + {X0: 10, X1: 90, Top: 0, Bottom: 30, Text: "Name", R: 0, C: 0, H: 1, + HLeft: 10, HRight: 200}, + {X0: 110, X1: 200, Top: 0, Bottom: 30, Text: "", R: 0, C: 1, SP: 1}, + {X0: 10, X1: 90, Top: 35, Bottom: 65, Text: "A", R: 1, C: 0}, + {X0: 110, X1: 200, Top: 35, Bottom: 65, Text: "B", R: 1, C: 1}, + } + rows := GroupBoxesByRC(boxes) + // The result should have colspan=2 for cell [0,0] and skip [0,1]. + // Currently groupBoxesByRC produces a flat grid without span info. + if len(rows) >= 1 && len(rows[0]) >= 2 && rows[0][1].Text == "" { + t.Log("KNOWN LIMITATION: colspan not computed — cell [0,1] is empty instead of merged") + } + _ = rows +} + +func TestCompressRowIndices(t *testing.T) { + // 6 boxes with 6 different R values → 6 rows. + boxes := []pdf.TextBox{ + {X0: 10, X1: 90, Top: 0, Bottom: 30, Text: "A", R: 0, C: 0}, + {X0: 110, X1: 190, Top: 0, Bottom: 30, Text: "B", R: 1, C: 1}, + {X0: 210, X1: 290, Top: 0, Bottom: 30, Text: "C", R: 2, C: 2}, + {X0: 10, X1: 90, Top: 35, Bottom: 65, Text: "D", R: 3, C: 0}, + {X0: 110, X1: 190, Top: 35, Bottom: 65, Text: "E", R: 4, C: 1}, + {X0: 210, X1: 290, Top: 35, Bottom: 65, Text: "F", R: 5, C: 2}, + } + // Sort first (as GroupBoxesByRC does) + sort.Slice(boxes, func(i, j int) bool { + if boxes[i].R != boxes[j].R { + return boxes[i].R < boxes[j].R + } + if boxes[i].Top != boxes[j].Top { + return boxes[i].Top < boxes[j].Top + } + return boxes[i].X0 < boxes[j].X0 + }) + + rowMap, compressed := compressRowIndices(boxes) + if compressed != 5 { // 0-based, 6 elements → max index 5 + t.Errorf("compressed = %d, want 5", compressed) + } + if rowMap[0] != 0 || rowMap[1] != 1 || rowMap[2] != 2 || rowMap[3] != 3 || rowMap[4] != 4 || rowMap[5] != 5 { + t.Errorf("rowMap mapping incorrect: %v", rowMap) + } +} + +func TestCompressRowIndices_SameR(t *testing.T) { + // Multiple boxes with same R → same compressed row. + boxes := []pdf.TextBox{ + {X0: 10, X1: 90, Top: 0, Bottom: 30, Text: "A", R: 0, C: 0}, + {X0: 110, X1: 190, Top: 0, Bottom: 30, Text: "B", R: 0, C: 1}, + {X0: 10, X1: 90, Top: 35, Bottom: 65, Text: "C", R: 1, C: 0}, + {X0: 110, X1: 190, Top: 35, Bottom: 65, Text: "D", R: 1, C: 1}, + } + // Sort first + sort.Slice(boxes, func(i, j int) bool { + if boxes[i].R != boxes[j].R { + return boxes[i].R < boxes[j].R + } + if boxes[i].Top != boxes[j].Top { + return boxes[i].Top < boxes[j].Top + } + return boxes[i].X0 < boxes[j].X0 + }) + + rowMap, compressed := compressRowIndices(boxes) + if compressed != 1 { // 0-based, 2 rows → max index 1 + t.Errorf("compressed = %d, want 1", compressed) + } + if rowMap[0] != 0 || rowMap[1] != 1 { + t.Errorf("rowMap mapping incorrect: %v", rowMap) + } +} + +func TestCollectBoxesPerRow(t *testing.T) { + boxes := []pdf.TextBox{ + {X0: 10, X1: 90, Top: 0, Bottom: 30, Text: "A", R: 0, C: 0}, + {X0: 110, X1: 190, Top: 0, Bottom: 30, Text: "B", R: 0, C: 1}, + {X0: 10, X1: 90, Top: 35, Bottom: 65, Text: "C", R: 1, C: 0}, + {X0: 110, X1: 190, Top: 35, Bottom: 65, Text: "D", R: 1, C: 1}, + } + // Sort first + sort.Slice(boxes, func(i, j int) bool { + if boxes[i].R != boxes[j].R { + return boxes[i].R < boxes[j].R + } + if boxes[i].Top != boxes[j].Top { + return boxes[i].Top < boxes[j].Top + } + return boxes[i].X0 < boxes[j].X0 + }) + + rowMap, _ := compressRowIndices(boxes) + cmap, maxCols := collectBoxesPerRow(boxes, rowMap) + + if len(cmap) != 2 { + t.Errorf("cmap has %d rows, want 2", len(cmap)) + } + if cmap[0][0].txt != "A" || cmap[0][1].txt != "B" { + t.Errorf("row 0 incorrect: %v, %v", cmap[0][0], cmap[0][1]) + } + if cmap[1][0].txt != "C" || cmap[1][1].txt != "D" { + t.Errorf("row 1 incorrect: %v, %v", cmap[1][0], cmap[1][1]) + } + if maxCols[0] != 1 || maxCols[1] != 1 { + t.Errorf("maxCols incorrect: %v", maxCols) + } +} + +func TestCollectBoxesPerRow_MergeSameCell(t *testing.T) { + boxes := []pdf.TextBox{ + {X0: 10, X1: 50, Top: 0, Bottom: 30, Text: "Hello", R: 0, C: 0}, + {X0: 50, X1: 90, Top: 0, Bottom: 30, Text: "World", R: 0, C: 0}, + } + // Sort first + sort.Slice(boxes, func(i, j int) bool { + if boxes[i].R != boxes[j].R { + return boxes[i].R < boxes[j].R + } + if boxes[i].Top != boxes[j].Top { + return boxes[i].Top < boxes[j].Top + } + return boxes[i].X0 < boxes[j].X0 + }) + + rowMap, _ := compressRowIndices(boxes) + cmap, _ := collectBoxesPerRow(boxes, rowMap) + + if cmap[0][0].txt != "Hello World" { + t.Errorf("merged text incorrect: %q", cmap[0][0].txt) + } +} + +func TestCompressColIndices(t *testing.T) { + boxes := []pdf.TextBox{ + {X0: 10, X1: 90, Top: 0, Bottom: 30, Text: "A", R: 0, C: 0}, + {X0: 110, X1: 190, Top: 0, Bottom: 30, Text: "B", R: 0, C: 1}, + {X0: 10, X1: 90, Top: 35, Bottom: 65, Text: "C", R: 1, C: 0}, + {X0: 110, X1: 190, Top: 35, Bottom: 65, Text: "D", R: 1, C: 1}, + } + // Sort first + sort.Slice(boxes, func(i, j int) bool { + if boxes[i].R != boxes[j].R { + return boxes[i].R < boxes[j].R + } + if boxes[i].Top != boxes[j].Top { + return boxes[i].Top < boxes[j].Top + } + return boxes[i].X0 < boxes[j].X0 + }) + + rowMap, compressed := compressRowIndices(boxes) + cCompressed, cMaxCol := compressColIndices(boxes, rowMap, compressed) + + if cCompressed[0][0] != 0 || cCompressed[0][1] != 1 { + t.Errorf("row 0 compression incorrect: %v", cCompressed[0]) + } + if cCompressed[1][0] != 0 || cCompressed[1][1] != 1 { + t.Errorf("row 1 compression incorrect: %v", cCompressed[1]) + } + if cMaxCol[0] != 1 || cMaxCol[1] != 1 { + t.Errorf("cMaxCol incorrect: %v", cMaxCol) + } +} + +func TestCompressColIndices_OverlapMerge(t *testing.T) { + boxes := []pdf.TextBox{ + {X0: 10, X1: 100, Top: 0, Bottom: 30, Text: "A", R: 0, C: 0}, + {X0: 90, X1: 190, Top: 0, Bottom: 30, Text: "B", R: 0, C: 1}, // Overlaps with A + } + // Sort first + sort.Slice(boxes, func(i, j int) bool { + if boxes[i].R != boxes[j].R { + return boxes[i].R < boxes[j].R + } + if boxes[i].Top != boxes[j].Top { + return boxes[i].Top < boxes[j].Top + } + return boxes[i].X0 < boxes[j].X0 + }) + + rowMap, compressed := compressRowIndices(boxes) + cCompressed, cMaxCol := compressColIndices(boxes, rowMap, compressed) + + if cCompressed[0][0] != 0 || cCompressed[0][1] != 0 { + t.Errorf("overlap merge incorrect: %v", cCompressed[0]) + } + // After fix: cMaxCol tracks unique compressed columns, not original C keys. + // 2 overlapping boxes → 1 compressed column → cMaxCol[0] == 0 + if cMaxCol[0] != 0 { + t.Errorf("cMaxCol incorrect: got %d, want 0 (2 overlapping boxes → 1 compressed column)", cMaxCol[0]) + } +} + +func TestBuildGrid(t *testing.T) { + boxes := []pdf.TextBox{ + {X0: 10, X1: 90, Top: 0, Bottom: 30, Text: "A", R: 0, C: 0}, + {X0: 110, X1: 190, Top: 0, Bottom: 30, Text: "B", R: 0, C: 1}, + {X0: 10, X1: 90, Top: 35, Bottom: 65, Text: "C", R: 1, C: 0}, + {X0: 110, X1: 190, Top: 35, Bottom: 65, Text: "D", R: 1, C: 1}, + } + // Sort first + sort.Slice(boxes, func(i, j int) bool { + if boxes[i].R != boxes[j].R { + return boxes[i].R < boxes[j].R + } + if boxes[i].Top != boxes[j].Top { + return boxes[i].Top < boxes[j].Top + } + return boxes[i].X0 < boxes[j].X0 + }) + + rowMap, compressed := compressRowIndices(boxes) + cmap, _ := collectBoxesPerRow(boxes, rowMap) + cCompressed, cMaxCol := compressColIndices(boxes, rowMap, compressed) + + rows := buildGrid(cmap, cCompressed, cMaxCol, compressed) + + if len(rows) != 2 { + t.Fatalf("expected 2 rows, got %d", len(rows)) + } + if len(rows[0]) != 2 || len(rows[1]) != 2 { + t.Errorf("expected 2 cols each, got %d and %d", len(rows[0]), len(rows[1])) + } + if rows[0][0].Text != "A" || rows[0][1].Text != "B" { + t.Errorf("row 0 wrong: %q %q", rows[0][0].Text, rows[0][1].Text) + } + if rows[1][0].Text != "C" || rows[1][1].Text != "D" { + t.Errorf("row 1 wrong: %q %q", rows[1][0].Text, rows[1][1].Text) + } +} diff --git a/internal/deepdoc/parser/pdf/table/table_construct.go b/internal/deepdoc/parser/pdf/table/table_construct.go index bb7509e95d..ca785b30ce 100644 --- a/internal/deepdoc/parser/pdf/table/table_construct.go +++ b/internal/deepdoc/parser/pdf/table/table_construct.go @@ -1,8 +1,6 @@ package table import ( - "fmt" - "html" "math" "regexp" "sort" @@ -11,119 +9,6 @@ import ( pdf "ragflow/internal/deepdoc/parser/pdf/type" ) -// ── construct table ───────────────────────────────────────────────────── - -// MergeTablesAcrossPages merges TableItems on consecutive pages with -// overlapping X and close Y proximity. Matches Python's -// _extract_table_figure table merge (pdf_parser.py:1061-1080). -func MergeTablesAcrossPages(tables []pdf.TableItem, medianHeights map[int]float64) []pdf.TableItem { - if len(tables) <= 1 { - return tables - } - // Sort by position for deterministic adjacency. - type indexed struct { - idx int - pg int - top float64 - } - var items []indexed - for i, tbl := range tables { - if len(tbl.Positions) == 0 { - continue - } - p := tbl.Positions[0] - pg := 0 - if len(p.PageNumbers) > 0 { - pg = p.PageNumbers[0] - } - items = append(items, indexed{i, pg, p.Top}) - } - sort.Slice(items, func(a, b int) bool { - if items[a].pg != items[b].pg { - return items[a].pg < items[b].pg - } - return items[a].top < items[b].top - }) - - merged := make([]bool, len(tables)) - var result []pdf.TableItem - - for _, it := range items { - if merged[it.idx] { - continue - } - anchor := tables[it.idx] - merged[it.idx] = true - - // Python nomerge_lout_no: tables whose box is followed by a - // caption/title/reference should not be merged cross-page. - if anchor.NoMerge { - result = append(result, anchor) - continue - } - - anchorPg := it.pg - anchorBott := anchor.Positions[0].Bottom - - // Look for consecutive-page continuations. - for _, jt := range items { - if merged[jt.idx] || jt.pg <= anchorPg { - continue - } - // Python nomerge_lout_no: skip continuation candidates - // tagged as no-merge. - if tables[jt.idx].NoMerge { - continue - } - if jt.pg-anchorPg > 1 { - break // pages must be consecutive - } - if len(tables[jt.idx].Positions) == 0 { - continue - } - bp := tables[jt.idx].Positions[0] - bpg := 0 - if len(bp.PageNumbers) > 0 { - bpg = bp.PageNumbers[0] - } - if bpg != anchorPg+1 { - continue - } - // Check X overlap. - ap := anchor.Positions[0] - if ap.Right < bp.Left || bp.Right < ap.Left { - continue - } - // Check Y proximity: page 1 table top should be close below - // page 0 table bottom. Python: y_dis ≤ mh * 23. - mh := 10.0 - if medianHeights != nil { - if h, ok := medianHeights[anchorPg]; ok && h > 0 { - mh = h - } - } - yDis := (bp.Top + bp.Bottom - anchorBott - ap.Bottom) / 2 - if yDis > mh*23 { - continue - } - // Merge: combine cells and positions. - anchor.Cells = append(anchor.Cells, tables[jt.idx].Cells...) - anchor.Positions = append(anchor.Positions, tables[jt.idx].Positions...) - if tables[jt.idx].Caption != "" { - if anchor.Caption != "" { - anchor.Caption += " " - } - anchor.Caption += tables[jt.idx].Caption - } - merged[jt.idx] = true - anchorPg = bpg - anchorBott = bp.Bottom - } - result = append(result, anchor) - } - return result -} - // 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. @@ -157,8 +42,8 @@ func ConstructTable(cells []pdf.TSRCell, boxes []pdf.TextBox, caption string, it 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). + // 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 @@ -212,342 +97,6 @@ func BoxHeaderSet(rows [][]pdf.TSRCell, boxes []pdf.TextBox) map[int]bool { return hdrs } -func HasAnyText(cells []pdf.TSRCell) bool { - for _, c := range cells { - if strings.TrimSpace(c.Text) != "" { - return true - } - } - return false -} - -// groupBoxesByRC groups text boxes into a cell grid by R/C annotations. -// Matches Python's construct_table: sort by R, merge nearby rows by Y proximity, -// sort by C within each row, merge nearby columns by X proximity. -func GroupBoxesByRC(boxes []pdf.TextBox) [][]pdf.TSRCell { - if len(boxes) == 0 { - return nil - } - // If no real R/C annotations (maxR <= 0), fall back to YX coordinate - // grouping — matching Python's construct_table when all R=-1. - maxR := 0 - for _, b := range boxes { - if b.R > maxR { - maxR = b.R - } - } - if maxR <= 0 { - return GroupBoxesByYX(boxes) - } - // Sort by R index first (Python: sort_R_firstly), then Y, then X. - sort.Slice(boxes, func(i, j int) bool { - if boxes[i].R != boxes[j].R { - return boxes[i].R < boxes[j].R - } - if boxes[i].Top != boxes[j].Top { - return boxes[i].Top < boxes[j].Top - } - return boxes[i].X0 < boxes[j].X0 - }) - - // Compress R indices: Python's sort_R_firstly grouping. - // R differs → always a new row. Same R + Y gap → also new row. - rowMap := make(map[int]int) // original R → compressed row index - compressed := 0 - rowMap[boxes[0].R] = 0 - lastR := boxes[0].R - btm := boxes[0].Bottom - for i := 1; i < len(boxes); i++ { - // Python: b["R"] != last_R → new row. - // Same R → always same row (Python doesn't check Y for same R). - if boxes[i].R != lastR { - compressed++ - rowMap[boxes[i].R] = compressed - lastR = boxes[i].R - btm = boxes[i].Bottom - } else { - // Same R → same physical row. - rowMap[boxes[i].R] = compressed - btm = (btm + boxes[i].Bottom) / 2.0 - } - } - - // Collect boxes per row, sort by C within each row. - type rb struct { - row, col int - txt string - x0, y0, x1, y1 float64 - label string - } - cmap := make(map[int]map[int]*rb) // row → col → entry - maxCols := make(map[int]int) - for _, b := range boxes { - t := strings.TrimSpace(b.Text) - // Keep boxes with SP/H annotations even if text is empty — - // their coordinates are needed for colspan/rowspan calculation. - if t == "" && b.H <= 0 && b.SP <= 0 { - continue - } - r := rowMap[b.R] - c := b.C - if cmap[r] == nil { - cmap[r] = make(map[int]*rb) - } - x0, y0, x1, y1, label := cellPosFromBox(b) - if v, ok := cmap[r][c]; ok { - v.txt += " " + t - // Merge spanning coordinates (use widest extent). - if b.H > 0 || b.SP > 0 { - v.label = cellLabelFromBox(b) - if v.x0 > x0 { - v.x0 = x0 - } - if v.y0 > y0 { - v.y0 = y0 - } - if v.x1 < x1 { - v.x1 = x1 - } - if v.y1 < y1 { - v.y1 = y1 - } - } - } else { - cmap[r][c] = &rb{r, c, t, x0, y0, x1, y1, label} - } - if c > maxCols[r] { - maxCols[r] = c - } - } - - // Compress C indices per row: sort boxes by X0 within the row, - // group disjoint X ranges into separate columns. This is equivalent - // to Python's sort_C_firstly but uses X0 ordering instead of C labels. - cCompressed := make(map[int]map[int]int) // row → (original C → compressed col) - cMaxCol := make(map[int]int) - for ri := 0; ri <= compressed; ri++ { - rowEntries := cmap[ri] - if rowEntries == nil { - continue - } - // Collect all boxes in this row, sorted by X0. - type rowBox struct { - c, idx int - x0, x1 float64 - txt string - } - var rowBoxes []rowBox - for i, b := range boxes { - if rowMap[b.R] == ri && (strings.TrimSpace(b.Text) != "" || b.H > 0 || b.SP > 0) { - rowBoxes = append(rowBoxes, rowBox{c: b.C, idx: i, x0: b.X0, x1: b.X1, txt: b.Text}) - } - } - sort.Slice(rowBoxes, func(i, j int) bool { return rowBoxes[i].x0 < rowBoxes[j].x0 }) - // Assign compressed column by X-order (disjoint X → new col). - cMap := make(map[int]int) // original C → compressed col - right := 0.0 - for _, rb := range rowBoxes { - if len(cMap) == 0 || rb.x0 >= right { - cc := len(cMap) - cMap[rb.c] = cc - right = rb.x1 - } else { - // Overlapping X → merge into last column. - cMap[rb.c] = len(cMap) - 1 - if rb.x1 > right { - right = rb.x1 - } - } - } - cCompressed[ri] = cMap - cMaxCol[ri] = len(cMap) - 1 - } - - // Build grid. - rows := make([][]pdf.TSRCell, compressed+1) - for ri := 0; ri <= compressed; ri++ { - maxC := cMaxCol[ri] - rows[ri] = make([]pdf.TSRCell, maxC+1) - for ci, v := range cmap[ri] { - cci := cCompressed[ri][ci] - if cci <= maxC { - rows[ri][cci].Text = v.txt - rows[ri][cci].X0 = v.x0 - rows[ri][cci].Y0 = v.y0 - rows[ri][cci].X1 = v.x1 - rows[ri][cci].Y1 = v.y1 - rows[ri][cci].Label = v.label - } - } - } - return rows -} - -// cellPosFromBox returns the position coordinates and label for a cell -// derived from a text box. Header cells use HLeft/HRight/HTop/HBott -// for spanning-aware positions; regular cells use the box's own bounds. -func cellPosFromBox(b pdf.TextBox) (x0, y0, x1, y1 float64, label string) { - x0, y0, x1, y1 = b.X0, b.Top, b.X1, b.Bottom - if b.H > 0 { - label = "table header" - if b.HLeft != 0 || b.HRight != 0 { - if b.HLeft != 0 { - x0 = b.HLeft - } - if b.HRight != 0 { - x1 = b.HRight - } - } - if b.HTop != 0 { - y0 = b.HTop - } - if b.HBott != 0 { - y1 = b.HBott - } - } else if b.SP > 0 { - label = "table spanning cell" - } - return -} - -// cellLabelFromBox returns the TSR label for a box based on H/SP annotations. -// Used when merging multiple boxes into one cell — preserves the spanning label. -func cellLabelFromBox(b pdf.TextBox) string { - if b.H > 0 { - return "table header" - } - if b.SP > 0 { - return "table spanning cell" - } - return "" -} - -// groupBoxesByYX groups boxes into a cell grid by Y/X coordinates, -// matching Python's construct_table which uses sort_R_firstly and -// sort_C_firstly when R/C annotations are absent. -// This is test-only — used by table_parity_test.go to verify pipeline -// parity with Python boxes that lack R/C annotations. -func GroupBoxesByYX(boxes []pdf.TextBox) [][]pdf.TSRCell { - if len(boxes) == 0 { - return nil - } - // Sort by (page, top, x0) — same as Python sort_R_firstly with R=-1. - sort.Slice(boxes, func(i, j int) bool { - if boxes[i].PageNumber != boxes[j].PageNumber { - return boxes[i].PageNumber < boxes[j].PageNumber - } - if boxes[i].Top != boxes[j].Top { - return boxes[i].Top < boxes[j].Top - } - return boxes[i].X0 < boxes[j].X0 - }) - - // Group into rows by Y proximity (Python's row grouping). - type rowGroup struct { - boxes []pdf.TextBox - top, btm float64 - } - var rowGroups []rowGroup - rowGroups = append(rowGroups, rowGroup{boxes: []pdf.TextBox{boxes[0]}, top: boxes[0].Top, btm: boxes[0].Bottom}) - for i := 1; i < len(boxes); i++ { - prev := &rowGroups[len(rowGroups)-1] - // Python: same row if top < prev.btm (Y overlaps) and same page. - if boxes[i].PageNumber == prev.boxes[0].PageNumber && boxes[i].Top < prev.btm { - prev.boxes = append(prev.boxes, boxes[i]) - if boxes[i].Top < prev.top { - prev.top = boxes[i].Top - } - if boxes[i].Bottom > prev.btm { - prev.btm = boxes[i].Bottom - } - } else { - rowGroups = append(rowGroups, rowGroup{boxes: []pdf.TextBox{boxes[i]}, top: boxes[i].Top, btm: boxes[i].Bottom}) - } - } - - // Within each row, group into columns by X proximity. - rows := make([][]pdf.TSRCell, len(rowGroups)) - for ri, rg := range rowGroups { - // Sort by X0. - sort.Slice(rg.boxes, func(i, j int) bool { return rg.boxes[i].X0 < rg.boxes[j].X0 }) - // Group by X overlap. - var cols []struct { - boxes []pdf.TextBox - x1 float64 - } - cols = append(cols, struct { - boxes []pdf.TextBox - x1 float64 - }{boxes: []pdf.TextBox{rg.boxes[0]}, x1: rg.boxes[0].X1}) - for i := 1; i < len(rg.boxes); i++ { - prev := &cols[len(cols)-1] - if rg.boxes[i].X0 < prev.x1 { - prev.boxes = append(prev.boxes, rg.boxes[i]) - if rg.boxes[i].X1 > prev.x1 { - prev.x1 = rg.boxes[i].X1 - } - } else { - cols = append(cols, struct { - boxes []pdf.TextBox - x1 float64 - }{boxes: []pdf.TextBox{rg.boxes[i]}, x1: rg.boxes[i].X1}) - } - } - rows[ri] = make([]pdf.TSRCell, len(cols)) - for ci, col := range cols { - var sb strings.Builder - for _, b := range col.boxes { - t := strings.TrimSpace(b.Text) - if t == "" { - continue - } - if sb.Len() > 0 { - sb.WriteByte(' ') - } - sb.WriteString(t) - } - rows[ri][ci].Text = sb.String() - } - } - return rows -} - -func BoxesHaveAnnotations(boxes []pdf.TextBox) bool { - maxR, maxC := 0, 0 - for _, b := range boxes { - if b.R > maxR { - maxR = b.R - } - if b.C > maxC { - maxC = b.C - } - } - // True if at least 2 rows or 2 cols (R/C are 0-based, so maxR>0 means ≥2 rows). - return maxR > 0 || maxC > 0 -} - -func HasText(rows [][]pdf.TSRCell) bool { - for _, row := range rows { - for _, c := range row { - if strings.TrimSpace(c.Text) != "" { - return true - } - } - } - return false -} - -func RowsToStrings(rows [][]pdf.TSRCell) [][]string { - out := make([][]string, len(rows)) - for ri, row := range rows { - out[ri] = make([]string, len(row)) - for ci, c := range row { - out[ri][ci] = c.Text - } - } - return out -} - // 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. @@ -578,7 +127,9 @@ func FillCellTextFromAnnotations(rows [][]pdf.TSRCell, boxes []pdf.TextBox) { 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 }) + 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, " ")) @@ -613,8 +164,10 @@ func tableRegionBox(tbl *pdf.TableItem, ref *pdf.TextBox, html string) pdf.TextB // 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, + X0: tbl.RegionLeft, + X1: tbl.RegionRight, + Top: tbl.RegionTop, + Bottom: tbl.RegionBottom, Text: html, PageNumber: pg, LayoutType: pdf.LayoutTypeTable, @@ -623,7 +176,10 @@ func tableRegionBox(tbl *pdf.TableItem, ref *pdf.TextBox, html string) pdf.TextB // 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, + X0: x0, + X1: x1, + Top: top, + Bottom: bot, Text: html, PageNumber: pg, LayoutType: pdf.LayoutTypeTable, @@ -651,213 +207,10 @@ func minRectangleDistance(left1, right1, top1, bottom1, left2, right2, top2, bot return math.Sqrt(dx*dx + dy*dy) } -func RowsToHTML(rows [][]pdf.TSRCell, caption string, headerRows map[int]bool, spanInfo map[[2]int][2]int, covered map[[2]int]bool) string { - var b strings.Builder - b.WriteString("
") - if caption != "" { - b.WriteString("") - } - for ri, row := range rows { - b.WriteString("") - for ci, cell := range row { - if covered[[2]int{ri, ci}] { - continue - } - tag := "td" - if headerRows[ri] { - tag = "th" - } - b.WriteString("<") - b.WriteString(tag) - sp := "" - if s, ok := spanInfo[[2]int{ri, ci}]; ok { - if s[0] > 1 { - sp = fmt.Sprintf("colspan=%d", s[0]) - } - if s[1] > 1 { - if sp != "" { - sp += " " - } - sp += fmt.Sprintf("rowspan=%d", s[1]) - } - } - if sp != "" { - b.WriteString(" ") - b.WriteString(sp) - } - b.WriteString(" >") - b.WriteString(cell.Text) - b.WriteString("") - } - b.WriteString("") - } - b.WriteString("
") - b.WriteString(caption) - b.WriteString("
") - return b.String() -} +// Orphan column/row cleanup (Python: construct_table lines 256-368) -// SimpleRowsToHTML converts plain string-based table data to an HTML table. -// The first row is treated as a header (). Used by DOCX, XLSX, PPTX, -// and HTML parsers that produce [][]string directly. -func SimpleRowsToHTML(rows [][]string) string { - if len(rows) == 0 { - return "
" - } - nCols := 0 - for _, row := range rows { - if len(row) > nCols { - nCols = len(row) - } - } - var b strings.Builder - b.WriteString("") - for ri, row := range rows { - b.WriteString("") - tag := "td" - if ri == 0 { - tag = "th" - } - for ci := 0; ci < nCols; ci++ { - text := "" - if ci < len(row) { - text = row[ci] - } - b.WriteString("<") - b.WriteString(tag) - b.WriteString(" >") - b.WriteString(html.EscapeString(text)) - b.WriteString("") - } - b.WriteString("") - } - b.WriteString("
") - return b.String() -} - -// Span computation (Python: __cal_spans) ── - -// calSpans computes colspan and rowspan for spanning cells in the grid. -// Returns spanInfo (row,col → colspan,rowspan) and covered (cells hidden by spans). -// Matches Python's __cal_spans (table_structure_recognizer.py:535). -// flattenGrid flattens a 2D grid into a 1D slice for fillCellTextFromBoxes. -func FlattenGrid(grid [][]pdf.TSRCell) []pdf.TSRCell { - n := 0 - for _, row := range grid { - n += len(row) - } - flat := make([]pdf.TSRCell, 0, n) - for _, row := range grid { - flat = append(flat, row...) - } - return flat -} - -func CalSpans(rows [][]pdf.TSRCell) (map[[2]int][2]int, map[[2]int]bool) { - spanInfo := make(map[[2]int][2]int) - covered := make(map[[2]int]bool) - if len(rows) == 0 || len(rows[0]) == 0 { - return spanInfo, covered - } - - // Compute column center positions. - nCols := len(rows[0]) - colLeft := make([]float64, nCols) - colRight := make([]float64, nCols) - for j := 0; j < nCols; j++ { - colLeft[j] = 1e9 - colRight[j] = -1e9 - } - nRows := len(rows) - rowTop := make([]float64, nRows) - rowBott := make([]float64, nRows) - for i := 0; i < nRows; i++ { - rowTop[i] = 1e9 - rowBott[i] = -1e9 - } - - for i, row := range rows { - for j, cell := range row { - if j >= nCols { - continue - } - // Exclude spanning cells from column/row boundary calculations. - // Use label-based detection (O(1), no dependency on column midpoints). - if strings.Contains(cell.Label, "spanning") { - continue - } - if cell.X0 < colLeft[j] { - colLeft[j] = cell.X0 - } - if cell.X1 > colRight[j] { - colRight[j] = cell.X1 - } - if cell.Y0 < rowTop[i] { - rowTop[i] = cell.Y0 - } - if cell.Y1 > rowBott[i] { - rowBott[i] = cell.Y1 - } - } - } - - // For each spanning cell, compute how many cols/rows it covers. - for i, row := range rows { - for j, cell := range row { - if j >= nCols || covered[[2]int{i, j}] { - continue - } - // Skip cells without position data (they can't span). - if cell.X0 == 0 && cell.X1 == 0 && cell.Y0 == 0 && cell.Y1 == 0 { - continue - } - cs, rs := 1, 1 - // Count columns whose center is inside this cell's X range. - for k := j + 1; k < nCols; k++ { - // Skip columns with no non-spanning cells (initial values unchanged). - if colLeft[k] == 1e9 && colRight[k] == -1e9 { - continue - } - colCenter := (colLeft[k] + colRight[k]) / 2 - if colCenter >= cell.X0 && colCenter <= cell.X1 { - cs++ - } - } - // Count rows whose center is inside this cell's Y range. - for k := i + 1; k < nRows; k++ { - // Skip rows with no non-spanning cells. - if rowTop[k] == 1e9 && rowBott[k] == -1e9 { - continue - } - rowCenter := (rowTop[k] + rowBott[k]) / 2 - if rowCenter >= cell.Y0 && rowCenter <= cell.Y1 { - rs++ - } - } - if cs > 1 || rs > 1 { - spanInfo[[2]int{i, j}] = [2]int{cs, rs} - // Mark covered cells. - for ri := i; ri < i+rs && ri < nRows; ri++ { - for cj := j; cj < j+cs && cj < nCols; cj++ { - if ri != i || cj != j { - covered[[2]int{ri, cj}] = true - } - } - } - } - } - } - return spanInfo, covered -} - -// ── 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. +// 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 @@ -865,84 +218,121 @@ func CleanupOrphanColumns(rows [][]pdf.TSRCell) [][]pdf.TSRCell { nCols := len(rows[0]) j := 0 -colLoop: for j < nCols { - e, ii := 0, 0 - for i := range rows { - if j < len(rows[i]) && strings.TrimSpace(rows[i][j].Text) != "" { - e++ - ii = i - } - if e > 1 { - j++ - continue colLoop - } - } - // Column j has only one non-empty cell at row ii. - // Check if adjacent columns have text for this row. - f := (j > 0 && j-1 < len(rows[ii]) && strings.TrimSpace(rows[ii][j-1].Text) != "") || j == 0 - ff := (j+1 < len(rows[ii]) && strings.TrimSpace(rows[ii][j+1].Text) != "") || j+1 >= len(rows[ii]) - if f && ff { - // Both adjacent columns are ok for merging — but this means - // there's text on both sides, keep column. + // Step 1: Count non-empty cells in column + e, ii := countNonEmptyCells(rows, j) + if e > 1 { j++ continue } - // Determine which side to merge into. - left := 1e9 - right := 1e9 - if j > 0 && !f { - for i := range rows { - if j-1 < len(rows[i]) && strings.TrimSpace(rows[i][j-1].Text) != "" { - // Distance from orphan cell to left neighbor. - if d := rows[ii][j].X0 - rows[i][j-1].X1; d < left { - left = d - } - } - } - } - if j+1 < nCols && !ff { - for i := range rows { - if j+1 < len(rows[i]) && strings.TrimSpace(rows[i][j+1].Text) != "" { - if d := rows[i][j+1].X0 - rows[ii][j].X1; d < right { - right = d - } - } - } + // Step 2: Check adjacent columns + hasLeftText, hasRightText := checkAdjacentColumns(rows, j, ii) + if hasLeftText && hasRightText { + j++ + continue } - if left < right && j > 0 { - // Merge into left column. - for i := range rows { - if j-1 < len(rows[i]) && j < len(rows[i]) { - if rows[i][j-1].Text == "" { - rows[i][j-1].Text = rows[i][j].Text - } else if rows[i][j].Text != "" { - rows[i][j-1].Text += " " + rows[i][j].Text - } - } - } + // 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 { - // Merge into right column. - for i := range rows { - if j < len(rows[i]) && j+1 < len(rows[i]) { - if rows[i][j+1].Text == "" { - rows[i][j+1].Text = rows[i][j].Text - } else if rows[i][j].Text != "" { - rows[i][j+1].Text = rows[i][j].Text + " " + rows[i][j+1].Text - } - } - } - } - // Remove column j. - for i := range rows { - if j < len(rows[i]) { - rows[i] = append(rows[i][:j], rows[i][j+1:]...) - } + 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 +} diff --git a/internal/deepdoc/parser/pdf/table/table_construct_test.go b/internal/deepdoc/parser/pdf/table/table_construct_test.go index e93c64d6b8..86de5e956d 100644 --- a/internal/deepdoc/parser/pdf/table/table_construct_test.go +++ b/internal/deepdoc/parser/pdf/table/table_construct_test.go @@ -1,9 +1,11 @@ + package table import ( - pdf "ragflow/internal/deepdoc/parser/pdf/type" "strings" "testing" + + pdf "ragflow/internal/deepdoc/parser/pdf/type" ) func TestCellTexts(t *testing.T) { @@ -17,8 +19,6 @@ func TestCellTexts(t *testing.T) { } } -// ── constructTable unit tests ───────────────────────────────────────── - func TestConstructTable_Simple3x2(t *testing.T) { // 3 columns × 2 rows — cells pre-filled (simulating extractTableBoxesFromImage). cells := []pdf.TSRCell{ @@ -140,10 +140,6 @@ func TestConstructTable_YBasedFallback(t *testing.T) { } } -// TestExtractTableAndReplace_CellTextFilled verifies that extractTableAndReplace -// fills cell text correctly with realistic coordinate transforms (Scale=3, CropOff≠0). -// This simulates the real pipeline where TSR cells are in crop pixel space and -// post-merge boxes are in PDF point space. func TestExtractTableAndReplace_CellTextFilled(t *testing.T) { // Simulate 公司差旅费 page 0 table coordinates. // DLA region: X0=217, X1=1584, Y0=985, Y1=1599 at 216 DPI → PDF: 72-528 x 328-533 @@ -168,9 +164,9 @@ func TestExtractTableAndReplace_CellTextFilled(t *testing.T) { // Cells pre-filled (extractTableBoxesFromImage already ran fillText + OCR). cells := []pdf.TSRCell{ {X0: 35, Y0: 441, X1: 456, Y1: 500, Text: "标职务", Label: "table row"}, - {X0: 460, Y0: 442, X1: 630, Y1: 500, Text: "飞机", Label: "table row"}, + {X0: 460, Y0: 441, X1: 630, Y1: 500, Text: "飞机", Label: "table row"}, {X0: 35, Y0: 501, X1: 456, Y1: 560, Text: "公司级领导", Label: "table row"}, - {X0: 460, Y0: 502, X1: 630, Y1: 560, Text: "经济舱位", Label: "table row"}, + {X0: 460, Y0: 501, X1: 630, Y1: 560, Text: "经济舱位", Label: "table row"}, } tables := []pdf.TableItem{{ @@ -208,8 +204,6 @@ func TestExtractTableAndReplace_CellTextFilled(t *testing.T) { } } -// TestConstructTable_FromBoxesRC builds HTML directly from boxes with R/C -// annotations, matching Python's construct_table. No cells needed for text. func TestConstructTable_FromBoxesRC(t *testing.T) { // Boxes with R (row) and C (col) annotations — like the output of // annotateTableBoxes after layout cleanup. @@ -245,10 +239,6 @@ func TestConstructTable_FromBoxesRC(t *testing.T) { t.Logf("HTML: %s", html) } -// TestFillCellTextFromBoxes_RCAnnotations fills text via R/C when spatial -// overlap is marginal. Real-world TSR cells and pdf_oxide boxes have pixel-level -// offsets — R/C annotations (set by annotateTableBoxes) are the Python-equivalent -// way to assign boxes to cells regardless of coordinate deviations. func TestFillCellTextFromBoxes_RCAnnotations(t *testing.T) { // Cells with real-world coordinate offsets (box shifted by 2px from cell). // Spatial overlap <30% for the shifted case — fillCellTextFromBoxes fails. @@ -260,19 +250,19 @@ func TestFillCellTextFromBoxes_RCAnnotations(t *testing.T) { } // Boxes have R/C annotations but their spatial overlap with cell rects - // is marginal (real-world scenario). R/C path should still fill text. + // is marginal (real-world scenario). R/C path should still fill text. boxes := []pdf.TextBox{ - {X0: 12, X1: 198, Top: 12, Bottom: 48, Text: "A", R: 0, C: 0}, // overlap ~92% → OK + {X0: 12, X1: 198, Top: 12, Bottom: 48, Text: "A", R: 0, C: 0}, // overlap ~92% → OK {X0: 215, X1: 395, Top: 12, Bottom: 48, Text: "B", R: 0, C: 1}, // overlap ~90% → OK - {X0: 12, X1: 198, Top: 58, Bottom: 92, Text: "C", R: 1, C: 0}, // overlap ~92% → OK - {X0: 215, X1: 350, Top: 58, Bottom: 90, Text: "D", R: 1, C: 1}, // overlap ~50% → MARGINAL + {X0: 12, X1: 198, Top: 58, Bottom: 92, Text: "C", R: 1, C: 0}, // overlap ~92% → OK + {X0: 215, X1: 350, Top: 58, Bottom: 92, Text: "D", R: 1, C: 1}, // overlap ~50% → MARGINAL } // This SHOULD fill all 4 cells via R/C, but spatial-only may fail on D. FillCellTextFromBoxes(cells, boxes) // When spatial overlap is marginal (box "D" at 50%), fillCellTextFromBoxes - // may still match because cell is empty (0.3 threshold). But the real + // may still match because cell is empty (0.3 threshold). But the real // problem is that fillCellTextFromBoxes depends on coordinates, while // R/C annotations don't. hasText := false @@ -293,7 +283,7 @@ func TestFillCellTextFromBoxes_RCAnnotations(t *testing.T) { {X0: 10, Y0: 55, X1: 200, Y1: 95}, {X0: 210, Y0: 55, X1: 400, Y1: 95}, } - rows := [][]pdf.TSRCell{{cells2[0], cells2[1]}, {cells2[2], cells2[3]}} + rows := GroupTSRCellsToRows(cells2) FillCellTextFromAnnotations(rows, boxes) if rows[0][0].Text != "A" { @@ -310,9 +300,6 @@ func TestFillCellTextFromBoxes_RCAnnotations(t *testing.T) { } } -// TestConstructTable_SingleRowMultiCol covers R=0 with multiple columns -// (table header pattern). boxesHaveAnnotations must detect valid annotations -// even though maxR=0. func TestConstructTable_SingleRowMultiCol(t *testing.T) { boxes := []pdf.TextBox{ {X0: 0, X1: 100, Top: 0, Bottom: 30, Text: "姓名", R: 0, C: 0}, @@ -329,9 +316,6 @@ func TestConstructTable_SingleRowMultiCol(t *testing.T) { } } -// TestConstructTable_MultiRowSingleCol covers C=0 with multiple rows -// (vertical list pattern). boxesHaveAnnotations must detect valid -// annotations even though maxC=0. func TestConstructTable_MultiRowSingleCol(t *testing.T) { boxes := []pdf.TextBox{ {X0: 0, X1: 100, Top: 0, Bottom: 30, Text: "第一行", R: 0, C: 0}, @@ -348,8 +332,6 @@ func TestConstructTable_MultiRowSingleCol(t *testing.T) { } } -// TestConstructTable_RCAfterMerge verifies that R/C annotations survive -// text merge. The merged box expands bounds but keeps the first box's R/C. func TestConstructTable_RCAfterMerge(t *testing.T) { // Simulate two adjacent fragments merged into one box. // The merged box keeps R/C from the first fragment. @@ -372,11 +354,6 @@ func TestConstructTable_RCAfterMerge(t *testing.T) { } } -// TestGroupTSRCellsToRowsLabeled_DefaultTableLabel verifies that cells with -// the real TSR default label "table" (class 0) are grouped correctly. -// The current deepDocReRowHdr regex only matches ".* (row|header)" — it misses -// the default "table" label, causing gatherTSR to return empty and forcing -// a fallback to pure Y-based grouping (which loses R/C annotations). func TestGroupTSRCellsToRowsLabeled_DefaultTableLabel(t *testing.T) { cells := []pdf.TSRCell{ {X0: 10, Y0: 0, X1: 100, Y1: 30, Label: "table"}, @@ -393,129 +370,6 @@ func TestGroupTSRCellsToRowsLabeled_DefaultTableLabel(t *testing.T) { } } -// TestGroupBoxesByRC_RDiffSplitsRows verifies that groupBoxesByRC -// creates separate rows for different R values (Python: R differs → new row). -// Even when boxes share the same Y, different R → different grid row. -func TestGroupBoxesByRC_RDiffSplitsRows(t *testing.T) { - // 6 boxes with 6 different R values → 6 rows (Python R-first splitting). - boxes := []pdf.TextBox{ - {X0: 10, X1: 90, Top: 0, Bottom: 30, Text: "A", R: 0, C: 0}, - {X0: 110, X1: 190, Top: 0, Bottom: 30, Text: "B", R: 1, C: 1}, - {X0: 210, X1: 290, Top: 0, Bottom: 30, Text: "C", R: 2, C: 2}, - {X0: 10, X1: 90, Top: 35, Bottom: 65, Text: "D", R: 3, C: 0}, - {X0: 110, X1: 190, Top: 35, Bottom: 65, Text: "E", R: 4, C: 1}, - {X0: 210, X1: 290, Top: 35, Bottom: 65, Text: "F", R: 5, C: 2}, - } - rows := GroupBoxesByRC(boxes) - // R=0,1,2,3,4,5 → 6 rows (Python: R differs → new row). - if len(rows) != 6 { - t.Fatalf("expected 6 rows (R differs → split), got %d", len(rows)) - } -} - -// TestGroupBoxesByRC_MergesCloseCols verifies that C compression works -// within each R group — merging different C values that are close in X. -func TestGroupBoxesByRC_MergesCloseCols(t *testing.T) { - // R=0 has C=0,1. R=1 has C=0,1. C compression → 2 cols each. - boxes := []pdf.TextBox{ - {X0: 10, X1: 90, Top: 0, Bottom: 30, Text: "A", R: 0, C: 0}, - {X0: 110, X1: 190, Top: 0, Bottom: 30, Text: "B", R: 0, C: 1}, - {X0: 10, X1: 90, Top: 35, Bottom: 65, Text: "C", R: 1, C: 0}, - {X0: 110, X1: 190, Top: 35, Bottom: 65, Text: "D", R: 1, C: 1}, - } - rows := GroupBoxesByRC(boxes) - if len(rows) != 2 { - t.Fatalf("expected 2 rows (R diff), got %d", len(rows)) - } - if len(rows[0]) != 2 || len(rows[1]) != 2 { - t.Errorf("expected 2 cols/row, got %d/%d", len(rows[0]), len(rows[1])) - } - if rows[0][0].Text != "A" || rows[0][1].Text != "B" { - t.Errorf("row0 wrong: %q %q", rows[0][0].Text, rows[0][1].Text) - } - if rows[1][0].Text != "C" || rows[1][1].Text != "D" { - t.Errorf("row1 wrong: %q %q", rows[1][0].Text, rows[1][1].Text) - } -} - -// TestGroupBoxesByRC_RDiffSplitsRow verifies that boxes with different R -// values are placed in separate rows even when their Y ranges overlap. -// Matches Python: R differs → new row unconditionally. -func TestGroupBoxesByRC_RDiffSplitsRow(t *testing.T) { - // R=0 and R=1 at same Y (overlapping) → two separate rows in the grid. - boxes := []pdf.TextBox{ - {X0: 10, X1: 90, Top: 0, Bottom: 30, Text: "A", R: 0, C: 0}, - {X0: 110, X1: 190, Top: 0, Bottom: 30, Text: "B", R: 1, C: 1}, - {X0: 10, X1: 90, Top: 35, Bottom: 65, Text: "C", R: 2, C: 0}, - {X0: 110, X1: 190, Top: 35, Bottom: 65, Text: "D", R: 3, C: 1}, - } - rows := GroupBoxesByRC(boxes) - // R=0,1,2,3 → 4 different R values → 4 rows (Python: R differs → new row). - if len(rows) != 4 { - t.Fatalf("expected 4 rows (R differs → split), got %d", len(rows)) - } - if rows[0][0].Text != "A" || rows[1][0].Text != "B" { - t.Errorf("row0/1 wrong: A=%q B=%q", rows[0][0].Text, rows[1][0].Text) - } -} - -// TestFillCellTextFromBoxes_RCOnly verifies that box text goes to exactly -// one cell via R/C annotations, not multiple cells via spatial overlap. -// A box overlapping two cells should only fill the one matching its R/C. -func TestFillCellTextFromBoxes_RCOnly(t *testing.T) { - cells := []pdf.TSRCell{ - {X0: 0, Y0: 0, X1: 100, Y1: 50, Label: "table"}, - {X0: 90, Y0: 0, X1: 200, Y1: 50, Label: "table"}, - } - // This box straddles cell 0 (X=0-100) and cell 1 (X=90-200). - // Spatial overlap: both match. R/C: should go to cell R=0, C=0 only. - boxes := []pdf.TextBox{ - {X0: 80, X1: 120, Top: 0, Bottom: 50, Text: "TEXT", LayoutType: "table", R: 0, C: 0}, - } - rows := GroupTSRCellsToRows(cells) - for _, b := range boxes { - t := strings.TrimSpace(b.Text) - if t == "" { - continue - } - if b.R >= 0 && b.R < len(rows) && b.C >= 0 && b.C < len(rows[b.R]) { - rows[b.R][b.C].Text = t - } - } - // Cell 0 should have text, cell 1 should NOT. - if rows[0][0].Text != "TEXT" { - t.Errorf("cell[0][0] = %q, want %q", rows[0][0].Text, "TEXT") - } - if rows[0][1].Text != "" { - t.Errorf("cell[0][1] = %q, should be empty (spatial overlap leak)", rows[0][1].Text) - } -} - -// TestRowsToHTML_HeaderRows verifies that header rows use instead of . -func TestRowsToHTML_HeaderRows(t *testing.T) { - cells := []pdf.TSRCell{ - {X0: 0, Y0: 0, X1: 100, Y1: 30, Text: "Name", Label: "table column header"}, - {X0: 101, Y0: 0, X1: 200, Y1: 30, Text: "Age", Label: "table column header"}, - {X0: 0, Y0: 35, X1: 100, Y1: 65, Text: "John", Label: "table row"}, - {X0: 101, Y0: 35, X1: 200, Y1: 65, Text: "30", Label: "table row"}, - } - // constructTable should produce for header row. - item := &pdf.TableItem{} - html := ConstructTable(cells, nil, "", item) - // Header row should use , data row . - if !strings.Contains(html, "") { - t.Errorf("expected for header row. HTML: %s", html) - } - if strings.Count(html, " cells, got %d. HTML: %s", strings.Count(html, " cells (data row), got %d", strings.Count(html, "30% each — spatial fills ALL). + // Box at X=30-270 overlaps all 3 cells (>30% each — spatial fills ALL). // With R/C, it belongs only to cell[1] (R=0, C=1). cells := []pdf.TSRCell{ {X0: 0, Y0: 0, X1: 100, Y1: 30, Label: "table"}, @@ -553,7 +404,7 @@ func TestFillCellText_RCOverSpatial(t *testing.T) { {X0: 30, X1: 270, Top: 0, Bottom: 30, Text: "TEXT", LayoutType: "table", R: 0, C: 1}, } - // Spatial fill: fills ALL overlapping cells —> duplication. + // Spatial fill: fills ALL overlapping cells → duplication. cellsCopy := make([]pdf.TSRCell, 3) copy(cellsCopy, cells) FillCellTextFromBoxes(cellsCopy, boxes) @@ -599,7 +450,7 @@ func TestIsCaptionBox(t *testing.T) { {"Table 1: Transport Levels", true}, {"图表 1: 测试", true}, {"公司领导班子成员、出差地", false}, // plain text, not caption - {"第十条到厂矿单位出差", false}, // normal paragraph + {"第十条到厂矿单位出差", false}, // normal paragraph {"", false}, } for _, tt := range tests { @@ -631,7 +482,7 @@ func TestFillCellTextFromBoxes_SkipsCaption(t *testing.T) { func TestFillCellText_RCPreventsCrossCellLeak(t *testing.T) { // Caption box at Y=0-15 overlaps BOTH cell rows (both are "empty"). - // Spatial fill: text leaks to both cells. R/C fill: only cell[0] gets text. + // Spatial fill: text leaks to cells[1]. R/C fill: only cell[0] gets text. cells := []pdf.TSRCell{ {X0: 0, Y0: 0, X1: 300, Y1: 30, Label: "table"}, {X0: 0, Y0: 35, X1: 300, Y1: 65, Label: "table"}, @@ -664,44 +515,6 @@ func TestFillCellText_RCPreventsCrossCellLeak(t *testing.T) { } } -func TestGroupBoxesByRC_FallbackToYXWhenNoAnnotations(t *testing.T) { - // When all boxes have R=-1 (Python's case: regex didn't match "table" label), - // groupBoxesByRC should fall back to YX coordinate grouping. - boxes := []pdf.TextBox{ - {X0: 10, X1: 90, Top: 0, Bottom: 30, Text: "A", R: -1, C: -1}, - {X0: 110, X1: 190, Top: 0, Bottom: 30, Text: "B", R: -1, C: -1}, - {X0: 10, X1: 90, Top: 35, Bottom: 65, Text: "C", R: -1, C: -1}, - {X0: 110, X1: 190, Top: 35, Bottom: 65, Text: "D", R: -1, C: -1}, - } - rows := GroupBoxesByRC(boxes) - // R=-1 for all → maxR = -1 → grid would be 0 rows. Must fall back to YX. - if len(rows) == 0 { - t.Fatal("groupBoxesByRC returned 0 rows when R=-1 — no YX fallback") - } - if len(rows) != 2 { - t.Errorf("expected 2 rows (Y-split), got %d", len(rows)) - } -} - -func TestRowsToHTML_Colspan(t *testing.T) { - // Box spanning 2 columns: SP annotation with HLeft/HRight covering cols 0-1. - boxes := []pdf.TextBox{ - {X0: 10, X1: 90, Top: 0, Bottom: 30, Text: "Name", R: 0, C: 0, H: 1, HLeft: 10, HRight: 190}, - {X0: 110, X1: 190, Top: 0, Bottom: 30, Text: "", R: 0, C: 1, SP: 1}, - {X0: 10, X1: 90, Top: 35, Bottom: 65, Text: "John", R: 1, C: 0}, - {X0: 110, X1: 190, Top: 35, Bottom: 65, Text: "30", R: 1, C: 1}, - } - rows := GroupBoxesByRC(boxes) - spans, covered := CalSpans(rows) - html := RowsToHTML(rows, "", nil, spans, covered) - if !strings.Contains(html, "colspan") { - t.Errorf("expected colspan attribute, got: %s", html) - } - t.Logf("HTML: %s", html) -} - -// TestStripCaptionFromCells verifies that caption-like text is cleared -// from TSR cells before the table HTML is built. func TestStripCaptionFromCells_ClearsCaptionPattern(t *testing.T) { cells := []pdf.TSRCell{ {X0: 0, Y0: 0, X1: 100, Y1: 50, Text: "表1：差旅费标准"}, @@ -718,8 +531,6 @@ func TestStripCaptionFromCells_ClearsCaptionPattern(t *testing.T) { } } -// TestStripCaptionFromCells_PreservesData verifies that non-caption -// cells are not cleared. func TestStripCaptionFromCells_PreservesData(t *testing.T) { cells := []pdf.TSRCell{ {X0: 0, Y0: 0, X1: 100, Y1: 50, Text: "姓名"}, @@ -735,19 +546,16 @@ func TestStripCaptionFromCells_PreservesData(t *testing.T) { StripCaptionFromCells(cells) for i := range cells { if cells[i].Text != orig[i] { - t.Errorf("cell[%d] changed: %q -> %q", i, orig[i], cells[i].Text) + t.Errorf("cell[%d] changed: %q → %q", i, orig[i], cells[i].Text) } } } -// TestStripCaptionFromCells_Empty is a no-op on empty cells. func TestStripCaptionFromCells_Empty(t *testing.T) { cells := []pdf.TSRCell{} StripCaptionFromCells(cells) // must not panic } -// TestConstructTable_StripsCaptionFromCells verifies that constructTable -// strips caption text from cells before building HTML. func TestConstructTable_StripsCaptionFromCells(t *testing.T) { // Cell[0] has caption text "表1：标题"; cell[1] has real data. cells := []pdf.TSRCell{ @@ -766,115 +574,6 @@ func TestConstructTable_StripsCaptionFromCells(t *testing.T) { t.Logf("HTML: %s", html) } -// TestCalSpans_NonSpanningCellsNotPolluted verifies that a regular cell -// at position [0,0] is NOT detected as spanning when a spanning cell at -// [0,1] extends to the left, polluting column boundary calculations. -// Bug: calSpans computed column boundaries from ALL cells including -// spanning cells. "部门开支汇总" at [0,1] with X0=0 extends colLeft[1] -// to 0 instead of 101, shifting the center and causing "Q1" at [0,0] -// to be incorrectly detected as spanning 2 columns. -func TestCalSpans_NonSpanningCellsNotPolluted(t *testing.T) { - // Simulate the SpannedTable test grid: row 0 has Q1(regular), 部门开支汇总(span), Q2(regular) - rows := [][]pdf.TSRCell{ - { - {X0: 0, Y0: 0, X1: 100, Y1: 30, Text: "Q1", Label: "table row"}, - {X0: 0, Y0: 0, X1: 200, Y1: 30, Text: "部门开支汇总", Label: "table spanning cell"}, - {X0: 101, Y0: 0, X1: 200, Y1: 30, Text: "Q2", Label: "table row"}, - }, - { - {X0: 0, Y0: 35, X1: 100, Y1: 65, Text: "100", Label: "table row"}, - {X0: 101, Y0: 35, X1: 200, Y1: 65, Text: "200", Label: "table row"}, - }, - } - - spans, covered := CalSpans(rows) - - // Q1 at [0,0] has X0=0, X1=100 which should only cover its own column. - // It should NOT get a colspan. - if s, ok := spans[[2]int{0, 0}]; ok { - t.Errorf("Q1 at [0,0] should NOT have colspan, got %v. "+ - "Spanning cell at [0,1] polluted column boundaries", s) - } - - // 部门开支汇总 at [0,1] has X0=0, X1=200 which DOES span columns 0 and 1. - if s, ok := spans[[2]int{0, 1}]; !ok { - t.Error("部门开支汇总 at [0,1] should have colspan=2 (covers X=0-200)") - } else if s[0] != 2 { - t.Errorf("部门开支汇总 colspan = %d, want 2", s[0]) - } - - // Q2 at [0,2] should be covered by the spanning cell (col 2 is within X=0-200). - if !covered[[2]int{0, 2}] { - t.Error("Q2 at [0,2] should be covered by spanning cell at [0,1]") - } - - t.Logf("spans: %v, covered: %v", spans, covered) -} - -// ── coordinate space conversion helpers ───────────────────────────────── -func TestRowsToHTML(t *testing.T) { - // rowsToHTML takes [][]pdf.TSRCell instead of [][]string (tableToHTML removed). - toCells := func(rows [][]string) [][]pdf.TSRCell { - out := make([][]pdf.TSRCell, len(rows)) - for ri, row := range rows { - out[ri] = make([]pdf.TSRCell, len(row)) - for ci, s := range row { - out[ri][ci] = pdf.TSRCell{Text: s} - } - } - return out - } - - t.Run("simple 2x2 table", func(t *testing.T) { - rows := toCells([][]string{ - {"姓名", "年龄"}, - {"张三", "25"}, - }) - html := RowsToHTML(rows, "", nil, nil, nil) - expected := "
姓名年龄
张三25
" - if html != expected { - t.Errorf("got %q\nwant %q", html, expected) - } - }) - - t.Run("empty table", func(t *testing.T) { - html := RowsToHTML(nil, "", nil, nil, nil) - if html != "
" { - t.Errorf("expected '
', got %q", html) - } - }) - - t.Run("single cell", func(t *testing.T) { - rows := toCells([][]string{{"X"}}) - html := RowsToHTML(rows, "", nil, nil, nil) - expected := "
X
" - if html != expected { - t.Errorf("got %q\nwant %q", html, expected) - } - }) - - t.Run("matches Python format for 公司差旅费", func(t *testing.T) { - rows := toCells([][]string{ - {"标职务", "飞机", "火车", "轮船", "其他交通工具（不含的士）"}, - {"公司级领导人员", "经济舱位", "火车软席", "二等舱位", "按实报销"}, - {"其他工作人员", "经济舱位", "火车硬席", "三等舱位", "按实报销"}, - }) - html := RowsToHTML(rows, "", nil, nil, nil) - if !strings.HasPrefix(html, "") || !strings.HasSuffix(html, "
") { - t.Errorf("not valid HTML: %s", html) - } - if !strings.Contains(html, "标职务") { - t.Errorf("missing cell '标职务': %s", html) - } - if strings.Count(html, "") != 3 { - t.Errorf("expected 3 rows, got %d", strings.Count(html, "")) - } - }) -} - -// TestExtractTableAndReplace verifies that extractTableAndReplace pops -// table boxes and replaces them with consolidated HTML, matching Python. - func TestExtractTableAndReplace(t *testing.T) { // Build boxes with table labels and a pdf.TableItem with cells. boxes := []pdf.TextBox{ @@ -906,15 +605,13 @@ func TestExtractTableAndReplace(t *testing.T) { } func TestBoxMatchesCell_FalsePositive(t *testing.T) { - // Cell: narrow table cell (40×20 px) + // Cell: narrow table cell (40x20 px) cell := pdf.TSRCell{X0: 0, Y0: 0, X1: 40, Y1: 20} - // Box A: entirely inside the cell → should match. + // Box A: entirely inside the cell → should match boxA := pdf.TextBox{X0: 5, X1: 35, Top: 2, Bottom: 18, Text: "标职务"} - // Box B: a wide body-text box that only slightly overlaps the cell. - // It covers x=30..200 but the cell is only x=0..40. - // Overlap: x=30..40 (10px), box width=170 → ratio=10/170=0.059 < 0.3. + // Box B: a wide body-text box that only slightly overlaps the cell boxB := pdf.TextBox{X0: 30, X1: 200, Top: 5, Bottom: 15, Text: "第二条出差人员应按规定等级乘坐交通工具..."} if !BoxMatchesCell(cell, boxA, true) { @@ -931,10 +628,6 @@ func TestBoxMatchesCell_FalsePositive(t *testing.T) { } } -// TestFillCellTextFromBoxes_PageGlobal verifies that fillCellTextFromBoxes -// correctly matches text boxes to cells when both use page-global 72 DPI -// coordinates, matching Python's construct_table approach. - func TestFillCellTextFromBoxes_PageGlobal(t *testing.T) { t.Run("exact alignment matches", func(t *testing.T) { cells := []pdf.TSRCell{ @@ -949,13 +642,13 @@ func TestFillCellTextFromBoxes_PageGlobal(t *testing.T) { } FillCellTextFromBoxes(cells, boxes) if cells[0].Text != "标职务" { - t.Errorf("cell[0] = %q, want '标职务'", cells[0].Text) + t.Errorf("cell[0] = %q, want %q", cells[0].Text, "标职务") } if cells[1].Text != "飞机" { - t.Errorf("cell[1] = %q, want '飞机'", cells[1].Text) + t.Errorf("cell[1] = %q, want %q", cells[1].Text, "飞机") } if cells[2].Text != "火车" { - t.Errorf("cell[2] = %q, want '火车'", cells[2].Text) + t.Errorf("cell[2] = %q, want %q", cells[2].Text, "火车") } }) @@ -967,7 +660,7 @@ func TestFillCellTextFromBoxes_PageGlobal(t *testing.T) { } FillCellTextFromBoxes(cells, boxes) if cells[0].Text != "标职务" { - t.Errorf("cell text = %q, want '标职务' (body text should not leak in)", cells[0].Text) + t.Errorf("cell text = %q, want %q (body text should not leak in)", cells[0].Text, "标职务") } }) @@ -984,109 +677,6 @@ func TestFillCellTextFromBoxes_PageGlobal(t *testing.T) { }) } -// spans and generates "@@5-6\t..." tags. - -func TestCrossPageTableMerge(t *testing.T) { - // Page 0 table: 2 cells, positioned at page 0. - pg0 := pdf.TableItem{ - Positions: []pdf.Position{ - {PageNumbers: []int{0}, Left: 50, Right: 500, Top: 100, Bottom: 800}, - }, - Scale: 1.0, - Cells: []pdf.TSRCell{ - {X0: 0, Y0: 0, X1: 100, Y1: 50, Text: "pg0_r0c0"}, - {X0: 100, Y0: 0, X1: 200, Y1: 50, Text: "pg0_r0c1"}, - }, - } - // Page 1 table: 2 cells, same X range, positioned at page 1. - pg1 := pdf.TableItem{ - Positions: []pdf.Position{ - {PageNumbers: []int{1}, Left: 50, Right: 500, Top: 100, Bottom: 300}, - }, - Scale: 1.0, - Cells: []pdf.TSRCell{ - {X0: 0, Y0: 0, X1: 100, Y1: 50, Text: "pg1_r0c0"}, - {X0: 100, Y0: 0, X1: 200, Y1: 50, Text: "pg1_r0c1"}, - }, - } - tables := []pdf.TableItem{pg0, pg1} - - // mergeTablesAcrossPages merges tables on consecutive pages with X overlap. - merged := MergeTablesAcrossPages(tables, nil) - if len(merged) != 1 { - t.Fatalf("expected 1 merged table, got %d", len(merged)) - } - if len(merged[0].Cells) != 4 { - t.Errorf("expected 4 merged cells, got %d", len(merged[0].Cells)) - } - if len(merged[0].Positions) != 2 { - t.Errorf("expected 2 merged positions, got %d", len(merged[0].Positions)) - } - t.Logf("Merged %d cells across %d pages", len(merged[0].Cells), len(merged[0].Positions)) -} - -// TestMergeTablesAcrossPages_NoOverlap verifies that non-adjacent or -// non-overlapping tables are NOT merged. - -func TestMergeTablesAcrossPages_NoOverlap(t *testing.T) { - // Tables with no X overlap should NOT be merged. - tables := []pdf.TableItem{ - { - Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 50, Right: 100, Top: 100, Bottom: 500}}, - Scale: 1.0, - Cells: []pdf.TSRCell{{Text: "left"}}, - }, - { - Positions: []pdf.Position{{PageNumbers: []int{1}, Left: 500, Right: 600, Top: 100, Bottom: 500}}, - Scale: 1.0, - Cells: []pdf.TSRCell{{Text: "right"}}, - }, - } - merged := MergeTablesAcrossPages(tables, nil) - if len(merged) != 2 { - t.Fatalf("non-overlapping tables: expected 2 tables, got %d", len(merged)) - } -} - -// TestMergeTablesAcrossPages_NonConsecutive verifies that tables on -// non-consecutive pages are NOT merged. - -func TestMergeTablesAcrossPages_NonConsecutive(t *testing.T) { - tables := []pdf.TableItem{ - { - Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 50, Right: 500, Top: 100, Bottom: 500}}, - Scale: 1.0, - Cells: []pdf.TSRCell{{Text: "page0"}}, - }, - { - Positions: []pdf.Position{{PageNumbers: []int{3}, Left: 50, Right: 500, Top: 100, Bottom: 500}}, - Scale: 1.0, - Cells: []pdf.TSRCell{{Text: "page3"}}, - }, - } - merged := MergeTablesAcrossPages(tables, nil) - if len(merged) != 2 { - t.Fatalf("non-consecutive pages: expected 2 tables, got %d", len(merged)) - } -} - -// TestMergeTablesAcrossPages_SingleTable verifies that a single table -// passes through unchanged. - -func TestMergeTablesAcrossPages_SingleTable(t *testing.T) { - tables := []pdf.TableItem{ - { - Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 50, Right: 500, Top: 100, Bottom: 500}}, - Scale: 1.0, - Cells: []pdf.TSRCell{{Text: "only"}}, - }, - } - merged := MergeTablesAcrossPages(tables, nil) - if len(merged) != 1 { - t.Fatalf("single table: expected 1 table, got %d", len(merged)) - } -} - func TestMergeCaptions_NeedsCaptionLayoutType(t *testing.T) { // Simulate what happens when DLA doesn't produce a "table caption" region: // a "text" section adjacent to a table is NOT treated as caption. @@ -1107,23 +697,200 @@ func TestMergeCaptions_NeedsCaptionLayoutType(t *testing.T) { } } -// TestGroupBoxesByRC_ColspanMissing exposes that groupBoxesByRC doesn't -// compute colspan/rowspan from SP annotations (__cal_spans in Python). +func TestCleanupOrphanColumns(t *testing.T) { + // Test 1: Less than 4 rows - no cleanup + t.Run("less than 4 rows", func(t *testing.T) { + rows := [][]pdf.TSRCell{ + {{Text: "a"}}, + {{Text: "b"}}, + {{Text: "c"}}, + } + result := CleanupOrphanColumns(rows) + if len(result) != 3 { + t.Errorf("expected 3 rows, got %d", len(result)) + } + }) -func TestGroupBoxesByRC_ColspanMissing(t *testing.T) { - // Box with SP annotation spanning 2 columns (HLeft→HRight covers cols 0-1). - boxes := []pdf.TextBox{ - {X0: 10, X1: 90, Top: 0, Bottom: 30, Text: "Name", R: 0, C: 0, H: 1, - HLeft: 10, HRight: 200}, - {X0: 110, X1: 200, Top: 0, Bottom: 30, Text: "", R: 0, C: 1, SP: 1}, - {X0: 10, X1: 90, Top: 35, Bottom: 65, Text: "A", R: 1, C: 0}, - {X0: 110, X1: 200, Top: 35, Bottom: 65, Text: "B", R: 1, C: 1}, - } - rows := GroupBoxesByRC(boxes) - // The result should have colspan=2 for cell [0,0] and skip [0,1]. - // Currently groupBoxesByRC produces a flat grid without span info. - if len(rows) >= 1 && len(rows[0]) >= 2 && rows[0][1].Text == "" { - t.Log("KNOWN LIMITATION: colspan not computed — cell [0,1] is empty instead of merged") - } - _ = rows + // Test 2: 4 rows, no orphan columns + t.Run("4 rows no orphans", func(t *testing.T) { + rows := [][]pdf.TSRCell{ + {{Text: "a"}, {Text: "b"}}, + {{Text: "c"}, {Text: "d"}}, + {{Text: "e"}, {Text: "f"}}, + {{Text: "g"}, {Text: "h"}}, + } + result := CleanupOrphanColumns(rows) + if len(result[0]) != 2 { + t.Errorf("expected 2 columns, got %d", len(result[0])) + } + }) + + // Test 3: 4 rows, one orphan column in the middle + t.Run("4 rows orphan column in middle kept", func(t *testing.T) { + rows := [][]pdf.TSRCell{ + {{Text: "a", X0: 0, X1: 10}, {Text: ""}, {Text: "b", X0: 30, X1: 40}}, + {{Text: "c", X0: 0, X1: 10}, {Text: ""}, {Text: "d", X0: 30, X1: 40}}, + {{Text: "e", X0: 0, X1: 10}, {Text: "orphan", X0: 15, X1: 25}, {Text: "f", X0: 30, X1: 40}}, + {{Text: "g", X0: 0, X1: 10}, {Text: ""}, {Text: "h", X0: 30, X1: 40}}, + } + result := CleanupOrphanColumns(rows) + if len(result[0]) != 3 { + t.Errorf("expected 3 columns (kept because both sides have text), got %d", len(result[0])) + } + }) +} + +func TestCountNonEmptyCells(t *testing.T) { + rows := [][]pdf.TSRCell{ + {{Text: "a"}, {Text: ""}}, + {{Text: ""}, {Text: ""}}, + {{Text: "b"}, {Text: "c"}}, + {{Text: ""}, {Text: ""}}, + } + + count, rowIdx := countNonEmptyCells(rows, 0) + if count != 2 { + t.Errorf("expected 2 non-empty cells in column 0, got %d", count) + } + if rowIdx != 2 { + t.Errorf("expected last non-empty cell at row 2, got %d", rowIdx) + } + + count, rowIdx = countNonEmptyCells(rows, 1) + if count != 1 { + t.Errorf("expected 1 non-empty cell in column 1, got %d", count) + } + if rowIdx != 2 { + t.Errorf("expected last non-empty cell at row 2, got %d", rowIdx) + } + + count, rowIdx = countNonEmptyCells(rows, 999) + if count != 0 { + t.Errorf("expected 0 non-empty cells for invalid column, got %d", count) + } +} + +func TestCheckAdjacentColumns(t *testing.T) { + rows := [][]pdf.TSRCell{ + {{Text: "left"}, {Text: "orphan"}, {Text: "right"}}, + } + + hasLeft, hasRight := checkAdjacentColumns(rows, 1, 0) + if !hasLeft { + t.Error("expected left column to have text") + } + if !hasRight { + t.Error("expected right column to have text") + } + + rows2 := [][]pdf.TSRCell{ + {{Text: ""}, {Text: "orphan"}, {Text: ""}}, + } + hasLeft, hasRight = checkAdjacentColumns(rows2, 1, 0) + if hasLeft { + t.Error("expected left column to be empty") + } + if hasRight { + t.Error("expected right column to be empty") + } + + // Test edge cases + rows3 := [][]pdf.TSRCell{ + {{Text: "only column"}}, + } + hasLeft, hasRight = checkAdjacentColumns(rows3, 0, 0) + if !hasLeft { // j == 0 should count hasLeft as true + t.Error("expected hasLeft to be true when j == 0") + } + if !hasRight { // j+1 >= len should count hasRight as true + t.Error("expected hasRight to be true when j+1 >= len") + } +} + +func TestCalculateMergeDistance(t *testing.T) { + rows := [][]pdf.TSRCell{ + {{Text: "left", X0: 0, X1: 10}, {Text: "orphan", X0: 15, X1: 25}, {Text: "right", X0: 30, X1: 40}}, + } + + leftDist, rightDist := calculateMergeDistance(rows, 1, 0, 3, false, false) + if leftDist != 5 { // 15 - 10 = 5 + t.Errorf("expected left distance 5, got %v", leftDist) + } + if rightDist != 5 { // 30 - 25 = 5 + t.Errorf("expected right distance 5, got %v", rightDist) + } +} + +func TestMergeColumn(t *testing.T) { + tests := []struct { + name string + mergeDir string // "left" or "right" + srcCol int + wantCol0 string + wantCol1 string + }{ + { + name: "merge left", + mergeDir: "left", + srcCol: 1, + wantCol0: "a b", + wantCol1: "b", + }, + { + name: "merge right", + mergeDir: "right", + srcCol: 0, + wantCol0: "a", + wantCol1: "a b", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + rows := [][]pdf.TSRCell{ + {{Text: "a"}, {Text: "b"}}, + {{Text: ""}, {Text: "d"}}, + {{Text: "e"}, {Text: ""}}, + } + + if tt.mergeDir == "left" { + mergeColumnIntoLeft(rows, tt.srcCol) + if rows[0][0].Text != tt.wantCol0 { + t.Errorf("expected '%s', got '%s'", tt.wantCol0, rows[0][0].Text) + } + if rows[1][0].Text != "d" { + t.Errorf("expected 'd', got '%s'", rows[1][0].Text) + } + if rows[2][0].Text != "e" { + t.Errorf("expected 'e', got '%s'", rows[2][0].Text) + } + } else { + mergeColumnIntoRight(rows, tt.srcCol) + if rows[0][1].Text != tt.wantCol1 { + t.Errorf("expected '%s' in right column, got '%s'", tt.wantCol1, rows[0][1].Text) + } + if rows[1][1].Text != "d" { + t.Errorf("expected 'd' in right column, got '%s'", rows[1][1].Text) + } + if rows[2][1].Text != "e" { + t.Errorf("expected 'e' in right column, got '%s'", rows[2][1].Text) + } + } + }) + } +} + +func TestRemoveColumn(t *testing.T) { + rows := [][]pdf.TSRCell{ + {{Text: "a"}, {Text: "b"}, {Text: "c"}}, + {{Text: "d"}, {Text: "e"}, {Text: "f"}}, + } + + result := removeColumn(rows, 1) + if len(result[0]) != 2 { + t.Errorf("expected 2 columns after removal, got %d", len(result[0])) + } + if result[0][0].Text != "a" || result[0][1].Text != "c" { + t.Errorf("unexpected column content after removal") + } } diff --git a/internal/deepdoc/parser/pdf/table/table_html.go b/internal/deepdoc/parser/pdf/table/table_html.go new file mode 100644 index 0000000000..8b3b320c3b --- /dev/null +++ b/internal/deepdoc/parser/pdf/table/table_html.go @@ -0,0 +1,112 @@ +package table + +import ( + "fmt" + "html" + "strings" + + pdf "ragflow/internal/deepdoc/parser/pdf/type" +) + +func RowsToHTML(rows [][]pdf.TSRCell, caption string, headerRows map[int]bool, spanInfo map[[2]int][2]int, covered map[[2]int]bool) string { + var b strings.Builder + b.WriteString("") + if caption != "" { + b.WriteString("") + } + for ri, row := range rows { + b.WriteString("") + for ci, cell := range row { + if covered[[2]int{ri, ci}] { continue } + tag := "td" + if headerRows[ri] { tag = "th" } + b.WriteString("<") + b.WriteString(tag) + sp := "" + if s, ok := spanInfo[[2]int{ri, ci}]; ok { + if s[0] > 1 { + sp = fmt.Sprintf("colspan=%d", s[0]) + } + if s[1] > 1 { + if sp != "" { sp += " " } + sp += fmt.Sprintf("rowspan=%d", s[1]) + } + } + if sp != "" { + b.WriteString(" ") + b.WriteString(sp) + } + b.WriteString(" >") + b.WriteString(html.EscapeString(cell.Text)) + b.WriteString("") + } + b.WriteString("") + } + b.WriteString("
") + b.WriteString(html.EscapeString(caption)) + b.WriteString("
") + return b.String() +} + +// SimpleRowsToHTML converts plain string-based table data to an HTML table. +// The first row is treated as a header (). Used by DOCX, XLSX, PPTX, +// and HTML parsers that produce [][]string directly. +func SimpleRowsToHTML(rows [][]string) string { + if len(rows) == 0 { + return "
" + } + nCols := 0 + for _, row := range rows { + if len(row) > nCols { nCols = len(row) } + } + var b strings.Builder + b.WriteString("") + for ri, row := range rows { + b.WriteString("") + tag := "td" + if ri == 0 { tag = "th" } + for ci := 0; ci < nCols; ci++ { + text := "" + if ci < len(row) { text = row[ci] } + b.WriteString("<") + b.WriteString(tag) + b.WriteString(" >") + b.WriteString(html.EscapeString(text)) + b.WriteString("") + } + b.WriteString("") + } + b.WriteString("
") + return b.String() +} + +func RowsToStrings(rows [][]pdf.TSRCell) [][]string { + out := make([][]string, len(rows)) + for ri, row := range rows { + out[ri] = make([]string, len(row)) + for ci, c := range row { + out[ri][ci] = c.Text + } + } + return out +} + +func HasText(rows [][]pdf.TSRCell) bool { + for _, row := range rows { + for _, c := range row { + if strings.TrimSpace(c.Text) != "" { return true } + } + } + return false +} + +func HasAnyText(cells []pdf.TSRCell) bool { + for _, c := range cells { + if strings.TrimSpace(c.Text) != "" { return true } + } + return false +} diff --git a/internal/deepdoc/parser/pdf/table/table_html_test.go b/internal/deepdoc/parser/pdf/table/table_html_test.go new file mode 100644 index 0000000000..4c96706848 --- /dev/null +++ b/internal/deepdoc/parser/pdf/table/table_html_test.go @@ -0,0 +1,108 @@ + +package table + +import ( + "strings" + "testing" + + pdf "ragflow/internal/deepdoc/parser/pdf/type" +) + +func TestRowsToHTML(t *testing.T) { + // rowsToHTML takes [][]pdf.TSRCell instead of [][]string (tableToHTML removed). + toCells := func(rows [][]string) [][]pdf.TSRCell { + out := make([][]pdf.TSRCell, len(rows)) + for ri, row := range rows { + out[ri] = make([]pdf.TSRCell, len(row)) + for ci, s := range row { + out[ri][ci] = pdf.TSRCell{Text: s} + } + } + return out + } + + t.Run("simple 2x2 table", func(t *testing.T) { + rows := toCells([][]string{ + {"姓名", "年龄"}, + {"张三", "25"}, + }) + html := RowsToHTML(rows, "", nil, nil, nil) + expected := "
姓名年龄
张三25
" + if html != expected { + t.Errorf("got %q\nwant %q", html, expected) + } + }) + + t.Run("empty table", func(t *testing.T) { + html := RowsToHTML(nil, "", nil, nil, nil) + if html != "
" { + t.Errorf("expected '
', got %q", html) + } + }) + + t.Run("single cell", func(t *testing.T) { + rows := toCells([][]string{{"X"}}) + html := RowsToHTML(rows, "", nil, nil, nil) + expected := "
X
" + if html != expected { + t.Errorf("got %q\nwant %q", html, expected) + } + }) + + t.Run("matches Python format for 公司差旅费", func(t *testing.T) { + rows := toCells([][]string{ + {"标职务", "飞机", "火车", "轮船", "其他交通工具（不含的士）"}, + {"公司级领导人员", "经济舱位", "火车软席", "二等舱位", "按实报销"}, + {"其他工作人员", "经济舱位", "火车硬席", "三等舱位", "按实报销"}, + }) + html := RowsToHTML(rows, "", nil, nil, nil) + if !strings.HasPrefix(html, "") || !strings.HasSuffix(html, "
") { + t.Errorf("not valid HTML: %s", html) + } + if !strings.Contains(html, "标职务") { + t.Errorf("missing cell '标职务': %s", html) + } + if strings.Count(html, "") != 3 { + t.Errorf("expected 3 rows, got %d", strings.Count(html, "")) + } + }) +} + +func TestRowsToHTML_HeaderRows(t *testing.T) { + cells := []pdf.TSRCell{ + {X0: 0, Y0: 0, X1: 100, Y1: 30, Text: "Name", Label: "table column header"}, + {X0: 101, Y0: 0, X1: 200, Y1: 30, Text: "Age", Label: "table column header"}, + {X0: 0, Y0: 35, X1: 100, Y1: 65, Text: "John", Label: "table row"}, + {X0: 101, Y0: 35, X1: 200, Y1: 65, Text: "30", Label: "table row"}, + } + // constructTable should produce for header row. + item := &pdf.TableItem{} + html := ConstructTable(cells, nil, "", item) + // Header row should use , data row . + if !strings.Contains(html, "") { + t.Errorf("expected for header row. HTML: %s", html) + } + if strings.Count(html, " cells, got %d. HTML: %s", strings.Count(html, " cells (data row), got %d", strings.Count(html, " 0 { + pg = p.PageNumbers[0] + } + items = append(items, indexed{i, pg, p.Top}) + } + sort.Slice(items, func(i, j int) bool { + if items[i].pg != items[j].pg { + return items[i].pg < items[j].pg + } + return items[i].top < items[j].top + }) + + merged := make([]bool, len(tables)) + var result []pdf.TableItem + + for _, it := range items { + if merged[it.idx] { + continue + } + anchor := tables[it.idx] + merged[it.idx] = true + + // Python nomerge_lout_no: tables whose box is followed by a + // caption/title/reference should not be merged cross-page. + if anchor.NoMerge { + result = append(result, anchor) + continue + } + + anchorPg := it.pg + anchorBtm := anchor.Positions[0].Bottom + + // Look for consecutive-page continuations. + for _, jt := range items { + if merged[jt.idx] || jt.pg <= anchorPg { + continue + } + // Python nomerge_lout_no: skip continuation candidates + // tagged as no-merge. + if tables[jt.idx].NoMerge { + continue + } + if jt.pg-anchorPg > 1 { + break // pages must be consecutive + } + if len(tables[jt.idx].Positions) == 0 { + continue + } + bp := tables[jt.idx].Positions[0] + bpg := 0 + if len(bp.PageNumbers) > 0 { + bpg = bp.PageNumbers[0] + } + if bpg != anchorPg+1 { + continue + } + // Check X overlap. + ap := anchor.Positions[0] + if ap.Right < bp.Left || bp.Right < ap.Left { + continue + } + // Check Y proximity: page 1 table top should be close below + // page 0 table bottom. Python: y_dis <= mh * 23. + mh := 10.0 + if medianHeights != nil { + if h, ok := medianHeights[anchorPg]; ok && h > 0 { + mh = h + } + } + yDis := (bp.Top + bp.Bottom - anchorBtm - ap.Bottom) / 2 + if yDis > mh*23 { + continue + } + // Merge: combine cells and positions. + anchor.Cells = append(anchor.Cells, tables[jt.idx].Cells...) + anchor.Positions = append(anchor.Positions, tables[jt.idx].Positions...) + if tables[jt.idx].Caption != "" { + if anchor.Caption != "" { + anchor.Caption += " " + } + anchor.Caption += tables[jt.idx].Caption + } + merged[jt.idx] = true + anchorPg = bpg + anchorBtm = bp.Bottom + ap = anchor.Positions[len(anchor.Positions)-1] + } + result = append(result, anchor) + } + // Append unprocessed tables (those with empty Positions) so they + // are not silently dropped from the output. + for i := range tables { + if !merged[i] { + result = append(result, tables[i]) + } + } + return result +} + diff --git a/internal/deepdoc/parser/pdf/table/table_merge_test.go b/internal/deepdoc/parser/pdf/table/table_merge_test.go new file mode 100644 index 0000000000..28d8e214e9 --- /dev/null +++ b/internal/deepdoc/parser/pdf/table/table_merge_test.go @@ -0,0 +1,179 @@ +package table + +import ( + "testing" + + pdf "ragflow/internal/deepdoc/parser/pdf/type" +) + +func TestCrossPageTableMerge(t *testing.T) { + // Page 0 table: 2 cells, positioned at page 0. + pg0 := pdf.TableItem{ + Positions: []pdf.Position{ + {PageNumbers: []int{0}, Left: 50, Right: 500, Top: 100, Bottom: 800}, + }, + Scale: 1.0, + Cells: []pdf.TSRCell{ + {X0: 0, Y0: 0, X1: 100, Y1: 50, Text: "pg0_r0c0"}, + {X0: 100, Y0: 0, X1: 200, Y1: 50, Text: "pg0_r0c1"}, + }, + } + // Page 1 table: 2 cells, same X range, positioned at page 1. + pg1 := pdf.TableItem{ + Positions: []pdf.Position{ + {PageNumbers: []int{1}, Left: 50, Right: 500, Top: 100, Bottom: 300}, + }, + Scale: 1.0, + Cells: []pdf.TSRCell{ + {X0: 0, Y0: 0, X1: 100, Y1: 50, Text: "pg1_r0c0"}, + {X0: 100, Y0: 0, X1: 200, Y1: 50, Text: "pg1_r0c1"}, + }, + } + tables := []pdf.TableItem{pg0, pg1} + + // mergeTablesAcrossPages merges tables on consecutive pages with X overlap. + merged := MergeTablesAcrossPages(tables, nil) + if len(merged) != 1 { + t.Fatalf("expected 1 merged table, got %d", len(merged)) + } + if len(merged[0].Cells) != 4 { + t.Errorf("expected 4 merged cells, got %d", len(merged[0].Cells)) + } + if len(merged[0].Positions) != 2 { + t.Errorf("expected 2 merged positions, got %d", len(merged[0].Positions)) + } + t.Logf("Merged %d cells across %d pages", len(merged[0].Cells), len(merged[0].Positions)) +} + +// TestMergeTablesAcrossPages_NoOverlap verifies that non-adjacent or +// non-overlapping tables are NOT merged. +func TestMergeTablesAcrossPages_NoOverlap(t *testing.T) { + // Tables with no X overlap should NOT be merged. + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 50, Right: 100, Top: 100, Bottom: 500}}, + Scale: 1.0, + Cells: []pdf.TSRCell{{Text: "left"}}, + }, + { + Positions: []pdf.Position{{PageNumbers: []int{1}, Left: 500, Right: 600, Top: 100, Bottom: 500}}, + Scale: 1.0, + Cells: []pdf.TSRCell{{Text: "right"}}, + }, + } + merged := MergeTablesAcrossPages(tables, nil) + if len(merged) != 2 { + t.Fatalf("non-overlapping tables: expected 2 tables, got %d", len(merged)) + } +} + +// TestMergeTablesAcrossPages_NonConsecutive verifies that tables on +// non-consecutive pages are NOT merged. +func TestMergeTablesAcrossPages_NonConsecutive(t *testing.T) { + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 50, Right: 500, Top: 100, Bottom: 500}}, + Scale: 1.0, + Cells: []pdf.TSRCell{{Text: "page0"}}, + }, + { + Positions: []pdf.Position{{PageNumbers: []int{3}, Left: 50, Right: 500, Top: 100, Bottom: 500}}, + Scale: 1.0, + Cells: []pdf.TSRCell{{Text: "page3"}}, + }, + } + merged := MergeTablesAcrossPages(tables, nil) + if len(merged) != 2 { + t.Fatalf("non-consecutive pages: expected 2 tables, got %d", len(merged)) + } +} + +// TestMergeTablesAcrossPages_SingleTable verifies that a single table +// passes through unchanged. +func TestMergeTablesAcrossPages_SingleTable(t *testing.T) { + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 50, Right: 500, Top: 100, Bottom: 500}}, + Scale: 1.0, + Cells: []pdf.TSRCell{{Text: "only"}}, + }, + } + merged := MergeTablesAcrossPages(tables, nil) + if len(merged) != 1 { + t.Fatalf("single table: expected 1 table, got %d", len(merged)) + } +} + +func TestMergeTablesAcrossPages_EmptyPositions(t *testing.T) { + // Tables with empty Positions should be preserved (not dropped). + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{}, + Cells: []pdf.TSRCell{{Text: "posless"}}, + }, + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 50, Right: 500, Top: 100, Bottom: 500}}, + Scale: 1.0, + Cells: []pdf.TSRCell{{Text: "normal"}}, + }, + } + merged := MergeTablesAcrossPages(tables, nil) + if len(merged) != 2 { + t.Fatalf("empty Positions: expected 2 tables (preserved), got %d", len(merged)) + } + // Tables with Positions come first (from items list), positionless tables are appended. + if len(merged[0].Positions) == 0 { + t.Error("expected table with Positions first in result") + } + if len(merged[1].Positions) != 0 { + t.Error("expected positionless table second in result") + } + if merged[1].Cells[0].Text != "posless" { + t.Errorf("positionless table content lost: got %q", merged[1].Cells[0].Text) + } +} + +func TestMergeTablesAcrossPages_LargeYGap(t *testing.T) { + // Tables with large Y gap should NOT be merged. + medianHeights := map[int]float64{0: 10} + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 50, Right: 500, Top: 100, Bottom: 150}}, + Scale: 1.0, + Cells: []pdf.TSRCell{{Text: "page0"}}, + }, + { + Positions: []pdf.Position{{PageNumbers: []int{1}, Left: 50, Right: 500, Top: 5000, Bottom: 5100}}, + Scale: 1.0, + Cells: []pdf.TSRCell{{Text: "page1_far"}}, + }, + } + merged := MergeTablesAcrossPages(tables, medianHeights) + if len(merged) != 2 { + t.Fatalf("large Y gap: expected 2 tables (not merged), got %d", len(merged)) + } +} + +func TestMergeTablesAcrossPages_NoMedianHeights(t *testing.T) { + // Without medianHeights, mh defaults to 10 and threshold is 230. + // yDis = (10 + 120 - 150 - 150) / 2 = -85, which is <= 230, so they merge. + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 50, Right: 500, Top: 100, Bottom: 150}}, + Scale: 1.0, + Cells: []pdf.TSRCell{{Text: "page0"}}, + }, + { + Positions: []pdf.Position{{PageNumbers: []int{1}, Left: 50, Right: 500, Top: 10, Bottom: 120}}, + Scale: 1.0, + Cells: []pdf.TSRCell{{Text: "page1_near"}}, + }, + } + merged := MergeTablesAcrossPages(tables, nil) + if len(merged) != 1 { + t.Fatalf("no medianHeights: expected 1 merged table, got %d", len(merged)) + } + if len(merged[0].Cells) != 2 { + t.Errorf("expected 2 cells after merge, got %d", len(merged[0].Cells)) + } +} diff --git a/internal/deepdoc/parser/pdf/table/table_post.go b/internal/deepdoc/parser/pdf/table/table_post.go index f97c5eb61a..2acaf74031 100644 --- a/internal/deepdoc/parser/pdf/table/table_post.go +++ b/internal/deepdoc/parser/pdf/table/table_post.go @@ -8,6 +8,192 @@ import ( pdf "ragflow/internal/deepdoc/parser/pdf/type" ) +// FilterBoxesByRemoveSet filters boxes by index set +// removeSet: key is index to remove, value=true means remove +func FilterBoxesByRemoveSet(boxes []pdf.TextBox, removeSet map[int]bool) []pdf.TextBox { + if len(removeSet) == 0 { + return boxes + } + if len(boxes) == 0 { + return boxes + } + // Pre-allocate: estimate final size to avoid resizing + // Use max to prevent negative capacity when len(removeSet) > len(boxes) + estimatedCap := len(boxes) - len(removeSet) + if estimatedCap < 0 { + estimatedCap = 0 + } + out := make([]pdf.TextBox, 0, estimatedCap) + for i, b := range boxes { + if !removeSet[i] { + out = append(out, b) + } + } + return out +} + +// createTableBoxFromItem creates HTML-containing TextBox from TableItem +func createTableBoxFromItem(tbl *pdf.TableItem, html string) pdf.TextBox { + pg := 0 + if len(tbl.Positions) > 0 && len(tbl.Positions[0].PageNumbers) > 0 { + pg = tbl.Positions[0].PageNumbers[0] + } + x0, x1, top, bottom := tbl.RegionLeft, tbl.RegionRight, tbl.RegionTop, tbl.RegionBottom + if x0 == 0 && x1 == 0 && top == 0 && bottom == 0 && len(tbl.Positions) > 0 { + p := tbl.Positions[0] + x0, x1, top, bottom = p.Left, p.Right, p.Top, p.Bottom + } + return pdf.TextBox{ + X0: x0, + X1: x1, + Top: top, + Bottom: bottom, + Text: html, + PageNumber: pg, + LayoutType: pdf.LayoutTypeTable, + } +} + +// handleImageOnlyPDFs handles cases with no boxes but tables (Image-only PDF) +func handleImageOnlyPDFs(tables []pdf.TableItem) []pdf.TextBox { + var out []pdf.TextBox + for ti := range tables { + if len(tables[ti].Cells) == 0 { + continue + } + s := tables[ti].Scale + pageGlobalCells := CellSliceToPageSpace(tables[ti].Cells, tables[ti].CropOffX, tables[ti].CropOffY, s) + var tableBoxes []pdf.TextBox + html := ConstructTable(pageGlobalCells, tableBoxes, tables[ti].Caption, &tables[ti]) + if html != "" { + out = append(out, createTableBoxFromItem(&tables[ti], html)) + } + } + return out +} + +// findTableAnchors finds the best insertion position for each table by finding +// the spatially nearest non-table text box. Returns a list of (tableIndex, position) +// pairs sorted by position. +func findTableAnchors(boxes []pdf.TextBox, tables []pdf.TableItem) []struct{ ti, pos int } { + replacedByTable := make(map[int]int) + + for ti := range tables { + if len(tables[ti].Cells) == 0 { + continue + } + tbl := &tables[ti] + tblLeft, tblRight := tbl.RegionLeft, tbl.RegionRight + tblTop, tblBottom := tbl.RegionTop, tbl.RegionBottom + tblPg := 0 + if len(tbl.Positions) > 0 { + p := tbl.Positions[0] + if len(p.PageNumbers) > 0 { + tblPg = p.PageNumbers[0] + } + if tblLeft == 0 && tblRight == 0 && tblTop == 0 && tblBottom == 0 { + tblLeft, tblRight = p.Left, p.Right + tblTop, tblBottom = p.Top, p.Bottom + } + } + bestDist := math.MaxFloat64 + bestIdx := -1 + for i, b := range boxes { + if b.LayoutType == pdf.LayoutTypeTable || b.LayoutType == pdf.LayoutTypeFigure { + continue + } + if b.PageNumber != tblPg { + continue + } + dist := minRectangleDistance( + b.X0, b.X1, b.Top, b.Bottom, + tblLeft, tblRight, tblTop, tblBottom, + ) + if dist < bestDist { + bestDist = dist + bestIdx = i + } + } + if bestIdx >= 0 { + if boxes[bestIdx].Bottom < tblTop { + bestIdx++ + } + replacedByTable[ti] = bestIdx + } + } + + // Build the anchor list and sort by position + anchorList := make([]struct{ ti, pos int }, 0, len(replacedByTable)) + for ti, pos := range replacedByTable { + anchorList = append(anchorList, struct{ ti, pos int }{ti, pos}) + } + sort.Slice(anchorList, func(i, j int) bool { return anchorList[i].pos < anchorList[j].pos }) + return anchorList +} + +// buildTableHTMLs constructs HTML for each table, converting cells to page space first. +// Returns a map from table index to HTML string. +func buildTableHTMLs(boxes []pdf.TextBox, tables []pdf.TableItem) map[int]string { + htmls := make(map[int]string) + for ti := range tables { + if len(tables[ti].Cells) == 0 { + continue + } + // Convert TSR cells from crop-pixel space to page-global 72 DPI + s := tables[ti].Scale + pageGlobalCells := CellSliceToPageSpace(tables[ti].Cells, tables[ti].CropOffX, tables[ti].CropOffY, s) + // Collect only table-labelled boxes + var tableBoxes []pdf.TextBox + for i := range boxes { + if boxes[i].LayoutType != pdf.LayoutTypeTable { + continue + } + for _, tp := range tables[ti].Positions { + if boxOverlapsPosition(boxes[i], tp) { + tableBoxes = append(tableBoxes, boxes[i]) + break + } + } + } + slog.Debug("extractTableAndReplace constructTable", "table", ti, "cells", len(pageGlobalCells), "boxes", len(tableBoxes)) + htmls[ti] = ConstructTable(pageGlobalCells, tableBoxes, tables[ti].Caption, &tables[ti]) + } + return htmls +} + +// insertTableBoxes filters out boxes in removeSet and inserts table HTML boxes at anchor positions. +func insertTableBoxes(boxes []pdf.TextBox, tables []pdf.TableItem, removeSet map[int]bool, + anchors []struct{ ti, pos int }, htmls map[int]string) []pdf.TextBox { + + out := make([]pdf.TextBox, 0, len(boxes)-len(removeSet)+len(anchors)) + anchorIdx := 0 + for i, b := range boxes { + // Insert any HTML boxes whose anchor position is before or at i + for anchorIdx < len(anchors) && anchors[anchorIdx].pos <= i { + ti := anchors[anchorIdx].ti + if html, ok := htmls[ti]; ok && html != "" { + tbl := &tables[ti] + out = append(out, tableRegionBox(tbl, &b, html)) + } + anchorIdx++ + } + if !removeSet[i] { + out = append(out, b) + } + } + // Insert remaining anchors after last box + for anchorIdx < len(anchors) { + ti := anchors[anchorIdx].ti + if html, ok := htmls[ti]; ok && html != "" { + tbl := &tables[ti] + last := &boxes[len(boxes)-1] + out = append(out, tableRegionBox(tbl, last, html)) + } + anchorIdx++ + } + return out +} + // extractTableAndReplace pops table boxes and replaces them with consolidated // HTML boxes (one per table). This matches Python's _extract_table_figure which // pops all boxes inside a table DLA region and inserts a single HTML box. @@ -53,15 +239,22 @@ type replacement struct { boxIdx int } -// buildReplacements scans for data-source-attribution boxes to remove and maps -// each table to overlapping table-layout boxes, producing the replacement list. -func buildReplacements(boxes []pdf.TextBox, tables []pdf.TableItem) (map[int]bool, []replacement) { +// buildRemoveSet scans for data-source-attribution boxes to remove. +// Does NOT depend on table indices — safe to call before MergeTablesAcrossPages. +func buildRemoveSet(boxes []pdf.TextBox) map[int]bool { removeSet := make(map[int]bool) for i := range boxes { if boxes[i].LayoutType == pdf.LayoutTypeTable && isDataSourceBox(boxes[i].Text) { removeSet[i] = true } } + return removeSet +} + +// buildReplacementsAfterMerge maps each table to overlapping table-layout boxes, +// producing the replacement list. Must be called AFTER MergeTablesAcrossPages so +// that tableIdx in each replacement refers to the correct merged-table slot. +func buildReplacementsAfterMerge(boxes []pdf.TextBox, tables []pdf.TableItem, removeSet map[int]bool) []replacement { var reps []replacement for ti := range tables { for i := range boxes { @@ -76,213 +269,115 @@ func buildReplacements(boxes []pdf.TextBox, tables []pdf.TableItem) (map[int]boo } } } + return reps +} + +// buildReplacements scans for data-source-attribution boxes to remove and maps +// each table to overlapping table-layout boxes, producing the replacement list. +// Deprecated: pre-merge variant kept for compatibility; prefer calling +// buildRemoveSet + MergeTablesAcrossPages + buildReplacementsAfterMerge. +func buildReplacements(boxes []pdf.TextBox, tables []pdf.TableItem) (map[int]bool, []replacement) { + removeSet := buildRemoveSet(boxes) + reps := buildReplacementsAfterMerge(boxes, tables, removeSet) return removeSet, reps } func ExtractTableAndReplace(boxes []pdf.TextBox, tables []pdf.TableItem) []pdf.TextBox { + removeSet := buildRemoveSet(boxes) if len(tables) == 0 { - return boxes + return FilterBoxesByRemoveSet(boxes, removeSet) } - // Pre-merge nomerge detection: match Python's nomerge_lout_no. - // Traverse boxes in page order. When a caption/title/reference is - // found, mark the preceding table group as NoMerge, preventing - // cross-page merge when a caption ends a table group. - // Python: if is_caption(c) or layout_type in ["table caption", "title", - // "figure caption", "reference"]: nomerge_lout_no.append(lst_lout_no) - MarkNoMergeTables(boxes, tables) - // Merge same-layoutno tables across consecutive pages (Python _extract_table_figure). + MarkNoMergeTables(boxes, tables) tables = MergeTablesAcrossPages(tables, nil) - // Pre-scan: mark data-source-attribution table boxes for removal. - // Python: if re.match(r"(数据|资料|图表)*来源[:： ]", self.boxes[i]["text"]): - // self.boxes.pop(i); continue — box discarded, no HTML replacement. - removeSet, replacements := buildReplacements(boxes, tables) + // Build replacements AFTER merge so tableIdx refers to the merged slice. + replacements := buildReplacementsAfterMerge(boxes, tables, removeSet) - // Image-only PDFs (0 boxes) may have tables with cells but no - // overlapping LayoutType=="table" boxes — generate HTML directly. if len(replacements) == 0 && len(boxes) == 0 { - var out []pdf.TextBox - for ti := range tables { - if len(tables[ti].Cells) == 0 { - continue - } - s := tables[ti].Scale - pageGlobalCells := CellSliceToPageSpace(tables[ti].Cells, tables[ti].CropOffX, tables[ti].CropOffY, s) - var tableBoxes []pdf.TextBox - html := ConstructTable(pageGlobalCells, tableBoxes, tables[ti].Caption, &tables[ti]) - if html != "" { - out = append(out, pdf.TextBox{ - Text: html, LayoutType: "table", PageNumber: 0, - }) - } - } - return out + return handleImageOnlyPDFs(tables) } if len(replacements) == 0 { - // No HTML replacements, but data-source boxes still need removal. - if len(removeSet) == 0 { - return boxes - } - out := make([]pdf.TextBox, 0, len(boxes)-len(removeSet)) - for i, b := range boxes { - if !removeSet[i] { - out = append(out, b) - } - } - return out + return FilterBoxesByRemoveSet(boxes, removeSet) } - // Distance-based anchor selection (Python's min_rectangle_distance). - // Find the spatially nearest non-table text box for each table and - // use that as the anchor, matching insert_table_figures behavior. - replacedByTable := make(map[int]int) - for ti := range tables { - if len(tables[ti].Cells) == 0 { - continue - } - tbl := &tables[ti] - tblLeft, tblRight := tbl.RegionLeft, tbl.RegionRight - tblTop, tblBottom := tbl.RegionTop, tbl.RegionBottom - tblPg := 0 - if len(tbl.Positions) > 0 { - p := tbl.Positions[0] - if len(p.PageNumbers) > 0 { - tblPg = p.PageNumbers[0] - } - if tblLeft == 0 && tblRight == 0 && tblTop == 0 && tblBottom == 0 { - tblLeft, tblRight = p.Left, p.Right - tblTop, tblBottom = p.Top, p.Bottom - } - } - bestDist := math.MaxFloat64 - bestIdx := -1 - for i, b := range boxes { - if b.LayoutType == pdf.LayoutTypeTable || b.LayoutType == pdf.LayoutTypeFigure { - continue - } - if b.PageNumber != tblPg { - continue - } - dist := minRectangleDistance( - b.X0, b.X1, b.Top, b.Bottom, - tblLeft, tblRight, tblTop, tblBottom, - ) - if dist < bestDist { - bestDist = dist - bestIdx = i - } - } - if bestIdx >= 0 { - if boxes[bestIdx].Bottom < tblTop { - bestIdx++ - } - replacedByTable[ti] = bestIdx - } else { - for _, r := range replacements { - if r.tableIdx == ti { - if _, ok := replacedByTable[ti]; !ok || r.boxIdx < replacedByTable[ti] { - replacedByTable[ti] = r.boxIdx - } - } - } - } + return processTablesWithReplacements(boxes, tables, removeSet, replacements) +} + +// buildAndSortAnchors creates and sorts anchor list +func buildAndSortAnchors(anchors map[int]int) []struct{ ti, pos int } { + result := make([]struct{ ti, pos int }, 0, len(anchors)) + for ti, pos := range anchors { + result = append(result, struct{ ti, pos int }{ti: ti, pos: pos}) } + sort.Slice(result, func(i, j int) bool { return result[i].pos < result[j].pos }) + return result +} + +// processTablesWithReplacements handles normal flow with replacements +func processTablesWithReplacements( + boxes []pdf.TextBox, + tables []pdf.TableItem, + removeSet map[int]bool, + replacements []replacement, +) []pdf.TextBox { for _, r := range replacements { removeSet[r.boxIdx] = true } + anchors := findTableAnchorsWithReplacements(boxes, tables, replacements) + htmls := buildTableHTMLs(boxes, tables) + anchorList := buildAndSortAnchors(anchors) + return insertTableBoxes(boxes, tables, removeSet, anchorList, htmls) +} - // Build HTML for each table using post-merge boxes converted to crop space. - htmlByTable := make(map[int]string) +// findTableAnchorsWithReplacements is like findTableAnchors but falls back to +// replacement positions when no text box anchor is found. +func findTableAnchorsWithReplacements(boxes []pdf.TextBox, tables []pdf.TableItem, + replacements []replacement) map[int]int { + + // First get anchors from findTableAnchors + anchorList := findTableAnchors(boxes, tables) + result := make(map[int]int, len(anchorList)) + for _, a := range anchorList { + result[a.ti] = a.pos + } + + // Fill in any missing tables using replacements for ti := range tables { - if len(tables[ti].Cells) == 0 { + if _, has := result[ti]; has { continue } - // Convert TSR cells from crop-pixel space to page-global 72 DPI, - // matching Python's coordinate space. Text boxes are already in - // page-global 72 DPI (from ocrMergeChars), so no conversion needed. - s := tables[ti].Scale - pageGlobalCells := CellSliceToPageSpace(tables[ti].Cells, tables[ti].CropOffX, tables[ti].CropOffY, s) - // Collect only table-labelled boxes (Python: filters by layout_type). - var tableBoxes []pdf.TextBox - for i := range boxes { - if boxes[i].LayoutType != pdf.LayoutTypeTable { - continue - } - for _, tp := range tables[ti].Positions { - if boxOverlapsPosition(boxes[i], tp) { - tableBoxes = append(tableBoxes, boxes[i]) - break + // Find the earliest replacement for this table + for _, r := range replacements { + if r.tableIdx == ti { + if _, ok := result[ti]; !ok || r.boxIdx < result[ti] { + result[ti] = r.boxIdx } } } - slog.Debug("extractTableAndReplace constructTable", "table", ti, "cells", len(pageGlobalCells), "boxes", len(tableBoxes)) - htmlByTable[ti] = ConstructTable(pageGlobalCells, tableBoxes, tables[ti].Caption, &tables[ti]) } - - // Sort anchors by position for stable insertion. - anchorList := make([]struct{ ti, pos int }, 0, len(replacedByTable)) - for ti, pos := range replacedByTable { - anchorList = append(anchorList, struct{ ti, pos int }{ti, pos}) - } - sort.Slice(anchorList, func(i, j int) bool { return anchorList[i].pos < anchorList[j].pos }) - - out := make([]pdf.TextBox, 0, len(boxes)-len(removeSet)+len(replacedByTable)) - anchorIdx := 0 - for i, b := range boxes { - // Insert any HTML boxes whose anchor position is before or at i. - for anchorIdx < len(anchorList) && anchorList[anchorIdx].pos <= i { - ti := anchorList[anchorIdx].ti - html := htmlByTable[ti] - if html != "" { - tbl := &tables[ti] - out = append(out, tableRegionBox(tbl, &b, html)) - } - anchorIdx++ - } - if !removeSet[i] { - out = append(out, b) - } - } - // Remaining anchors after last box. - for anchorIdx < len(anchorList) { - ti := anchorList[anchorIdx].ti - html := htmlByTable[ti] - if html != "" { - tbl := &tables[ti] - last := &boxes[len(boxes)-1] - out = append(out, tableRegionBox(tbl, last, html)) - } - anchorIdx++ - } - return out + return result } -// consolidateFigures merges figure boxes that share the same LayoutNo -// (i.e., belong to the same DLA figure region) into a single pdf.TextBox. -// Matches Python's _extract_table_figure + insert_table_figures which pops -// individual figure boxes and re-inserts one consolidated figure block -// per DLA region with combined text. -// -// Figure boxes whose text matches the data-source discard pattern -// (r"(数据|资料|图表)*来源[:： ]") are removed entirely — matching Python's -// _extract_table_figure discard behavior (pdf_parser.py:1050-1052). -func ConsolidateFigures(boxes []pdf.TextBox) []pdf.TextBox { - // Pre-scan: mark data-source-attribution figure boxes for removal. - // Python: if re.match(r"(数据|资料|图表)*来源[:： ]", self.boxes[i]["text"]): - // self.boxes.pop(i); continue — box discarded. +// figKey groups figure boxes by page and layout number +type figKey struct { + page int + ln string +} + +// markDataSourceBoxesForRemoval marks data source attribution figure boxes for removal +func markDataSourceBoxesForRemoval(boxes []pdf.TextBox) map[int]bool { removeSet := make(map[int]bool) for i, b := range boxes { if b.LayoutType == pdf.LayoutTypeFigure && isDataSourceBox(b.Text) { removeSet[i] = true } } + return removeSet +} - // Group figure boxes by (page, layoutno). - type figKey struct { - page int - ln string - } +// groupFigureBoxes groups figure boxes by (page, layoutno) +func groupFigureBoxes(boxes []pdf.TextBox, removeSet map[int]bool) map[figKey][]int { groups := make(map[figKey][]int) for i, b := range boxes { if b.LayoutType != pdf.LayoutTypeFigure || removeSet[i] { @@ -291,27 +386,15 @@ func ConsolidateFigures(boxes []pdf.TextBox) []pdf.TextBox { key := figKey{b.PageNumber, b.LayoutNo} groups[key] = append(groups[key], i) } + return groups +} - if len(groups) == 0 { - // Still need to filter out data-source figure boxes. - if len(removeSet) == 0 { - return boxes - } - out := make([]pdf.TextBox, 0, len(boxes)-len(removeSet)) - for i, b := range boxes { - if !removeSet[i] { - out = append(out, b) - } - } - return out - } - - // Collect indices to remove (all group members except the first). +// mergeFigureGroups merges figure boxes within groups +func mergeFigureGroups(boxes []pdf.TextBox, groups map[figKey][]int, removeSet map[int]bool) { for _, indices := range groups { if len(indices) <= 1 { continue } - // Merge into the first box of the group. anchor := indices[0] for _, idx := range indices[1:] { b := boxes[idx] @@ -323,18 +406,26 @@ func ConsolidateFigures(boxes []pdf.TextBox) []pdf.TextBox { removeSet[idx] = true } } +} - if len(removeSet) == 0 { - return boxes +// ConsolidateFigures merges figure boxes that share the same LayoutNo +// (i.e., belong to the same DLA figure region) into a single pdf.TextBox. +// Matches Python's _extract_table_figure + insert_table_figures which pops +// individual figure boxes and re-inserts one consolidated figure block +// per DLA region with combined text. +// +// Figure boxes whose text matches the data-source discard pattern +// (r"(数据|资料|图表)*来源[:： ]") are removed entirely — matching Python's +// _extract_table_figure discard behavior (pdf_parser.py:1050-1052). +func ConsolidateFigures(boxes []pdf.TextBox) []pdf.TextBox { + removeSet := markDataSourceBoxesForRemoval(boxes) + groups := groupFigureBoxes(boxes, removeSet) + + if len(groups) > 0 { + mergeFigureGroups(boxes, groups, removeSet) } - out := make([]pdf.TextBox, 0, len(boxes)-len(removeSet)) - for i, b := range boxes { - if !removeSet[i] { - out = append(out, b) - } - } - return out + return FilterBoxesByRemoveSet(boxes, removeSet) } // boxOverlapsPosition checks if a pdf.TextBox overlaps a pdf.Position with margin. diff --git a/internal/deepdoc/parser/pdf/table/table_post_test.go b/internal/deepdoc/parser/pdf/table/table_post_test.go new file mode 100644 index 0000000000..f2cae473c9 --- /dev/null +++ b/internal/deepdoc/parser/pdf/table/table_post_test.go @@ -0,0 +1,802 @@ +package table + +import ( + "strings" + "testing" + + pdf "ragflow/internal/deepdoc/parser/pdf/type" +) + +// ============================================ +// 第一部分：findTableAnchors 的测试 +// ============================================ + +func TestFindTableAnchors_SingleTable(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "before", LayoutType: pdf.LayoutTypeText, PageNumber: 0, X0: 10, X1: 100, Top: 10, Bottom: 50}, + {Text: "table1", LayoutType: pdf.LayoutTypeTable, PageNumber: 0, X0: 10, X1: 400, Top: 60, Bottom: 200}, + {Text: "after", LayoutType: pdf.LayoutTypeText, PageNumber: 0, X0: 10, X1: 100, Top: 210, Bottom: 250}, + } + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 10, Right: 400, Top: 60, Bottom: 200}}, + RegionLeft: 10, RegionRight: 400, RegionTop: 60, RegionBottom: 200, + Cells: []pdf.TSRCell{{Text: "cell"}}, + }, + } + + anchors := findTableAnchors(boxes, tables) + if len(anchors) != 1 { + t.Errorf("expected 1 anchor, got %d", len(anchors)) + } + if anchors[0].pos != 1 { + t.Errorf("expected anchor at pos 1, got %d", anchors[0].pos) + } +} + +func TestFindTableAnchors_NoBoxes(t *testing.T) { + boxes := []pdf.TextBox{} + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 10, Right: 400, Top: 60, Bottom: 200}}, + Cells: []pdf.TSRCell{{Text: "cell"}}, + }, + } + + anchors := findTableAnchors(boxes, tables) + if len(anchors) != 0 { + t.Errorf("expected 0 anchors, got %d", len(anchors)) + } +} + +func TestFindTableAnchors_MultipleTables(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "text1", LayoutType: pdf.LayoutTypeText, PageNumber: 0, X0: 10, X1: 100, Top: 10, Bottom: 30}, + {Text: "table1", LayoutType: pdf.LayoutTypeTable, PageNumber: 0, X0: 10, X1: 400, Top: 40, Bottom: 100}, + {Text: "text2", LayoutType: pdf.LayoutTypeText, PageNumber: 0, X0: 10, X1: 100, Top: 110, Bottom: 140}, + {Text: "table2", LayoutType: pdf.LayoutTypeTable, PageNumber: 0, X0: 10, X1: 400, Top: 150, Bottom: 210}, + } + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 10, Right: 400, Top: 40, Bottom: 100}}, + RegionLeft: 10, RegionRight: 400, RegionTop: 40, RegionBottom: 100, + Cells: []pdf.TSRCell{{Text: "cell1"}}, + }, + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 10, Right: 400, Top: 150, Bottom: 210}}, + RegionLeft: 10, RegionRight: 400, RegionTop: 150, RegionBottom: 210, + Cells: []pdf.TSRCell{{Text: "cell2"}}, + }, + } + + anchors := findTableAnchors(boxes, tables) + if len(anchors) != 2 { + t.Errorf("expected 2 anchors, got %d", len(anchors)) + } +} + +func TestFindTableAnchors_AnchorAboveTable(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "above", LayoutType: pdf.LayoutTypeText, PageNumber: 0, X0: 10, X1: 100, Top: 10, Bottom: 30}, + {Text: "table", LayoutType: pdf.LayoutTypeTable, PageNumber: 0, X0: 10, X1: 400, Top: 40, Bottom: 100}, + } + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 10, Right: 400, Top: 40, Bottom: 100}}, + RegionLeft: 10, RegionRight: 400, RegionTop: 40, RegionBottom: 100, + Cells: []pdf.TSRCell{{Text: "cell"}}, + }, + } + + anchors := findTableAnchors(boxes, tables) + if len(anchors) != 1 { + t.Errorf("expected 1 anchor, got %d", len(anchors)) + } + if anchors[0].pos != 1 { + t.Errorf("expected anchor at pos 1 (insert after above text), got %d", anchors[0].pos) + } +} + +func TestFindTableAnchors_DifferentPage(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "page0", LayoutType: pdf.LayoutTypeText, PageNumber: 0, X0: 10, X1: 100, Top: 10, Bottom: 30}, + {Text: "page1", LayoutType: pdf.LayoutTypeText, PageNumber: 1, X0: 10, X1: 100, Top: 10, Bottom: 30}, + } + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 10, Right: 400, Top: 40, Bottom: 100}}, + RegionLeft: 10, RegionRight: 400, RegionTop: 40, RegionBottom: 100, + Cells: []pdf.TSRCell{{Text: "cell"}}, + }, + } + + anchors := findTableAnchors(boxes, tables) + if len(anchors) != 1 { + t.Errorf("expected 1 anchor, got %d", len(anchors)) + } + // The text box is above the table, so pos is incremented to 1 + if anchors[0].pos != 1 { + t.Errorf("expected anchor at pos 1, got %d", anchors[0].pos) + } +} + +// ============================================ +// 第二部分：buildTableHTMLs 的测试 +// ============================================ + +func TestBuildTableHTMLs_SingleTable(t *testing.T) { + boxes := []pdf.TextBox{} + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 10, Right: 400, Top: 60, Bottom: 200}}, + Scale: 1.0, + Cells: []pdf.TSRCell{ + {X0: 0, Y0: 0, X1: 100, Y1: 50, Text: "A"}, + {X0: 100, Y0: 0, X1: 200, Y1: 50, Text: "B"}, + }, + }, + } + + htmls := buildTableHTMLs(boxes, tables) + if len(htmls) != 1 { + t.Errorf("expected 1 HTML entry, got %d", len(htmls)) + } + if htmls[0] == "" { + t.Error("expected non-empty HTML") + } +} + +func TestBuildTableHTMLs_NoCells(t *testing.T) { + boxes := []pdf.TextBox{} + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 10, Right: 400, Top: 60, Bottom: 200}}, + Cells: []pdf.TSRCell{}, + }, + } + + htmls := buildTableHTMLs(boxes, tables) + if len(htmls) != 0 { + t.Errorf("expected 0 HTML entries for no cells, got %d", len(htmls)) + } +} + +func TestBuildTableHTMLs_WithTableBoxes(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "cell text", LayoutType: pdf.LayoutTypeTable, PageNumber: 0, X0: 10, X1: 100, Top: 60, Bottom: 100}, + } + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 10, Right: 400, Top: 60, Bottom: 200}}, + Scale: 1.0, + Cells: []pdf.TSRCell{ + {X0: 0, Y0: 0, X1: 100, Y1: 50, Text: "A"}, + }, + }, + } + + htmls := buildTableHTMLs(boxes, tables) + if len(htmls) != 1 { + t.Errorf("expected 1 HTML entry, got %d", len(htmls)) + } +} + +// ============================================ +// 第三部分：insertTableBoxes 的测试 +// ============================================ + +func TestInsertTableBoxes_Basic(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "before", PageNumber: 0}, + {Text: "to replace", LayoutType: pdf.LayoutTypeTable, PageNumber: 0}, + {Text: "after", PageNumber: 0}, + } + removeSet := map[int]bool{1: true} + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 10, Right: 400, Top: 60, Bottom: 200}}, + RegionLeft: 10, RegionRight: 400, RegionTop: 60, RegionBottom: 200, + }, + } + anchors := []struct{ ti, pos int }{{ti: 0, pos: 1}} + htmls := map[int]string{0: "test
"} + + result := insertTableBoxes(boxes, tables, removeSet, anchors, htmls) + if len(result) != 3 { + t.Errorf("expected 3 boxes (before + table + after), got %d", len(result)) + } + if result[1].Text != "test
" { + t.Errorf("expected table HTML at position 1, got %q", result[1].Text) + } +} + +func TestInsertTableBoxes_NoRemove(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "before", PageNumber: 0}, + {Text: "after", PageNumber: 0}, + } + removeSet := map[int]bool{} + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 10, Right: 400, Top: 60, Bottom: 200}}, + RegionLeft: 10, RegionRight: 400, RegionTop: 60, RegionBottom: 200, + }, + } + anchors := []struct{ ti, pos int }{{ti: 0, pos: 1}} + htmls := map[int]string{0: "test
"} + + result := insertTableBoxes(boxes, tables, removeSet, anchors, htmls) + if len(result) != 3 { + t.Errorf("expected 3 boxes, got %d", len(result)) + } +} + +func TestInsertTableBoxes_AtEnd(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "first", PageNumber: 0}, + {Text: "second", PageNumber: 0}, + } + removeSet := map[int]bool{} + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 10, Right: 400, Top: 60, Bottom: 200}}, + RegionLeft: 10, RegionRight: 400, RegionTop: 60, RegionBottom: 200, + }, + } + anchors := []struct{ ti, pos int }{{ti: 0, pos: 2}} + htmls := map[int]string{0: "end
"} + + result := insertTableBoxes(boxes, tables, removeSet, anchors, htmls) + if len(result) != 3 { + t.Errorf("expected 3 boxes, got %d", len(result)) + } + if result[2].Text != "end
" { + t.Errorf("expected table at end, got %q", result[2].Text) + } +} + +func TestInsertTableBoxes_MultipleAnchors(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "1", PageNumber: 0}, + {Text: "2", PageNumber: 0}, + {Text: "3", PageNumber: 0}, + } + removeSet := map[int]bool{} + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 10, Right: 400, Top: 60, Bottom: 200}}, + RegionLeft: 10, RegionRight: 400, RegionTop: 60, RegionBottom: 200, + }, + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 10, Right: 400, Top: 210, Bottom: 350}}, + RegionLeft: 10, RegionRight: 400, RegionTop: 210, RegionBottom: 350, + }, + } + anchors := []struct{ ti, pos int }{{ti: 0, pos: 1}, {ti: 1, pos: 3}} + htmls := map[int]string{0: "A
", 1: "B
"} + + result := insertTableBoxes(boxes, tables, removeSet, anchors, htmls) + if len(result) != 5 { + t.Errorf("expected 5 boxes, got %d", len(result)) + } + if result[1].Text != "A
" { + t.Errorf("expected table A at pos 1") + } + if result[4].Text != "B
" { + t.Errorf("expected table B at pos 4") + } +} + +func TestInsertTableBoxes_EmptyHTML(t *testing.T) { + boxes := []pdf.TextBox{{Text: "text", PageNumber: 0}} + removeSet := map[int]bool{} + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 10, Right: 400, Top: 60, Bottom: 200}}, + RegionLeft: 10, RegionRight: 400, RegionTop: 60, RegionBottom: 200, + }, + } + anchors := []struct{ ti, pos int }{{ti: 0, pos: 1}} + htmls := map[int]string{0: ""} + + result := insertTableBoxes(boxes, tables, removeSet, anchors, htmls) + if len(result) != 1 { + t.Errorf("expected 1 box (no empty HTML inserted), got %d", len(result)) + } +} + +// ============================================ +// 第四部分：集成测试 - 验证重构后功能保持一致 +// ============================================ + +func TestExtractTableAndReplace_Integration(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "intro", LayoutType: pdf.LayoutTypeText, PageNumber: 0, X0: 10, X1: 100, Top: 10, Bottom: 30}, + {Text: "table box", LayoutType: pdf.LayoutTypeTable, PageNumber: 0, X0: 10, X1: 400, Top: 40, Bottom: 150}, + {Text: "outro", LayoutType: pdf.LayoutTypeText, PageNumber: 0, X0: 10, X1: 100, Top: 160, Bottom: 190}, + } + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 10, Right: 400, Top: 40, Bottom: 150}}, + RegionLeft: 10, RegionRight: 400, RegionTop: 40, RegionBottom: 150, + Scale: 1.0, + Cells: []pdf.TSRCell{ + {X0: 0, Y0: 0, X1: 100, Y1: 50, Text: "A"}, + {X0: 100, Y0: 0, X1: 200, Y1: 50, Text: "B"}, + }, + }, + } + + result := ExtractTableAndReplace(boxes, tables) + if len(result) != 3 { + t.Errorf("expected 3 boxes, got %d", len(result)) + } +} + +func TestExtractTableAndReplace_NoTables(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "text1", PageNumber: 0}, + {Text: "text2", PageNumber: 0}, + } + tables := []pdf.TableItem{} + + result := ExtractTableAndReplace(boxes, tables) + if len(result) != 2 { + t.Errorf("expected 2 boxes unchanged, got %d", len(result)) + } +} + +func TestExtractTableAndReplace_DataSourceBox(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "数据来源: somewhere", LayoutType: pdf.LayoutTypeTable, PageNumber: 0}, + {Text: "normal text", LayoutType: pdf.LayoutTypeText, PageNumber: 0}, + } + tables := []pdf.TableItem{} + + result := ExtractTableAndReplace(boxes, tables) + if len(result) != 1 { + t.Errorf("expected 1 box (data source removed), got %d", len(result)) + } + if result[0].Text != "normal text" { + t.Errorf("expected normal text to remain, got %q", result[0].Text) + } +} + +func TestExtractTableAndReplace_ZeroBoxesWithTables(t *testing.T) { + boxes := []pdf.TextBox{} + tables := []pdf.TableItem{ + { + Scale: 1.0, + Cells: []pdf.TSRCell{ + {X0: 0, Y0: 0, X1: 100, Y1: 50, Text: "A"}, + }, + }, + } + + result := ExtractTableAndReplace(boxes, tables) + if len(result) != 1 { + t.Errorf("expected 1 table box for zero input boxes, got %d", len(result)) + } +} + +// ============================================ +// 第五部分：FilterBoxesByRemoveSet 单元测试 +// ============================================ + +func TestFilterBoxesByRemoveSet_EmptyRemoveSet(t *testing.T) { + boxes := []pdf.TextBox{{Text: "a"}, {Text: "b"}, {Text: "c"}} + removeSet := map[int]bool{} + + result := FilterBoxesByRemoveSet(boxes, removeSet) + if len(result) != 3 { + t.Errorf("expected all boxes to remain, got %d", len(result)) + } +} + +func TestFilterBoxesByRemoveSet_RemoveSome(t *testing.T) { + boxes := []pdf.TextBox{{Text: "keep0"}, {Text: "remove1"}, {Text: "keep2"}, {Text: "remove3"}} + removeSet := map[int]bool{1: true, 3: true} + + result := FilterBoxesByRemoveSet(boxes, removeSet) + if len(result) != 2 { + t.Errorf("expected 2 boxes, got %d", len(result)) + } + if result[0].Text != "keep0" || result[1].Text != "keep2" { + t.Errorf("unexpected filtered result: %+v", result) + } +} + +func TestFilterBoxesByRemoveSet_RemoveAll(t *testing.T) { + boxes := []pdf.TextBox{{Text: "a"}, {Text: "b"}} + removeSet := map[int]bool{0: true, 1: true} + + result := FilterBoxesByRemoveSet(boxes, removeSet) + if len(result) != 0 { + t.Errorf("expected empty result, got %d", len(result)) + } +} + +func TestFilterBoxesByRemoveSet_EmptyInput(t *testing.T) { + var boxes []pdf.TextBox + removeSet := map[int]bool{0: true} + + result := FilterBoxesByRemoveSet(boxes, removeSet) + if len(result) != 0 { + t.Errorf("expected empty result for empty input, got %d", len(result)) + } +} + +func TestFilterBoxesByRemoveSet_Preallocation(t *testing.T) { + // 验证容量预分配是否合理 + boxes := make([]pdf.TextBox, 100) + removeSet := map[int]bool{} + for i := 0; i < 30; i++ { + removeSet[i] = true // 标记 30 个要移除 + } + + result := FilterBoxesByRemoveSet(boxes, removeSet) + if len(result) != 70 { + t.Errorf("expected 70 boxes, got %d", len(result)) + } + // 验证容量至少为 70 + if cap(result) < 70 { + t.Errorf("expected capacity >= 70, got %d", cap(result)) + } +} + +// ============================================ +// 第六部分：createTableBoxFromItem 单元测试 +// ============================================ + +func TestCreateTableBoxFromItem_Basic(t *testing.T) { + table := &pdf.TableItem{ + RegionLeft: 10, + RegionRight: 400, + RegionTop: 60, + RegionBottom: 200, + Positions: []pdf.Position{{ + PageNumbers: []int{1}, + }}, + } + + box := createTableBoxFromItem(table, "test
") + if box.Text != "test
" { + t.Errorf("expected HTML text, got %q", box.Text) + } + if box.LayoutType != pdf.LayoutTypeTable { + t.Errorf("expected table layout, got %v", box.LayoutType) + } + if box.PageNumber != 1 { + t.Errorf("expected page 1, got %d", box.PageNumber) + } + if box.X0 != 10 || box.X1 != 400 || box.Top != 60 || box.Bottom != 200 { + t.Errorf("expected correct coordinates, got (%.0f,%.0f,%.0f,%.0f)", box.X0, box.X1, box.Top, box.Bottom) + } +} + +func TestCreateTableBoxFromItem_FallbackToPosition(t *testing.T) { + // Region 字段为空时，使用 Position 的坐标 + table := &pdf.TableItem{ + Positions: []pdf.Position{{ + PageNumbers: []int{2}, + Left: 20, + Right: 300, + Top: 50, + Bottom: 150, + }}, + } + + box := createTableBoxFromItem(table, "fallback
") + if box.X0 != 20 || box.X1 != 300 || box.Top != 50 || box.Bottom != 150 { + t.Errorf("expected fallback coordinates, got (%.0f,%.0f,%.0f,%.0f)", box.X0, box.X1, box.Top, box.Bottom) + } +} + +func TestCreateTableBoxFromItem_EmptyPositions(t *testing.T) { + // 没有 Positions 时也能工作 + table := &pdf.TableItem{ + RegionLeft: 10, + RegionRight: 100, + RegionTop: 10, + RegionBottom: 100, + } + + box := createTableBoxFromItem(table, "empty-pos
") + if box.PageNumber != 0 { + t.Errorf("expected page 0, got %d", box.PageNumber) + } +} + +// ============================================ +// 第七部分：handleImageOnlyPDFs 单元测试 +// ============================================ + +func TestHandleImageOnlyPDFs_EmptyTables(t *testing.T) { + result := handleImageOnlyPDFs([]pdf.TableItem{}) + if len(result) != 0 { + t.Errorf("expected empty result, got %d boxes", len(result)) + } +} + +func TestHandleImageOnlyPDFs_EmptyCells(t *testing.T) { + tables := []pdf.TableItem{ + {Cells: []pdf.TSRCell{}}, // 没有 cell 的 table + } + result := handleImageOnlyPDFs(tables) + if len(result) != 0 { + t.Errorf("expected no boxes for empty cells, got %d", len(result)) + } +} + +func TestHandleImageOnlyPDFs_SingleTable(t *testing.T) { + tables := []pdf.TableItem{ + { + Scale: 1.0, + CropOffX: 0, + CropOffY: 0, + RegionLeft: 10, + RegionRight: 200, + RegionTop: 20, + RegionBottom: 100, + Positions: []pdf.Position{{PageNumbers: []int{0}}}, + Cells: []pdf.TSRCell{ + {X0: 0, Y0: 0, X1: 100, Y1: 50, Text: "cell1"}, + }, + }, + } + result := handleImageOnlyPDFs(tables) + if len(result) != 1 { + t.Errorf("expected 1 box, got %d", len(result)) + } + if result[0].LayoutType != pdf.LayoutTypeTable { + t.Error("expected table layout type") + } +} + +func TestHandleImageOnlyPDFs_MultipleTables(t *testing.T) { + tables := []pdf.TableItem{ + { + Scale: 1.0, + RegionLeft: 10, RegionRight: 200, + RegionTop: 20, RegionBottom: 100, + Positions: []pdf.Position{{PageNumbers: []int{0}}}, + Cells: []pdf.TSRCell{{Text: "table1"}}, + }, + { + // 没有 cell 的 table，应该被跳过 + Cells: []pdf.TSRCell{}, + }, + { + Scale: 1.0, + RegionLeft: 10, RegionRight: 200, + RegionTop: 120, RegionBottom: 200, + Positions: []pdf.Position{{PageNumbers: []int{1}}}, + Cells: []pdf.TSRCell{{Text: "table2"}}, + }, + } + result := handleImageOnlyPDFs(tables) + if len(result) != 2 { + t.Errorf("expected 2 boxes, got %d", len(result)) + } +} + +// ============================================ +// 阶段 2: buildAndSortAnchors 和 processTablesWithReplacements +// ============================================ + +func TestBuildAndSortAnchors(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "text1", PageNumber: 0, Top: 10}, + {Text: "table1", LayoutType: pdf.LayoutTypeTable, PageNumber: 0, Top: 50}, + {Text: "text2", PageNumber: 0, Top: 100}, + } + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 10, Right: 400, Top: 50, Bottom: 80}}, + RegionLeft: 10, RegionRight: 400, RegionTop: 50, RegionBottom: 80, + Cells: []pdf.TSRCell{{Text: "cell1"}}, + }, + } + _, replacements := buildReplacements(boxes, tables) + + anchors := findTableAnchorsWithReplacements(boxes, tables, replacements) + result := buildAndSortAnchors(anchors) + + if len(result) != 1 { + t.Errorf("expected 1 anchor, got %d", len(result)) + } +} + +// ============================================ +// 第八部分：ConsolidateFigures 子函数的单元测试 +// ============================================ + +func TestMarkDataSourceBoxesForRemoval(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "数据来源: test", LayoutType: pdf.LayoutTypeFigure, PageNumber: 0}, + {Text: "资料来源：abc", LayoutType: pdf.LayoutTypeFigure, PageNumber: 0}, + {Text: "图表来源 def", LayoutType: pdf.LayoutTypeFigure, PageNumber: 0}, + {Text: "正常图片内容", LayoutType: pdf.LayoutTypeFigure, PageNumber: 0}, + {Text: "数据来源: 不应该移除", LayoutType: pdf.LayoutTypeText, PageNumber: 0}, // 不是 figure 类型 + } + + removeSet := markDataSourceBoxesForRemoval(boxes) + if len(removeSet) != 3 { + t.Errorf("expected 3 boxes marked for removal, got %d", len(removeSet)) + } + if !removeSet[0] || !removeSet[1] || !removeSet[2] { + t.Error("expected boxes 0, 1, 2 to be marked for removal") + } + if removeSet[3] || removeSet[4] { + t.Error("expected boxes 3 and 4 NOT to be marked") + } +} + +func TestGroupFigureBoxes(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "fig1-part1", LayoutType: pdf.LayoutTypeFigure, PageNumber: 0, LayoutNo: "fig-0"}, + {Text: "fig1-part2", LayoutType: pdf.LayoutTypeFigure, PageNumber: 0, LayoutNo: "fig-0"}, + {Text: "fig2", LayoutType: pdf.LayoutTypeFigure, PageNumber: 0, LayoutNo: "fig-1"}, + {Text: "fig3", LayoutType: pdf.LayoutTypeFigure, PageNumber: 1, LayoutNo: "fig-0"}, // 不同页面 + {Text: "text", LayoutType: pdf.LayoutTypeText, PageNumber: 0}, + } + removeSet := map[int]bool{} + + groups := groupFigureBoxes(boxes, removeSet) + if len(groups) != 3 { + t.Errorf("expected 3 groups, got %d", len(groups)) + } + + // 验证组的内容 + key1 := figKey{page: 0, ln: "fig-0"} + if len(groups[key1]) != 2 { + t.Errorf("expected 2 boxes in fig-0 group, got %d", len(groups[key1])) + } + key2 := figKey{page: 0, ln: "fig-1"} + if len(groups[key2]) != 1 { + t.Errorf("expected 1 box in fig-1 group, got %d", len(groups[key2])) + } + key3 := figKey{page: 1, ln: "fig-0"} + if len(groups[key3]) != 1 { + t.Errorf("expected 1 box in page 1 fig-0 group, got %d", len(groups[key3])) + } +} + +func TestMergeFigureGroups(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "part1", LayoutType: pdf.LayoutTypeFigure, PageNumber: 0, LayoutNo: "fig-0", + X0: 10, X1: 100, Top: 10, Bottom: 50}, + {Text: "part2", LayoutType: pdf.LayoutTypeFigure, PageNumber: 0, LayoutNo: "fig-0", + X0: 50, X1: 150, Top: 30, Bottom: 80}, + {Text: "single", LayoutType: pdf.LayoutTypeFigure, PageNumber: 0, LayoutNo: "fig-1", + X0: 200, X1: 300, Top: 10, Bottom: 50}, + } + removeSet := make(map[int]bool) + groups := map[figKey][]int{ + {page: 0, ln: "fig-0"}: {0, 1}, + {page: 0, ln: "fig-1"}: {2}, + } + + mergeFigureGroups(boxes, groups, removeSet) + + // 验证合并后的结果 + if boxes[0].Text != "part1\npart2" { + t.Errorf("expected merged text, got %q", boxes[0].Text) + } + if boxes[0].X0 != 10 || boxes[0].X1 != 150 || boxes[0].Top != 10 || boxes[0].Bottom != 80 { + t.Error("expected merged bounding box") + } + if !removeSet[1] { + t.Error("expected box 1 to be marked for removal") + } + if removeSet[2] { + t.Error("expected single box NOT to be marked for removal") + } +} + +func TestConsolidateFigures_Integration(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "数据来源: test", LayoutType: pdf.LayoutTypeFigure, PageNumber: 0, LayoutNo: "fig-0"}, + {Text: "fig1-part1", LayoutType: pdf.LayoutTypeFigure, PageNumber: 0, LayoutNo: "fig-0", + X0: 10, X1: 100, Top: 10, Bottom: 50}, + {Text: "fig1-part2", LayoutType: pdf.LayoutTypeFigure, PageNumber: 0, LayoutNo: "fig-0", + X0: 50, X1: 150, Top: 30, Bottom: 80}, + {Text: "normal text", LayoutType: pdf.LayoutTypeText, PageNumber: 0}, + } + + result := ConsolidateFigures(boxes) + + // 验证结果 + if len(result) != 2 { // 合并后的 figure + normal text + t.Errorf("expected 2 boxes, got %d", len(result)) + } + + // 检查 figure 是否正确合并 + var figureFound bool + for _, b := range result { + if b.LayoutType == pdf.LayoutTypeFigure { + figureFound = true + if b.Text != "fig1-part1\nfig1-part2" { + t.Errorf("expected merged figure text, got %q", b.Text) + } + } + } + if !figureFound { + t.Error("expected figure box in result") + } +} + +func TestConsolidateFigures_NoFigures(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "text1", LayoutType: pdf.LayoutTypeText, PageNumber: 0}, + {Text: "text2", LayoutType: pdf.LayoutTypeText, PageNumber: 0}, + } + + result := ConsolidateFigures(boxes) + if len(result) != 2 { + t.Errorf("expected 2 boxes unchanged, got %d", len(result)) + } +} + +func TestConsolidateFigures_OnlyDataSource(t *testing.T) { + boxes := []pdf.TextBox{ + {Text: "数据来源: test", LayoutType: pdf.LayoutTypeFigure, PageNumber: 0}, + } + + result := ConsolidateFigures(boxes) + if len(result) != 0 { + t.Errorf("expected 0 boxes (data source removed), got %d", len(result)) + } +} + +func TestExtractTableAndReplace_MergeTablesAcrossPages(t *testing.T) { + // Regression test: two tables on consecutive pages with overlapping X + // should be merged by MergeTablesAcrossPages, and buildReplacementsAfterMerge + // must correctly index into the merged slice (not the original pre-merge slice). + boxes := []pdf.TextBox{ + {Text: "intro", LayoutType: pdf.LayoutTypeText, PageNumber: 0, X0: 10, X1: 100, Top: 10, Bottom: 30}, + {Text: "table1", LayoutType: pdf.LayoutTypeTable, PageNumber: 0, X0: 10, X1: 400, Top: 40, Bottom: 150}, + {Text: "middle", LayoutType: pdf.LayoutTypeText, PageNumber: 0, X0: 10, X1: 100, Top: 160, Bottom: 190}, + {Text: "table2", LayoutType: pdf.LayoutTypeTable, PageNumber: 1, X0: 10, X1: 400, Top: 10, Bottom: 120}, + {Text: "outro", LayoutType: pdf.LayoutTypeText, PageNumber: 1, X0: 10, X1: 100, Top: 130, Bottom: 160}, + } + tables := []pdf.TableItem{ + { + Positions: []pdf.Position{{PageNumbers: []int{0}, Left: 10, Right: 400, Top: 40, Bottom: 150}}, + RegionLeft: 10, RegionRight: 400, RegionTop: 40, RegionBottom: 150, + Scale: 1.0, + Cells: []pdf.TSRCell{ + {X0: 0, Y0: 0, X1: 100, Y1: 30, Text: "Page0_A"}, + {X0: 100, Y0: 0, X1: 200, Y1: 30, Text: "Page0_B"}, + }, + }, + { + Positions: []pdf.Position{{PageNumbers: []int{1}, Left: 10, Right: 400, Top: 10, Bottom: 120}}, + RegionLeft: 10, RegionRight: 400, RegionTop: 10, RegionBottom: 120, + Scale: 1.0, + Cells: []pdf.TSRCell{ + {X0: 0, Y0: 50, X1: 100, Y1: 80, Text: "Page1_C"}, + {X0: 100, Y0: 50, X1: 200, Y1: 80, Text: "Page1_D"}, + }, + }, + } + + result := ExtractTableAndReplace(boxes, tables) + if len(result) == 0 { + t.Fatal("expected non-empty result") + } + // After merge: 2 table boxes replaced by 1 merged HTML box. + // Original 5 boxes → 4 expected (intro, merged_table, middle, outro). + if len(result) != 4 { + t.Errorf("expected 4 boxes after merge+replace, got %d", len(result)) + } + // The merged HTML box should contain cells from both pages. + htmlBox := result[1] + if !strings.Contains(htmlBox.Text, "Page0") || !strings.Contains(htmlBox.Text, "Page1") { + t.Errorf("merged HTML should contain cells from both pages, got: %s", htmlBox.Text[:min(100, len(htmlBox.Text))]) + } + // Verify the original text boxes are preserved in the right order. + if result[0].Text != "intro" || result[2].Text != "middle" || result[3].Text != "outro" { + t.Error("non-table boxes should be preserved in original order") + } +} diff --git a/internal/deepdoc/parser/pdf/table/table_spans.go b/internal/deepdoc/parser/pdf/table/table_spans.go new file mode 100644 index 0000000000..9fc79f7712 --- /dev/null +++ b/internal/deepdoc/parser/pdf/table/table_spans.go @@ -0,0 +1,90 @@ +package table + +import ( + "strings" + + pdf "ragflow/internal/deepdoc/parser/pdf/type" +) + +// calSpans computes colspan and rowspan for spanning cells in the grid. +// Returns spanInfo (row,col → colspan,rowspan) and covered (cells hidden by spans). +// Matches Python's __cal_spans (table_structure_recognizer.py:535). +func CalSpans(rows [][]pdf.TSRCell) (map[[2]int][2]int, map[[2]int]bool) { + spanInfo := make(map[[2]int][2]int) + covered := make(map[[2]int]bool) + if len(rows) == 0 || len(rows[0]) == 0 { return spanInfo, covered } + + // Compute column center positions. + nCols := len(rows[0]) + colLeft := make([]float64, nCols) + colRight := make([]float64, nCols) + for j := 0; j < nCols; j++ { + colLeft[j] = 1e9 + colRight[j] = -1e9 + } + nRows := len(rows) + rowTop := make([]float64, nRows) + rowBott := make([]float64, nRows) + for i := 0; i < nRows; i++ { + rowTop[i] = 1e9 + rowBott[i] = -1e9 + } + + for i, row := range rows { + for j, cell := range row { + if j >= nCols { continue } + // Exclude spanning cells from column/row boundary calculations. + // Use label-based detection (O(1), no dependency on column midpoints). + if strings.Contains(cell.Label, "spanning") { continue } + if cell.X0 < colLeft[j] { colLeft[j] = cell.X0 } + if cell.X1 > colRight[j] { colRight[j] = cell.X1 } + if cell.Y0 < rowTop[i] { rowTop[i] = cell.Y0 } + if cell.Y1 > rowBott[i] { rowBott[i] = cell.Y1 } + } + } + + // For each spanning cell, compute how many cols/rows it covers. + for i, row := range rows { + for j, cell := range row { + if j >= nCols || covered[[2]int{i,j}] { continue } + // Skip cells without position data (they can't span). + if cell.X0 == 0 && cell.X1 == 0 && cell.Y0 == 0 && cell.Y1 == 0 { continue } + cs, rs := 1, 1 + // Count columns whose center is inside this cell's X range. + for k := j+1; k < nCols; k++ { + // Skip columns with no non-spanning cells (initial values unchanged). + if colLeft[k] == 1e9 && colRight[k] == -1e9 { continue } + colCenter := (colLeft[k] + colRight[k]) / 2 + if colCenter >= cell.X0 && colCenter <= cell.X1 { cs++ } + } + // Count rows whose center is inside this cell's Y range. + for k := i+1; k < nRows; k++ { + // Skip rows with no non-spanning cells. + if rowTop[k] == 1e9 && rowBott[k] == -1e9 { continue } + rowCenter := (rowTop[k] + rowBott[k]) / 2 + if rowCenter >= cell.Y0 && rowCenter <= cell.Y1 { rs++ } + } + if cs > 1 || rs > 1 { + spanInfo[[2]int{i,j}] = [2]int{cs, rs} + // Mark covered cells. + for ri := i; ri < i+rs && ri < nRows; ri++ { + for cj := j; cj < j+cs && cj < nCols; cj++ { + if ri != i || cj != j { + covered[[2]int{ri, cj}] = true + } + } + } + } + } + } + return spanInfo, covered +} + +// flattenGrid flattens a 2D grid into a 1D slice for fillCellTextFromBoxes. +func FlattenGrid(grid [][]pdf.TSRCell) []pdf.TSRCell { + n := 0 + for _, row := range grid { n += len(row) } + flat := make([]pdf.TSRCell, 0, n) + for _, row := range grid { flat = append(flat, row...) } + return flat +} diff --git a/internal/deepdoc/parser/pdf/table/table_spans_test.go b/internal/deepdoc/parser/pdf/table/table_spans_test.go new file mode 100644 index 0000000000..389b13dc98 --- /dev/null +++ b/internal/deepdoc/parser/pdf/table/table_spans_test.go @@ -0,0 +1,46 @@ + +package table + +import ( + "testing" + + pdf "ragflow/internal/deepdoc/parser/pdf/type" +) + +func TestCalSpans_NonSpanningCellsNotPolluted(t *testing.T) { + // Simulate the SpannedTable test grid: row 0 has Q1(regular), 部门开支汇总(span), Q2(regular) + rows := [][]pdf.TSRCell{ + { + {X0: 0, Y0: 0, X1: 100, Y1: 30, Text: "Q1", Label: "table row"}, + {X0: 0, Y0: 0, X1: 200, Y1: 30, Text: "部门开支汇总", Label: "table spanning cell"}, + {X0: 101, Y0: 0, X1: 200, Y1: 30, Text: "Q2", Label: "table row"}, + }, + { + {X0: 0, Y0: 35, X1: 100, Y1: 65, Text: "100", Label: "table row"}, + {X0: 101, Y0: 35, X1: 200, Y1: 65, Text: "200", Label: "table row"}, + }, + } + + spans, covered := CalSpans(rows) + + // Q1 at [0,0] has X0=0, X1=100 which should only cover its own column. + // It should NOT get a colspan. + if s, ok := spans[[2]int{0, 0}]; ok { + t.Errorf("Q1 at [0,0] should NOT have colspan, got %v. "+ + "Spanning cell at [0,1] polluted column boundaries", s) + } + + // 部门开支汇总 at [0,1] has X0=0, X1=200 which DOES span columns 0 and 1. + if s, ok := spans[[2]int{0, 1}]; !ok { + t.Error("部门开支汇总 at [0,1] should have colspan=2 (covers X=0-200)") + } else if s[0] != 2 { + t.Errorf("部门开支汇总 colspan = %d, want 2", s[0]) + } + + // Q2 at [0,2] should be covered by the spanning cell (col 2 is within X=0-200). + if !covered[[2]int{0, 2}] { + t.Error("Q2 at [0,2] should be covered by spanning cell at [0,1]") + } + + t.Logf("spans: %v, covered: %v", spans, covered) +}