zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-06 11:28:38 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor: pdf parser (#16625)

Browse Source 
### Summary

PDF parser refactor
This commit is contained in:
Jack
2026-07-05 20:45:35 +08:00
committed by GitHub
parent 014c3f634f
commit 1d3c100acb
17 changed files with 3988 additions and 1758 deletions
Download Patch File Download Diff File Expand all files Collapse all files

461
internal/deepdoc/parser/pdf/layout/layout.go
View File

@@ -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 {

435
internal/deepdoc/parser/pdf/layout/layout_test.go
View File

@@ -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)
}
}

381
internal/deepdoc/parser/pdf/parser.go
View File

@@ -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...)
}

63
internal/deepdoc/parser/pdf/parser_ocr.go
View File

@@ -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 {

212
internal/deepdoc/parser/pdf/parser_test.go
View File

@@ -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: 交通工具等级<table>") {
if !strings.HasPrefix(sections[0].Text, "表1: 交通工具等级<table") {
t.Errorf("expected caption before table HTML, got %q", sections[0].Text)
}
}
@@ -622,3 +640,175 @@ func TestParseRaw_ZeroZoom_NoNaN(t *testing.T) {
t.Fatal("expected at least one position to validate")
}
}
// ── Test for refactored helper functions ──────────────────────────────
// These are simple tests for the helper functions to ensure they work.
func TestParser_getBatchSize(t *testing.T) {
tests := []struct {
name string
batchSize int
want int
}{
{
name: "positive batch size",
batchSize: 100,
want: 100,
},
{
name: "zero batch size",
batchSize: 0,
want: 50,
},
{
name: "negative batch size",
batchSize: -1,
want: 50,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
p := &Parser{Config: pdf.ParserConfig{BatchSize: tt.batchSize}}
got := p.getBatchSize()
if got != tt.want {
t.Errorf("getBatchSize() = %v, want %v", got, tt.want)
}
})
}
}
func TestParser_setupPageConcurrency(t *testing.T) {
tests := []struct {
name string
maxConcurrency int
wantCap int
}{
{
name: "positive concurrency",
maxConcurrency: 5,
wantCap: 5,
},
{
name: "zero concurrency",
maxConcurrency: 0,
wantCap: 1,
},
{
name: "negative concurrency",
maxConcurrency: -1,
wantCap: 1,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
p := &Parser{Config: pdf.ParserConfig{MaxOCRConcurrency: tt.maxConcurrency}}
sem, wg := p.setupPageConcurrency()
if cap(sem) != tt.wantCap {
t.Errorf("setupPageConcurrency() sem cap = %v, want %v", cap(sem), tt.wantCap)
}
if wg == nil {
t.Error("setupPageConcurrency() wg should not be nil")
}
})
}
}
func TestParser_prescanPages(t *testing.T) {
// Provide at least 30 ASCII chars per page so DetectEnglish finds a run ≥30.
charsPage0 := make([]pdf.TextChar, 50)
for i := range charsPage0 {
charsPage0[i] = pdf.TextChar{
Text: "a", PageNumber: 0, Top: 10, Bottom: 20,
X0: 50 + float64(i*10), X1: 58 + float64(i*10),
}
}
charsPage1 := make([]pdf.TextChar, 50)
for i := range charsPage1 {
charsPage1[i] = pdf.TextChar{
Text: "b", PageNumber: 1, Top: 10, Bottom: 20,
X0: 50 + float64(i*10), X1: 58 + float64(i*10),
}
}
eng := &MockEngine{
NumPages: 2,
Chars: map[int][]pdf.TextChar{
0: charsPage0,
1: charsPage1,
},
}
p := &Parser{Config: pdf.DefaultParserConfig()}
prescanChars, prescanMedianH, prescanMedianW, isEnglish, scanNoise := p.prescanPages(context.Background(), eng, 0, 1, 2)
if _, ok := prescanChars[0]; !ok {
t.Error("prescanPages() prescanChars should contain page 0")
}
if _, ok := prescanMedianH[0]; !ok {
t.Error("prescanPages() prescanMedianH should contain page 0")
}
if _, ok := prescanMedianW[0]; !ok {
t.Error("prescanPages() prescanMedianW should contain page 0")
}
if isEnglish != true {
t.Errorf("prescanPages() isEnglish = %v, want true", isEnglish)
}
if scanNoise != false {
t.Errorf("prescanPages() scanNoise = %v, want false", scanNoise)
}
}
func TestParser_mergeBatchResults(t *testing.T) {
p := &Parser{Config: pdf.DefaultParserConfig()}
result := &pdf.ParseResult{
Sections: []pdf.Section{{Text: "section1"}},
Tables: []pdf.TableItem{{}},
Metrics: pdf.PipelineMetrics{
BoxesInitial: 10,
BoxesTextMerge: 8,
BoxesVertMerge: 5,
BoxesFinal: 3,
TablesCount: 1,
},
}
batch := &pdf.ParseResult{
Sections: []pdf.Section{{Text: "section2"}},
Tables: []pdf.TableItem{{}},
PageImages: map[int]image.Image{1: image.NewRGBA(image.Rect(0, 0, 10, 10))},
Metrics: pdf.PipelineMetrics{
BoxesInitial: 20,
BoxesTextMerge: 15,
BoxesVertMerge: 10,
BoxesFinal: 6,
TablesCount: 2,
},
}
p.mergeBatchResults(result, batch)
if len(result.Sections) != 2 {
t.Errorf("mergeBatchResults() sections length = %v, want 2", len(result.Sections))
}
if len(result.Tables) != 2 {
t.Errorf("mergeBatchResults() tables length = %v, want 2", len(result.Tables))
}
if result.Metrics.BoxesInitial != 30 {
t.Errorf("mergeBatchResults() BoxesInitial = %v, want 30", result.Metrics.BoxesInitial)
}
if result.Metrics.BoxesTextMerge != 23 {
t.Errorf("mergeBatchResults() BoxesTextMerge = %v, want 23", result.Metrics.BoxesTextMerge)
}
if result.Metrics.BoxesVertMerge != 15 {
t.Errorf("mergeBatchResults() BoxesVertMerge = %v, want 15", result.Metrics.BoxesVertMerge)
}
if result.Metrics.BoxesFinal != 9 {
t.Errorf("mergeBatchResults() BoxesFinal = %v, want 9", result.Metrics.BoxesFinal)
}
if result.Metrics.TablesCount != 3 {
t.Errorf("mergeBatchResults() TablesCount = %v, want 3", result.Metrics.TablesCount)
}
if _, ok := result.PageImages[1]; !ok {
t.Error("mergeBatchResults() pageImages should contain page 1")
}
}

