diff --git a/internal/ingestion/component/chunker/json_global_merge_test.go b/internal/ingestion/component/chunker/json_global_merge_test.go new file mode 100644 index 0000000000..c42d438964 --- /dev/null +++ b/internal/ingestion/component/chunker/json_global_merge_test.go @@ -0,0 +1,125 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, +// including without limitation the rights to use, copy, modify, merge, +// publish, distribute, sublicense, and/or sell copies of the Software, +// and to permit persons to whom the Software is furnished to do so, +// subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included +// in all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE +// LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +// WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +package chunker + +// TestJSONGlobalMergeMatchesPython pins the TokenChunker json path's +// cross-item merge: adjacent text chunks across JSON items collapse into one +// chunk when they jointly fit chunk_token_size, matching Python's +// _merge_text_chunks_by_token_size. +import ( + "context" + "strings" + "testing" +) + +func TestJSONGlobalMergeMatchesPython(t *testing.T) { + const budget = 128 + comp, err := NewTokenChunker(map[string]any{ + "chunk_token_size": float64(budget), + }) + if err != nil { + t.Fatalf("construct TokenChunker: %v", err) + } + + input := map[string]any{ + "name": "t", + "output_format": "json", + "json": []map[string]any{ + { + "text": "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu nu xi omicron pi rho sigma tau upsilon phi chi psi omega", + "doc_type_kwd": "text", + }, + { + "text": "another long item with many words that should be merged when the token budget allows merging across items on the python side but stays separate on the go side", + "doc_type_kwd": "text", + }, + }, + } + + out, err := comp.Invoke(context.Background(), nil, input) + if err != nil { + t.Fatalf("invoke TokenChunker: %v", err) + } + if msg, _ := out["_ERROR"].(string); msg != "" { + t.Fatalf("TokenChunker returned _ERROR: %s", msg) + } + chunks, _ := out["chunks"].([]map[string]any) + if len(chunks) != 1 { + t.Fatalf("want 1 chunk (global merge across items), got %d", len(chunks)) + } + text, _ := chunks[0]["text"].(string) + + first := "alpha beta gamma" + second := "another long item" + i1 := strings.Index(text, first) + i2 := strings.Index(text, second) + if i1 < 0 || i2 < 0 { + t.Errorf("merged chunk does not contain both items:\n%q", text) + } + // The two items must appear in source order within the single merged chunk. + if i1 >= i2 { + t.Errorf("merged chunk reordered items (want %q before %q):\n%q", first, second, text) + } +} + +// TestJSONSingleItemNotSubSplitMatchesPython pins the TokenChunker json path's +// handling of a single item that exceeds chunk_token_size: the over-budget +// item is kept whole (no sub-split), yielding exactly one chunk whose text +// equals the input verbatim. +func TestJSONSingleItemNotSubSplitMatchesPython(t *testing.T) { + const budget = 128 + comp, err := NewTokenChunker(map[string]any{ + "chunk_token_size": float64(budget), + }) + if err != nil { + t.Fatalf("construct TokenChunker: %v", err) + } + + // A single item far over the token budget (well above 128 tokens). + long := strings.Repeat("word ", 200) + input := map[string]any{ + "name": "t", + "output_format": "json", + "json": []map[string]any{ + {"text": long, "doc_type_kwd": "text"}, + }, + } + + out, err := comp.Invoke(context.Background(), nil, input) + if err != nil { + t.Fatalf("invoke TokenChunker: %v", err) + } + if msg, _ := out["_ERROR"].(string); msg != "" { + t.Fatalf("TokenChunker returned _ERROR: %s", msg) + } + chunks, _ := out["chunks"].([]map[string]any) + if len(chunks) != 1 { + t.Fatalf("want 1 chunk (over-budget single item kept whole), got %d", len(chunks)) + } + if got := strings.TrimSpace(chunks[0]["text"].(string)); got != strings.TrimSpace(long) { + t.Errorf("over-budget item not kept whole:\n got=%q\nwant=%q", got, strings.TrimSpace(long)) + } +} diff --git a/internal/ingestion/component/chunker/token.go b/internal/ingestion/component/chunker/token.go index 0bd85079ed..5a730f38da 100644 --- a/internal/ingestion/component/chunker/token.go +++ b/internal/ingestion/component/chunker/token.go @@ -326,7 +326,7 @@ func (c *TokenChunkerComponent) invokeTextPayload(_ context.Context, text string // Split-then-merge: split on delimiters, then greedily merge to // chunk_token_size with optional overlap. perItem := [][]schema.ChunkDoc{docs} - merged := mergeByTokenSizeFromJSON(perItem, c.param.ChunkTokenSize, c.param.OverlappedPercent) + merged := mergeByTokenSizeFromJSON(perItem, c.param.ChunkTokenSize, c.param.OverlappedPercent, true) return chunkOutputs(flatten(merged)) } @@ -635,7 +635,11 @@ func (c *TokenChunkerComponent) invokeJSONPayload(ctx context.Context, items []s // Otherwise split-then-merge: delimiter-split segments are greedily // merged to chunk_token_size with optional overlap. if !hasCustomDelim(c.param.Delimiters) { - attached = mergeByTokenSizeFromJSON(attached, c.param.ChunkTokenSize, c.param.OverlappedPercent) + // Python _merge_text_chunks_by_token_size merges adjacent text + // chunks across JSON items into one global token budget. Flatten the + // per-item structure into a single sequence first so the merge is + // global; non-text chunks still break the merge via their CKType. + attached = mergeByTokenSizeFromJSON([][]schema.ChunkDoc{flatten(attached)}, c.param.ChunkTokenSize, c.param.OverlappedPercent, false) } flat := flatten(attached) @@ -874,7 +878,7 @@ func takeFromStart(text string, tokens int) string { // hard cap (rag/nlp/__init__.py after the strict chunk_token_num fix). // Oversized text units are sub-split via splitOversizedUnit before merge; // overlap is applied only when overlap+segment still fits the budget. -func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, overlappedPct float64) [][]schema.ChunkDoc { +func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, overlappedPct float64, subSplitOversize bool) [][]schema.ChunkDoc { // overlappedPct is a [0,100] percentage. Clamp defensively because this // helper is also exercised directly by tests. if overlappedPct < 0 { @@ -947,7 +951,15 @@ func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, over addTextChunk(ck) continue } - // Hard-cap atomic oversize units before merge. + // Over-budget unit. + if !subSplitOversize { + // JSON path: Python keeps each over-budget item whole — it does + // not sub-split a single item, so emit it as one chunk. + addTextChunk(ck) + continue + } + // Text path: hard-cap atomic oversize units before merge, matching + // Python's _split_oversized_unit. slog.Debug("TokenChunker: splitting oversized JSON unit via splitOversizedUnit", "len", len(ck.Text), "tokens", tk, "chunk_token_size", chunkTokens) for _, piece := range splitOversizedUnit(ck.Text, chunkTokens) { diff --git a/internal/ingestion/component/chunker/token_batch1_test.go b/internal/ingestion/component/chunker/token_batch1_test.go index 0ec5fd41d4..eb4d659601 100644 --- a/internal/ingestion/component/chunker/token_batch1_test.go +++ b/internal/ingestion/component/chunker/token_batch1_test.go @@ -84,7 +84,7 @@ func TestMergeByTokenSizeFromJSON_OverlapStripsTags(t *testing.T) { {Text: bText, DocType: "text", CKType: "text", TKNums: intPtr(bN)}, }, } - got := mergeByTokenSizeFromJSON(items, budget, 30.0) + got := mergeByTokenSizeFromJSON(items, budget, 30.0, true) merged := got[0] if len(merged) != 2 { t.Fatalf("want 2 merged chunks (overlap path), got %d (a=%d b=%d budget=%d)", len(merged), aN, bN, budget) @@ -130,12 +130,12 @@ func clampOverlapFixture() [][]schema.ChunkDoc { } func TestMergeByTokenSizeFromJSON_ClampsOverlappedPct(t *testing.T) { - at100 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 100) + at100 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 100, true) if at100 == nil || len(at100) == 0 { t.Fatalf("overlappedPct=100: nil/empty result") } - at150 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 150) - atHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 1e300) + at150 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 150, true) + atHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 1e300, true) if !reflect.DeepEqual(at100, at150) { t.Errorf("overlappedPct=150 should clamp to 100; output differs from 100") } @@ -143,12 +143,12 @@ func TestMergeByTokenSizeFromJSON_ClampsOverlappedPct(t *testing.T) { t.Errorf("overlappedPct=1e300 should clamp to 100; output differs from 100") } - at0 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 0) + at0 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 0, true) if at0 == nil || len(at0) == 0 { t.Fatalf("overlappedPct=0: nil/empty result") } - atNeg := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -5) - atNegHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -1e300) + atNeg := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -5, true) + atNegHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -1e300, true) if !reflect.DeepEqual(at0, atNeg) { t.Errorf("overlappedPct=-5 should clamp to 0; output differs from 0") } @@ -169,7 +169,7 @@ func TestMergeByTokenSizeFromJSON_EmptyPrevKeepsChunk(t *testing.T) { {Text: "keepme", DocType: "text", CKType: "text", TKNums: intPtr(5)}, }, } - got := mergeByTokenSizeFromJSON(items, 128, 0) + got := mergeByTokenSizeFromJSON(items, 128, 0, true) merged := got[0] if len(merged) != 1 { t.Fatalf("want 1 merged chunk, got %d", len(merged)) diff --git a/internal/ingestion/component/chunker/token_pdfpos_test.go b/internal/ingestion/component/chunker/token_pdfpos_test.go index 0c348cf308..d1eb951c7d 100644 --- a/internal/ingestion/component/chunker/token_pdfpos_test.go +++ b/internal/ingestion/component/chunker/token_pdfpos_test.go @@ -38,7 +38,7 @@ func TestMergeByTokenSizeFromJSON_ExtendsPDFPositions(t *testing.T) { {Text: "beta", DocType: "text", CKType: "text", TKNums: intPtr(5), PDFPositions: posB}, }, } - got := mergeByTokenSizeFromJSON(items, 128, 0) + got := mergeByTokenSizeFromJSON(items, 128, 0, true) merged := got[0] if len(merged) != 1 { t.Fatalf("want 1 merged chunk, got %d", len(merged)) @@ -63,7 +63,7 @@ func TestMergeByTokenSizeFromJSON_ExtendsPositions(t *testing.T) { {Text: "b", DocType: "text", CKType: "text", TKNums: intPtr(5), Positions: posB}, }, } - got := mergeByTokenSizeFromJSON(items, 128, 0) + got := mergeByTokenSizeFromJSON(items, 128, 0, true) combined := string(got[0][0].Positions) if !strings.Contains(combined, "1,2,3") || !strings.Contains(combined, "4,5,6") { t.Errorf("merged chunk dropped/omitted `positions`: %s", combined) @@ -103,7 +103,7 @@ func TestMergeByTokenSizeFromJSON_PositionsDecodeToMatrix(t *testing.T) { {Text: "b", DocType: "text", CKType: "text", TKNums: intPtr(5), Positions: posB}, }, } - got := mergeByTokenSizeFromJSON(items, 128, 0) + got := mergeByTokenSizeFromJSON(items, 128, 0, true) m := got[0][0].ToMap() raw, ok := m["positions"] if !ok { diff --git a/internal/ingestion/component/chunker/token_strict_cap_test.go b/internal/ingestion/component/chunker/token_strict_cap_test.go index 5f70881d16..0f99156df4 100644 --- a/internal/ingestion/component/chunker/token_strict_cap_test.go +++ b/internal/ingestion/component/chunker/token_strict_cap_test.go @@ -108,7 +108,7 @@ func TestMergeByTokenSizeFromJSON_StrictCapNoOvershoot(t *testing.T) { Text: text, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(text)), }) } - got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 0) + got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 0, true) merged := got[0] if len(merged) < 3 { t.Fatalf("want >=3 chunks, got %d", len(merged)) @@ -131,7 +131,7 @@ func TestMergeByTokenSizeFromJSON_OverlapDroppedAtOverflow(t *testing.T) { Text: text, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(text)), }) } - got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 20) + got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 20, true) for i, ck := range got[0] { if n := tokenizeStr(ck.Text); n > budget { t.Errorf("chunk %d exceeds budget with overlap: tokens=%d", i, n) @@ -146,7 +146,7 @@ func TestMergeByTokenSizeFromJSON_OversizedUnitIsSubSplit(t *testing.T) { items := [][]schema.ChunkDoc{{ {Text: long, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(long))}, }} - got := mergeByTokenSizeFromJSON(items, budget, 0) + got := mergeByTokenSizeFromJSON(items, budget, 0, true) if len(got[0]) < 2 { t.Fatalf("oversized unit must yield multiple chunks, got %d", len(got[0])) }