371
internal/deepdoc/parser/pdf/table/group_boxes.go Normal file
View File

@@ -0,0 +1,371 @@
package table
import (
"sort"
"strings"
pdf "ragflow/internal/deepdoc/parser/pdf/type"
)
// rb is a row-box entry that holds cell data during grid construction.
type rb struct {
row, col int
txt string
x0, y0, x1, y1 float64
label string
}
// GroupBoxesByRC groups text boxes into a cell grid by R/C annotations.
// Matches Python's construct_table: sort by R, 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.
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
}

350
internal/deepdoc/parser/pdf/table/group_boxes_test.go Normal file
View File

@@ -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)
}
}

850
internal/deepdoc/parser/pdf/table/table_construct.go
View File

@@ -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("<table>")
if caption != "" {
b.WriteString("<caption>")
b.WriteString(caption)
b.WriteString("</caption>")
}
for ri, row := range rows {
b.WriteString("<tr>")
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(tag)
b.WriteString(">")
}
b.WriteString("</tr>")
}
b.WriteString("</table>")
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 (<th>). Used by DOCX, XLSX, PPTX,
// and HTML parsers that produce [][]string directly.
func SimpleRowsToHTML(rows [][]string) string {
if len(rows) == 0 {
return "<table></table>"
}
nCols := 0
for _, row := range rows {
if len(row) > nCols {
nCols = len(row)
}
}
var b strings.Builder
b.WriteString("<table>")
for ri, row := range rows {
b.WriteString("<tr>")
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(tag)
b.WriteString(">")
}
b.WriteString("</tr>")
}
b.WriteString("</table>")
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
}

669
internal/deepdoc/parser/pdf/table/table_construct_test.go
View File

@@ -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 <th > instead of <td >.
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 <th > for header row.
item := &pdf.TableItem{}
html := ConstructTable(cells, nil, "", item)
// Header row should use <th >, data row <td >.
if !strings.Contains(html, "<th >") {
t.Errorf("expected <th > for header row. HTML: %s", html)
}
if strings.Count(html, "<th ") != 2 {
t.Errorf("expected 2 <th > cells, got %d. HTML: %s", strings.Count(html, "<th "), html)
}
if strings.Count(html, "<td ") != 2 {
t.Errorf("expected 2 <td > cells (data row), got %d", strings.Count(html, "<td "))
}
}
// TestExtractTableAndReplace_OnlyTableBoxes verifies that only boxes with
// LayoutType=="table" are passed to constructTable (Python: filters by layout_type).
func TestExtractTableAndReplace_OnlyTableBoxes(t *testing.T) {
boxes := []pdf.TextBox{
{X0: 10, X1: 90, Top: 0, Bottom: 30, Text: "A", R: 0, C: 0, LayoutType: "table"},
@@ -538,11 +392,8 @@ func TestExtractTableAndReplace_OnlyTableBoxes(t *testing.T) {
}
}
// TestFillCellText_RCOverSpatial verifies that R/C-based fill puts a
// box into exactly one cell (matching Python), unlike spatial fill which
// puts it into all overlapping cells.
func TestFillCellText_RCOverSpatial(t *testing.T) {
// Box at X=30..270 overlaps all 3 cells (>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 := "<table><tr><td >姓名</td><td >年龄</td></tr><tr><td >张三</td><td >25</td></tr></table>"
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 != "<table></table>" {
t.Errorf("expected '<table></table>', got %q", html)
}
})
t.Run("single cell", func(t *testing.T) {
rows := toCells([][]string{{"X"}})
html := RowsToHTML(rows, "", nil, nil, nil)
expected := "<table><tr><td >X</td></tr></table>"
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, "<table>") || !strings.HasSuffix(html, "</table>") {
t.Errorf("not valid HTML: %s", html)
}
if !strings.Contains(html, "<td >标职务</td>") {
t.Errorf("missing cell '标职务': %s", html)
}
if strings.Count(html, "<tr>") != 3 {
t.Errorf("expected 3 rows, got %d", strings.Count(html, "<tr>"))
}
})
}
// 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")
}
}

112
internal/deepdoc/parser/pdf/table/table_html.go Normal file
View File

@@ -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("<table>")
if caption != "" {
b.WriteString("<caption>")
b.WriteString(html.EscapeString(caption))
b.WriteString("</caption>")
}
for ri, row := range rows {
b.WriteString("<tr>")
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(tag)
b.WriteString(">")
}
b.WriteString("</tr>")
}
b.WriteString("</table>")
return b.String()
}
// SimpleRowsToHTML converts plain string-based table data to an HTML table.
// The first row is treated as a header (<th>). Used by DOCX, XLSX, PPTX,
// and HTML parsers that produce [][]string directly.
func SimpleRowsToHTML(rows [][]string) string {
if len(rows) == 0 {
return "<table></table>"
}
nCols := 0
for _, row := range rows {
if len(row) > nCols { nCols = len(row) }
}
var b strings.Builder
b.WriteString("<table>")
for ri, row := range rows {
b.WriteString("<tr>")
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(tag)
b.WriteString(">")
}
b.WriteString("</tr>")
}
b.WriteString("</table>")
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
}

108
internal/deepdoc/parser/pdf/table/table_html_test.go Normal file
View File

@@ -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 := "<table><tr><td >姓名</td><td >年龄</td></tr><tr><td >张三</td><td >25</td></tr></table>"
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 != "<table></table>" {
t.Errorf("expected '<table></table>', got %q", html)
}
})
t.Run("single cell", func(t *testing.T) {
rows := toCells([][]string{{"X"}})
html := RowsToHTML(rows, "", nil, nil, nil)
expected := "<table><tr><td >X</td></tr></table>"
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, "<table>") || !strings.HasSuffix(html, "</table>") {
t.Errorf("not valid HTML: %s", html)
}
if !strings.Contains(html, "<td >标职务</td>") {
t.Errorf("missing cell '标职务': %s", html)
}
if strings.Count(html, "<tr>") != 3 {
t.Errorf("expected 3 rows, got %d", strings.Count(html, "<tr>"))
}
})
}
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 <th > for header row.
item := &pdf.TableItem{}
html := ConstructTable(cells, nil, "", item)
// Header row should use <th >, data row <td >.
if !strings.Contains(html, "<th >") {
t.Errorf("expected <th > for header row. HTML: %s", html)
}
if strings.Count(html, "<th ") != 2 {
t.Errorf("expected 2 <th > cells, got %d. HTML: %s", strings.Count(html, "<th "), html)
}
if strings.Count(html, "<td ") != 2 {
t.Errorf("expected 2 <td > cells (data row), got %d", strings.Count(html, "<td "))
}
}
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)
}

128
internal/deepdoc/parser/pdf/table/table_merge.go Normal file
View File

@@ -0,0 +1,128 @@
package table
import (
"sort"
pdf "ragflow/internal/deepdoc/parser/pdf/type"
)
// 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(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
}

179
internal/deepdoc/parser/pdf/table/table_merge_test.go Normal file
View File

@@ -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))
}
}

489
internal/deepdoc/parser/pdf/table/table_post.go
View File

@@ -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.

802
internal/deepdoc/parser/pdf/table/table_post_test.go Normal file
View File

@@ -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: "<table>test</table>"}
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 != "<table>test</table>" {
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: "<table>test</table>"}
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: "<table>end</table>"}
result := insertTableBoxes(boxes, tables, removeSet, anchors, htmls)
if len(result) != 3 {
t.Errorf("expected 3 boxes, got %d", len(result))
}
if result[2].Text != "<table>end</table>" {
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: "<table>A</table>", 1: "<table>B</table>"}
result := insertTableBoxes(boxes, tables, removeSet, anchors, htmls)
if len(result) != 5 {
t.Errorf("expected 5 boxes, got %d", len(result))
}
if result[1].Text != "<table>A</table>" {
t.Errorf("expected table A at pos 1")
}
if result[4].Text != "<table>B</table>" {
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, "<table>test</table>")
if box.Text != "<table>test</table>" {
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, "<table>fallback</table>")
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, "<table>empty-pos</table>")
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")
}
}

90
internal/deepdoc/parser/pdf/table/table_spans.go Normal file
View File

@@ -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
}

46
internal/deepdoc/parser/pdf/table/table_spans_test.go Normal file
View File

@@ -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)
}