diff --git a/internal/common/logger.go b/internal/common/logger.go index 3d0cae1180..6ef806dcae 100644 --- a/internal/common/logger.go +++ b/internal/common/logger.go @@ -421,8 +421,13 @@ func ErrorCtx(ctx context.Context, msg string, err error, fields ...zap.Field) { Logger.Error(detail, append(sessionFields(ctx), fields...)...) } -// IsDebugEnabled returns true if debug logging is enabled. +// IsDebugEnabled returns true if debug logging is enabled. It is safe to call +// before InitLogger (e.g. in unit tests): an uninitialized level reports false, +// matching the production default of info-level logging. func IsDebugEnabled() bool { + if Logger == nil { + return false + } return atomicLevel.Enabled(zapcore.DebugLevel) } diff --git a/internal/engine/elasticsearch/chunk.go b/internal/engine/elasticsearch/chunk.go index 4f7e383374..42211851b6 100644 --- a/internal/engine/elasticsearch/chunk.go +++ b/internal/engine/elasticsearch/chunk.go @@ -3575,7 +3575,12 @@ func parseOrderByExpr(orderBy *types.OrderByExpr) []map[string]interface{} { direction = "desc" } - // Skip id field (cannot order by text field) + // Skip id field (cannot order by text field). The logical "id" column is + // shimmed from the ES _id in the response path, but ordering on _id + // requires indices.id_field_data, which is disabled by default on the + // bundled Elasticsearch 8.x deployment and would make every id-ordered + // query fail at runtime. Callers that need deep pagination must supply a + // real stored, sortable field instead of relying on the logical id. if field.Field == "id" { continue } diff --git a/internal/engine/elasticsearch/parse_orderby_test.go b/internal/engine/elasticsearch/parse_orderby_test.go new file mode 100644 index 0000000000..adcefa949f --- /dev/null +++ b/internal/engine/elasticsearch/parse_orderby_test.go @@ -0,0 +1,61 @@ +package elasticsearch + +import ( + "reflect" + "testing" + + "ragflow/internal/engine/types" +) + +// TestParseOrderByExprSkipsLogicalID guards that the logical "id" column is NOT +// translated into an Elasticsearch _id sort. Ordering on _id requires +// indices.id_field_data, which is disabled by default on the bundled +// Elasticsearch 8.x deployment; emitting a _id sort would make every id-ordered +// query (knowledge_compile merge reads, dataset structure graph) fail at runtime. +// The caller that needs deep pagination must supply a real stored, sortable field +// instead of relying on the logical id. +func TestParseOrderByExprSkipsLogicalID(t *testing.T) { + cases := []struct { + name string + expr *types.OrderByExpr + want []map[string]interface{} + }{ + { + name: "asc id is skipped (no _id sort)", + expr: (&types.OrderByExpr{}).Asc("id"), + want: nil, + }, + { + name: "desc id is skipped (no _id sort)", + expr: (&types.OrderByExpr{}).Desc("id"), + want: nil, + }, + { + name: "id mixed with other fields keeps the others, drops id", + expr: (&types.OrderByExpr{}).Asc("weight_flt").Asc("id"), + want: []map[string]interface{}{ + {"weight_flt": map[string]interface{}{"order": "asc", "unmapped_type": "float"}}, + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := parseOrderByExpr(tc.expr) + if !reflect.DeepEqual(got, tc.want) { + t.Fatalf("parseOrderByExpr(%v) = %#v, want %#v", tc.expr, got, tc.want) + } + }) + } +} + +// TestParseOrderByExprIDNotDroppedToIDMetadata is the explicit regression guard: +// parseOrderByExpr must never emit a {"_id": ...} sort for the logical id column. +func TestParseOrderByExprIDNotDroppedToIDMetadata(t *testing.T) { + got := parseOrderByExpr((&types.OrderByExpr{}).Asc("id")) + for _, sort := range got { + if _, ok := sort["_id"]; ok { + t.Fatalf("parseOrderByExpr(id) emitted a _id sort %#v; this breaks on default ES 8.x", sort) + } + } +} diff --git a/internal/ingestion/knowledge_compile/consumer.go b/internal/ingestion/knowledge_compile/consumer.go index b999f5148c..251615035f 100644 --- a/internal/ingestion/knowledge_compile/consumer.go +++ b/internal/ingestion/knowledge_compile/consumer.go @@ -438,20 +438,25 @@ func (c *Consumer) RebuildDataset(ctx context.Context, tenant, kb, mode string) if err != nil { return fmt.Errorf("knowledge_compile: rebuild list docs: %w", err) } + // Recover each document's compile variants from its compiled products. The + // read is batched (LoadDocProductsByDocs) and processed in bounded chunks so + // the in-flight product set never exceeds one batch (docProductsBatchSize) + // instead of the entire KB — the previous loop loaded one document at a time + // and released it; a single full-KB batch read would have retained every + // document's products (including vectors) until recovery finished. + variantsByDoc, err := c.recoverDocVariantsBatch(ctx, tenant, kb, docs) + if err != nil { + return fmt.Errorf("knowledge_compile: rebuild recover variants: %w", err) + } for _, docID := range docs { - variants, taskTypes, rerr := c.recoverDocTypes(ctx, tenant, kb, docID) - if rerr != nil { - // O2a hard failure: abort the rebuild so an unknown template kind - // cannot leave the doc's dataset-level products unrebuilt. - return fmt.Errorf("knowledge_compile: rebuild recover variants %s: %w", docID, rerr) - } - if len(variants) == 0 { + vr := variantsByDoc[docID] + if len(vr.variants) == 0 { // Documents without compiled products do not contribute to any // dataset-level variant and must not enqueue an empty completion // event that falls back to the legacy all-products path. continue } - if perr := rs.Publish(ctx, tenant, kb, docID, string(EventTypeCompleted), variants, taskTypes); perr != nil { + if perr := rs.Publish(ctx, tenant, kb, docID, string(EventTypeCompleted), vr.variants, vr.taskTypes); perr != nil { return fmt.Errorf("knowledge_compile: rebuild republish %s: %w", docID, perr) } } @@ -475,27 +480,11 @@ func (c *Consumer) defaultDocLister(ctx context.Context, tenant, kb string) ([]s return dao.NewDocumentDAO().ListIDsByKBIDWithOptions(ctx, c.db, dao.DocumentListOptions{KbID: kb}) } -// recoverDocTypes reconstructs the compile-type and frontend task-type sets a -// document produced by reading its persisted doc-level products and mapping -// each product's canonical compilation kind through the -// shared kind mappings. -// It is used by RebuildDataset (B1) so the republished completed events carry -// the variants needed to route the dataset-level re-compile (tree nav, structure -// merge, wiki merge) — a nil/empty Variants would fall back to the legacy -// unified path and tree nav could not be rebuilt. -// -// KindToVariant (O2a) is a fixed whitelist that HARD-FAILS on unknown kinds: a -// product whose kind is not in the whitelist is a corrupt/unsupported state and -// the rebuild must abort (returning an error) rather than silently drop the -// variant and leave the doc's dataset-level products unrebuilt. A product -// without an authoritative kind falls back to its reverse-mapped variant (which -// itself must be whitelist-valid, enforced by KwdToVariant in the Reader). The -// Both returned slices are sorted and de-duplicated. -func (c *Consumer) recoverDocTypes(ctx context.Context, tenant, kb, docID string) ([]string, []string, error) { - products, err := c.reader.LoadDocProducts(ctx, tenant, kb, docID) - if err != nil { - return nil, nil, fmt.Errorf("knowledge_compile: recover variants load doc %s: %w", docID, err) - } +// recoverDocTypesFrom maps a document's already-loaded compiled products to its +// variant and frontend task-type sets. Splitting the mapping from the read lets +// the rebuild path reuse a single batch read (LoadDocProductsByDocs) instead of +// one read per document. +func recoverDocTypesFrom(products []kccommon.Product) ([]string, []string, error) { seenVariants := map[string]struct{}{} seenTaskTypes := map[string]struct{}{} var variants []string @@ -505,7 +494,7 @@ func (c *Consumer) recoverDocTypes(ctx context.Context, tenant, kb, docID string v, err := kccommon.KindToVariant(p.Kind) if err != nil { // O2a hard failure: do not continue with an incomplete variant set. - return nil, nil, fmt.Errorf("knowledge_compile: recover variants doc %s kind %q: %w", docID, p.Kind, err) + return nil, nil, fmt.Errorf("knowledge_compile: recover variants kind %q: %w", p.Kind, err) } if _, dup := seenVariants[string(v)]; !dup { seenVariants[string(v)] = struct{}{} @@ -540,6 +529,43 @@ func (c *Consumer) recoverDocTypes(ctx context.Context, tenant, kb, docID string return variants, taskTypes, nil } +// docVariantResult holds the recovered compile variants and frontend task types +// for one document. +type docVariantResult struct { + variants []string + taskTypes []string +} + +// recoverDocVariantsBatch loads the compiled products for the given documents and +// maps each to its variant/task-type sets. Documents are walked in +// docProductsBatchSize chunks so the products held in memory at any moment are +// limited to a single batch instead of the whole KB (the previous loop loaded one +// document at a time and released it; a single full-KB batch read would retain +// every document's products — including vectors — until recovery finished). +func (c *Consumer) recoverDocVariantsBatch(ctx context.Context, tenant, kb string, docs []string) (map[string]docVariantResult, error) { + out := make(map[string]docVariantResult, len(docs)) + for start := 0; start < len(docs); start += docProductsBatchSize { + end := start + docProductsBatchSize + if end > len(docs) { + end = len(docs) + } + productsByDoc, err := c.reader.LoadDocProductsByDocs(ctx, tenant, kb, docs[start:end]) + if err != nil { + return nil, err + } + for _, docID := range docs[start:end] { + variants, taskTypes, rerr := recoverDocTypesFrom(productsByDoc[docID]) + if rerr != nil { + // O2a hard failure: abort the rebuild so an unknown template kind + // cannot leave the doc's dataset-level products unrebuilt. + return nil, fmt.Errorf("knowledge_compile: rebuild recover variants %s: %w", docID, rerr) + } + out[docID] = docVariantResult{variants: variants, taskTypes: taskTypes} + } + } + return out, nil +} + // filterWikiPageCandidates returns ONLY the wiki page products the dataset-level // wiki merge may fold. It drops every non-wiki product (tree/structure/nav) so // those never enter the wiki Dedup/KNN/WriteMerged path, and within the wiki @@ -885,11 +911,27 @@ func (c *Consumer) processBatch(ctx context.Context, tenant, kb, token string, e return err } entityMerged = append(entityMerged, topicMerged...) - if len(wikiDiff.affectedKeys) > 0 { - if err := c.deleteMissingScopedWikiPages(ctx, tenant, kb, token, wikiDiff.affectedKeys, entityMerged); err != nil { + // Load the dataset's merged wiki pages once and reuse the snapshot for both + // the scoped-page deletion and the topic-key seed below, instead of reading + // the same pages twice per merge step. + var existingPages []kccommon.Product + if pageReader, ok := c.reader.(mergedWikiPageReader); ok { + existingPages, err = pageReader.LoadMergedWikiPages(ctx, tenant, kb) + if err != nil { return err } } + deletedIDs := []string{} + if len(wikiDiff.affectedKeys) > 0 { + deletedIDs, err = c.deleteMissingScopedWikiPages(ctx, tenant, kb, token, wikiDiff.affectedKeys, entityMerged, existingPages) + if err != nil { + return err + } + } + excluded := make(map[string]struct{}, len(deletedIDs)) + for _, id := range deletedIDs { + excluded[id] = struct{}{} + } c.reportProgress(ctx, tenant, kb, token, taskTypes, 0.58, "merging_by_slug", fmt.Sprintf("Merging knowledge pages by slug: %d", len(entityMerged))) // Entity pages have stable page identity and must bypass topic-mode KNN and // Page-specific entity/concept merging. Only topic-mode pages continue below. @@ -905,23 +947,14 @@ func (c *Consumer) processBatch(ctx context.Context, tenant, kb, token string, e // Diagnostics: after batch dedup, verify the per-doc products still carry // their embedding before the KNN/merge path consumes cand.Vector. If // candidates show 0 vectors while incoming had vectors, Dedup is dropping - // them and both the KNN lookups and the merged rows lose embeddings. - { - incomingVec, candVec := 0, 0 - for _, p := range incoming { - if len(p.Vector) > 0 { - incomingVec++ - } - } - for _, p := range candidates { - if len(p.Vector) > 0 { - candVec++ - } - } + // them and both the KNN lookups and the merged rows lose embeddings. Gated to + // debug: the counting below allocates even in quiet production. + if common.IsDebugEnabled() { + inVec, candVec := computeDedupVectorAudit(incoming, candidates) common.Info("knowledge_compile: dedup vector audit", zap.String("kb_id", kb), zap.Int("incoming", len(incoming)), - zap.Int("incoming_with_vector", incomingVec), + zap.Int("incoming_with_vector", inVec), zap.Int("candidates", len(candidates)), zap.Int("candidates_with_vector", candVec)) } @@ -931,17 +964,9 @@ func (c *Consumer) processBatch(ctx context.Context, tenant, kb, token string, e // against the wiki_page rows WriteMerged emits (they differ when KNN merges a // candidate into an existing merged row, when DecideBatch returns distinct new // rows, or when candidates of the same variant collapse in-memory). This is - // what explains "39 wiki_page written vs 21 KNN requests". - { - byVariant := map[string]int{} - var wikiPageSlugs []string - for _, cand := range candidates { - v := string(cand.Variant) - byVariant[v]++ - if cand.Variant == kccommon.VariantWiki { - wikiPageSlugs = append(wikiPageSlugs, candidateIdentity(cand)) - } - } + // what explains "39 wiki_page written vs 21 KNN requests". Gated to debug. + if common.IsDebugEnabled() { + byVariant, wikiPageSlugs := computeCandidatesBreakdown(candidates) common.Info("knowledge_compile: candidates breakdown", zap.String("kb_id", kb), zap.Int("candidates", len(candidates)), @@ -953,11 +978,6 @@ func (c *Consumer) processBatch(ctx context.Context, tenant, kb, token string, e // mirroring _struct_doc_storage_dedup_batch. Candidates that KNN-hit // the same existing merged row are grouped so the row is merged/updated once // instead of being rewritten per candidate. - type mergeGroup struct { - existing kccommon.Product - candidates []kccommon.Product - score float64 - } groupsByID := make(map[string]*mergeGroup, len(candidates)) var ( unmatchedMu sync.Mutex @@ -967,24 +987,10 @@ func (c *Consumer) processBatch(ctx context.Context, tenant, kb, token string, e // Exact topic recall is independent of page-content KNN. A page can use // different wording from the existing topic while still belonging to the // same topic, so consult the current dataset pages before running KNN. - topicPagesByKey := make(map[string]kccommon.Product) - if pageReader, ok := c.reader.(mergedWikiPageReader); ok { - existingPages, err := pageReader.LoadMergedWikiPages(ctx, tenant, kb) - if err != nil { - return err - } - for _, page := range existingPages { - if !isTopicPage(page) { - continue - } - key := topicKey(productTopic(page)) - if key != "" { - if _, exists := topicPagesByKey[key]; !exists { - topicPagesByKey[key] = page - } - } - } - } + // Seed with the surviving topic pages (excluding the ones this step just + // deleted), reproducing the post-deletion snapshot the previous double-read + // produced without a second engine query. + topicPagesByKey := buildTopicPagesByKey(existingPages, excluded) // The KNN pass is docengine-bounded (vector search), not CPU-bounded, so we // fan it out across the shared global compilerPool (vCPU-sized). Output order // is irrelevant: merged rows are upserted by their idempotent dataset-level @@ -1116,25 +1122,13 @@ func (c *Consumer) processBatch(ctx context.Context, tenant, kb, token string, e // Diagnostics: summarize the KNN groups before the LLM merge decision so // the reader can see which existing merged rows were hit and by how many // candidates (e.g. whether a "吕布" candidate hit an existing 吕布 row and - // was then judged by the LLM). - { - type groupDump struct { - Existing string `json:"existing"` - Candidates []string `json:"candidates"` - Score float64 `json:"score"` - } - dumps := make([]groupDump, 0, len(groupsByID)) - for _, g := range groupsByID { - cands := make([]string, 0, len(g.candidates)) - for _, cand := range g.candidates { - cands = append(cands, candidateIdentity(cand)) - } - dumps = append(dumps, groupDump{Existing: candidateIdentity(g.existing), Candidates: cands, Score: g.score}) - } + // was then judged by the LLM). Gated to debug: the slice/map construction + // below allocates even in quiet production and is only useful when tracing. + if common.IsDebugEnabled() { common.Info("knowledge_compile: KNN groups for DecideBatch", zap.String("kb_id", kb), zap.Int("groups", len(groupsByID)), - zap.Any("groups_detail", dumps)) + zap.Any("groups_detail", computeKNNGroupDumps(groupsByID))) } batched := make([]MergeGroup, 0, len(groupsByID)) for _, g := range groupsByID { @@ -1404,11 +1398,16 @@ type mergedWikiPageDeleter interface { DeleteMergedWikiPages(ctx context.Context, tenant, kb string, ids []string) error } -func (c *Consumer) deleteMissingScopedWikiPages(ctx context.Context, tenant, kb, token string, affected map[string]struct{}, current []kccommon.Product) error { - reader, readerOK := c.reader.(mergedWikiPageReader) +// deleteMissingScopedWikiPages removes the merged wiki pages whose contribution +// was dropped by this batch (present in affectedKeys but no longer in current). +// It takes the already-loaded merged pages (existingPages) so the merge step can +// reuse a single LoadMergedWikiPages snapshot; the caller passes the same slice +// it uses to seed the topic-key map, excluding the returned deleted IDs. +func (c *Consumer) deleteMissingScopedWikiPages(ctx context.Context, tenant, kb, token string, affected map[string]struct{}, current, existingPages []kccommon.Product) ([]string, error) { + _, readerOK := c.reader.(mergedWikiPageReader) deleter, writerOK := c.writer.(mergedWikiPageDeleter) if !readerOK || !writerOK { - return fmt.Errorf("knowledge_compile: scoped Wiki cleanup is unsupported") + return nil, fmt.Errorf("knowledge_compile: scoped Wiki cleanup is unsupported") } present := make(map[string]struct{}, len(current)) for _, page := range current { @@ -1416,12 +1415,8 @@ func (c *Consumer) deleteMissingScopedWikiPages(ctx context.Context, tenant, kb, present[key] = struct{}{} } } - existing, err := reader.LoadMergedWikiPages(ctx, tenant, kb) - if err != nil { - return err - } ids := make([]string, 0) - for _, page := range existing { + for _, page := range existingPages { key := wikiPageMergeKey(page) if _, scoped := affected[key]; !scoped { continue @@ -1431,11 +1426,39 @@ func (c *Consumer) deleteMissingScopedWikiPages(ctx context.Context, tenant, kb, } } if len(ids) == 0 { - return nil + return nil, nil } - return c.withWriteLock(ctx, kb, token, func() error { + if err := c.withWriteLock(ctx, kb, token, func() error { return deleter.DeleteMergedWikiPages(ctx, tenant, kb, ids) - }) + }); err != nil { + return nil, err + } + return ids, nil +} + +// buildTopicPagesByKey indexes the topic-mode merged pages by their topic key, +// skipping any page whose ID is in excludedIDs. excludedIDs are the page IDs this +// process deletes in the same merge step; skipping them reproduces the +// post-deletion snapshot the merge step previously obtained via a second +// LoadMergedWikiPages call, without a second read. +func buildTopicPagesByKey(pages []kccommon.Product, excludedIDs map[string]struct{}) map[string]kccommon.Product { + out := make(map[string]kccommon.Product) + for _, page := range pages { + if !isTopicPage(page) { + continue + } + if _, del := excludedIDs[page.ID]; del { + continue + } + key := topicKey(productTopic(page)) + if key == "" { + continue + } + if _, exists := out[key]; !exists { + out[key] = page + } + } + return out } func refreshWikiProductVectors(ctx context.Context, tenant string, products []kccommon.Product) []kccommon.Product { @@ -1889,3 +1912,83 @@ func (c *Consumer) removeNavLocked(ctx context.Context, tenant, kb, token string return nil }) } + +// mergeGroup groups the KNN/LLM merge candidates that resolve to the same +// existing merged row, so the row is merged/updated once instead of being +// rewritten per candidate. +type mergeGroup struct { + existing kccommon.Product + candidates []kccommon.Product + score float64 +} + +// groupDump is the debug-only serialization of a mergeGroup for the KNN +// pre-LLM-merge diagnostic log. +type groupDump struct { + Existing string `json:"existing"` + Candidates []string `json:"candidates"` + Score float64 `json:"score"` +} + +// computeDedupVectorAudit counts how many incoming and post-dedup candidates +// still carry an embedding, surfacing whether Dedup is dropping vectors before +// the KNN/merge path consumes cand.Vector. +func computeDedupVectorAudit(incoming, candidates []kccommon.Product) (int, int) { + incomingVec, candVec := 0, 0 + for _, p := range incoming { + if len(p.Vector) > 0 { + incomingVec++ + } + } + for _, p := range candidates { + if len(p.Vector) > 0 { + candVec++ + } + } + return incomingVec, candVec +} + +// computeCandidatesBreakdown breaks the KNN-input set down by variant and lists +// the wiki_page slugs, so a reader can reconcile candidate counts against the +// merged rows WriteMerged emits. +func computeCandidatesBreakdown(candidates []kccommon.Product) (map[string]int, []string) { + byVariant := map[string]int{} + var wikiPageSlugs []string + for _, cand := range candidates { + v := string(cand.Variant) + byVariant[v]++ + if cand.Variant == kccommon.VariantWiki { + wikiPageSlugs = append(wikiPageSlugs, candidateIdentity(cand)) + } + } + return byVariant, wikiPageSlugs +} + +// computeDocProductsVectorAudit counts how many loaded per-doc products carry an +// embedding and the (deduped) vector dimensions, surfacing when +// VectorFromChunkMap fails to restore the vector. +func computeDocProductsVectorAudit(products []kccommon.Product) (int, map[int]int) { + vecCount := 0 + dims := map[int]int{} + for _, p := range products { + if dim := len(p.Vector); dim > 0 { + vecCount++ + dims[dim]++ + } + } + return vecCount, dims +} + +// computeKNNGroupDumps serializes the KNN groups before the LLM merge decision +// for the debug-only diagnostic log. +func computeKNNGroupDumps(groups map[string]*mergeGroup) []groupDump { + dumps := make([]groupDump, 0, len(groups)) + for _, g := range groups { + cands := make([]string, 0, len(g.candidates)) + for _, cand := range g.candidates { + cands = append(cands, candidateIdentity(cand)) + } + dumps = append(dumps, groupDump{Existing: candidateIdentity(g.existing), Candidates: cands, Score: g.score}) + } + return dumps +} diff --git a/internal/ingestion/knowledge_compile/fake_engine_test.go b/internal/ingestion/knowledge_compile/fake_engine_test.go index e0f694698f..bce9b24e67 100644 --- a/internal/ingestion/knowledge_compile/fake_engine_test.go +++ b/internal/ingestion/knowledge_compile/fake_engine_test.go @@ -17,7 +17,9 @@ package knowledge_compile import ( "context" + "fmt" + "ragflow/internal/common" "ragflow/internal/engine/types" "gorm.io/gorm" @@ -35,6 +37,16 @@ type fakeEngine struct { searchChunks []map[string]interface{} // lastSearchReq captures the most recent SearchRequest for assertions. lastSearchReq *types.SearchRequest + // searchReqs accumulates every SearchRequest in call order, so tests can + // assert on multi-search paths (e.g. chunked batch reads). + searchReqs []*types.SearchRequest + // searchCalls counts Search invocations. + searchCalls int + // windowed emulates the backend max_result_window: a request with no + // OrderBy whose from+size exceeds common.MAX_RESULT_WINDOW is rejected, the + // way Elasticsearch rejects from+size beyond index.max_result_window. With a + // stable OrderBy set, deep pagination routes to search_after and is allowed. + windowed bool // lastDeleteCond captures the most recent DeleteChunks condition. lastDeleteCond map[string]interface{} // deleteConds accumulates every DeleteChunks condition in call order, so a @@ -48,7 +60,88 @@ type fakeEngine struct { func (f *fakeEngine) Search(_ context.Context, req *types.SearchRequest) (*types.SearchResult, error) { f.lastSearchReq = req - return &types.SearchResult{Chunks: f.searchChunks}, nil + f.searchReqs = append(f.searchReqs, req) + f.searchCalls++ + // Honour the subset of filters the production engine applies, so tests can + // assert real filtering behaviour (doc_id IN + exists) instead of a raw + // pass-through of every canned chunk. + chunks := applyEngineFilters(f.searchChunks, req.Filter) + if f.windowed { + // Faithfully emulate the Elasticsearch/OpenSearch index.max_result_window: + // a request whose from+size exceeds the window is rejected. The real + // engines only relax this via a stable Sort + search_after; our batch + // reader deliberately avoids that path, so a pathological batch simply + // fails loudly instead of silently truncating or looping. The rejection + // does NOT depend on OrderBy being present, because the engine would + // reject the deep from+size regardless once search_after is unavailable. + if req.Offset+req.Limit > common.MAX_RESULT_WINDOW { + return nil, fmt.Errorf("Result window is too large, from + size must be <= %d", common.MAX_RESULT_WINDOW) + } + start := req.Offset + if start > len(chunks) { + start = len(chunks) + } + end := start + req.Limit + if end > len(chunks) { + end = len(chunks) + } + return &types.SearchResult{Chunks: chunks[start:end]}, nil + } + return &types.SearchResult{Chunks: chunks}, nil +} + +// applyEngineFilters mirrors the production engine's bool.filter handling for the +// clauses the knowledge_compile reader relies on: doc_id IN (terms) and exists. +// Other clauses (scope_kwd, type_kwd, slug_kwd, available_int, …) are +// pass-through. It never mutates the input slice. +func applyEngineFilters(chunks []map[string]interface{}, filter map[string]interface{}) []map[string]interface{} { + if len(filter) == 0 { + return chunks + } + out := chunks + if raw, ok := filter["doc_id"]; ok { + ids := docIDSet(raw) + if len(ids) > 0 { + kept := make([]map[string]interface{}, 0, len(out)) + for _, c := range out { + if id, _ := c["doc_id"].(string); ids[id] { + kept = append(kept, c) + } + } + out = kept + } + } + if field, ok := filter["exists"].(string); ok && field != "" { + kept := make([]map[string]interface{}, 0, len(out)) + for _, c := range out { + if _, present := c[field]; present { + kept = append(kept, c) + } + } + out = kept + } + return out +} + +// docIDSet normalises a doc_id filter value (string, []string, or []interface{} of +// strings) into a lookup set. +func docIDSet(raw interface{}) map[string]bool { + out := map[string]bool{} + switch v := raw.(type) { + case string: + out[v] = true + case []string: + for _, s := range v { + out[s] = true + } + case []interface{}: + for _, s := range v { + if str, ok := s.(string); ok { + out[str] = true + } + } + } + return out } func (f *fakeEngine) DeleteChunks(_ context.Context, condition map[string]interface{}, _, _ string) (int64, error) { diff --git a/internal/ingestion/knowledge_compile/pr7_efficiency_test.go b/internal/ingestion/knowledge_compile/pr7_efficiency_test.go new file mode 100644 index 0000000000..22ec8c7e35 --- /dev/null +++ b/internal/ingestion/knowledge_compile/pr7_efficiency_test.go @@ -0,0 +1,273 @@ +// 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. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// 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. +// See the License for the specific language governing permissions and +// limitations under the License. + +package knowledge_compile + +import ( + "context" + "reflect" + "testing" + + "ragflow/internal/engine/types" + kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// --------------------------------------------------------------------------- +// Change 2: LoadDocProductsByDocs batch reader. +// --------------------------------------------------------------------------- + +func TestLoadDocProductsByDocsGroupsByDoc(t *testing.T) { + eng := &fakeEngine{searchChunks: []map[string]interface{}{ + {"id": "p1", "doc_id": "d1", "compile_kwd": "wiki_page", "available_int": 1, "content_with_weight": "c1"}, + {"id": "p2", "doc_id": "d2", "compile_kwd": "wiki_page", "available_int": 1, "content_with_weight": "c2"}, + {"id": "p3", "doc_id": "d1", "compile_kwd": "structure", "available_int": 1, "content_with_weight": "c3"}, + }} + r := engineReader{eng: eng} + got, err := r.LoadDocProductsByDocs(context.Background(), "tenant", "kb", []string{"d1", "d2"}) + if err != nil { + t.Fatalf("LoadDocProductsByDocs: %v", err) + } + filter, ok := eng.lastSearchReq.Filter["doc_id"] + if !ok { + t.Fatal("missing doc_id filter") + } + if !reflect.DeepEqual(filter, []interface{}{"d1", "d2"}) { + t.Fatalf("doc_id filter = %#v, want []interface{}{\"d1\",\"d2\"}", filter) + } + if len(got["d1"]) != 2 { + t.Fatalf("d1 products = %d, want 2", len(got["d1"])) + } + if len(got["d2"]) != 1 { + t.Fatalf("d2 products = %d, want 1", len(got["d2"])) + } +} + +func TestLoadDocProductsByDocsEmptySkipsSearch(t *testing.T) { + eng := &fakeEngine{searchChunks: []map[string]interface{}{ + {"id": "p1", "doc_id": "d1", "compile_kwd": "wiki_page", "available_int": 1, "content_with_weight": "c1"}, + }} + r := engineReader{eng: eng} + got, err := r.LoadDocProductsByDocs(context.Background(), "tenant", "kb", nil) + if err != nil { + t.Fatalf("LoadDocProductsByDocs: %v", err) + } + if len(got) != 0 { + t.Fatalf("empty docIDs must return empty map, got %d entries", len(got)) + } + if eng.searchCalls != 0 { + t.Fatalf("empty docIDs must not issue a Search, got %d calls", eng.searchCalls) + } +} + +func TestLoadDocProductsByDocsDocWithoutProducts(t *testing.T) { + eng := &fakeEngine{searchChunks: []map[string]interface{}{ + {"id": "p1", "doc_id": "d1", "compile_kwd": "wiki_page", "available_int": 1, "content_with_weight": "c1"}, + }} + r := engineReader{eng: eng} + // d2 has no matching rows but is requested: should be present with an empty slice. + got, err := r.LoadDocProductsByDocs(context.Background(), "tenant", "kb", []string{"d1", "d2"}) + if err != nil { + t.Fatalf("LoadDocProductsByDocs: %v", err) + } + if len(got["d1"]) != 1 { + t.Fatalf("d1 products = %d, want 1", len(got["d1"])) + } + if v := got["d2"]; len(v) != 0 { + t.Fatalf("d2 must be empty slice, got %d", len(v)) + } +} + +func TestLoadDocProductsByDocsChunksLargeSet(t *testing.T) { + eng := &fakeEngine{searchChunks: []map[string]interface{}{ + {"id": "p1", "doc_id": "d1", "compile_kwd": "wiki_page", "available_int": 1, "content_with_weight": "c1"}, + }} + r := engineReader{eng: eng} + docIDs := make([]string, 0, 1200) + for i := 0; i < 1200; i++ { + docIDs = append(docIDs, "doc-"+itoa(i)) + } + if _, err := r.LoadDocProductsByDocs(context.Background(), "tenant", "kb", docIDs); err != nil { + t.Fatalf("LoadDocProductsByDocs: %v", err) + } + // 1200 docs at a 500-doc batch must split into 3 searches. + if len(eng.searchReqs) != 3 { + t.Fatalf("search calls = %d, want 3 (chunked)", len(eng.searchReqs)) + } + wantSizes := []int{500, 500, 200} + for i, req := range eng.searchReqs { + filter, ok := req.Filter["doc_id"].([]interface{}) + if !ok { + t.Fatalf("req %d missing doc_id []interface{} filter", i) + } + if len(filter) != wantSizes[i] { + t.Fatalf("req %d filter size = %d, want %d", i, len(filter), wantSizes[i]) + } + } +} + +// --------------------------------------------------------------------------- +// Change 1: buildTopicPagesByKey excludes pages deleted in the same merge step. +// --------------------------------------------------------------------------- + +func TestBuildTopicPagesByKeyExcludesDeleted(t *testing.T) { + pages := []kccommon.Product{ + {ID: "keep", DocID: "kb", Meta: map[string]any{"page_type": "topic", "topic": "a"}}, + {ID: "del", DocID: "kb", Meta: map[string]any{"page_type": "topic", "topic": "b"}}, + {ID: "nt", DocID: "kb", Meta: map[string]any{"page_type": "entity", "topic": "c"}}, + } + // mark "del" as a topic page by isTopicPage; isTopicPage reads Meta page_type. + excluded := map[string]struct{}{"del": {}} + got := buildTopicPagesByKey(pages, excluded) + // "keep" topic page survives; "del" (deleted) and "nt" (non-topic) excluded. + if _, ok := got["a"]; !ok { + t.Fatalf("surviving topic page key \"a\" missing: %v", keysOf(got)) + } + if _, ok := got["b"]; ok { + t.Fatalf("deleted topic page key \"b\" must be excluded: %v", keysOf(got)) + } + if _, ok := got["c"]; ok { + t.Fatalf("non-topic page key \"c\" must be excluded: %v", keysOf(got)) + } +} + +// --------------------------------------------------------------------------- +// Change 3: recoverDocTypesFrom operates on pre-loaded products (batch path). +// --------------------------------------------------------------------------- + +func TestRecoverDocTypesFromEquivalent(t *testing.T) { + products := []kccommon.Product{ + {DocID: "d1", Kind: "structure"}, + {DocID: "d1", Kind: "tree"}, + {DocID: "d1", Kind: "structure"}, // duplicate variant, deduped + } + got, taskTypes, err := recoverDocTypesFrom(products) + if err != nil { + t.Fatalf("recoverDocTypesFrom: %v", err) + } + want := []string{string(kccommon.VariantStructure), string(kccommon.VariantTree)} + if !reflect.DeepEqual(got, want) { + t.Fatalf("variants = %v, want %v", got, want) + } + if !reflect.DeepEqual(taskTypes, []string{kccommon.TaskTypeGraph, kccommon.TaskTypeTree}) { + t.Fatalf("task types = %v, want [Graph Tree]", taskTypes) + } +} + +func TestRecoverDocTypesFromUnknownKindHardFails(t *testing.T) { + products := []kccommon.Product{ + {DocID: "d1", Kind: "structure"}, + {DocID: "d1", Kind: "garbage"}, + } + if _, _, err := recoverDocTypesFrom(products); err == nil { + t.Fatal("recoverDocTypesFrom must hard-fail on an unknown authoritative kind") + } +} + +func TestRecoverDocTypesFromEmpty(t *testing.T) { + got, taskTypes, err := recoverDocTypesFrom(nil) + if err != nil { + t.Fatalf("recoverDocTypesFrom(nil): %v", err) + } + if len(got) != 0 || len(taskTypes) != 0 { + t.Fatalf("empty products must yield empty variants/taskTypes, got %v/%v", got, taskTypes) + } +} + +// --------------------------------------------------------------------------- +// Change 4: diagnostic construction helpers (gated by IsDebugEnabled at call site). +// --------------------------------------------------------------------------- + +func TestComputeDedupVectorAudit(t *testing.T) { + incoming := []kccommon.Product{{Vector: []float32{1}}, {}} + candidates := []kccommon.Product{{Vector: []float32{1}}, {Vector: []float32{1}}, {}} + inVec, candVec := computeDedupVectorAudit(incoming, candidates) + if inVec != 1 || candVec != 2 { + t.Fatalf("computeDedupVectorAudit = (%d,%d), want (1,2)", inVec, candVec) + } +} + +func TestComputeCandidatesBreakdown(t *testing.T) { + candidates := []kccommon.Product{ + {DocID: "kb", Variant: kccommon.VariantWiki, Meta: map[string]any{"slug": "a"}}, + {DocID: "kb", Variant: kccommon.VariantWiki, Meta: map[string]any{"slug": "b"}}, + {DocID: "kb", Variant: kccommon.VariantTree}, + } + byVariant, slugs := computeCandidatesBreakdown(candidates) + if byVariant[string(kccommon.VariantWiki)] != 2 || byVariant[string(kccommon.VariantTree)] != 1 { + t.Fatalf("byVariant = %v, want wiki=2 tree=1", byVariant) + } + if len(slugs) != 2 { + t.Fatalf("wiki slugs = %d, want 2", len(slugs)) + } +} + +func TestComputeDocProductsVectorAudit(t *testing.T) { + products := []kccommon.Product{{Vector: make([]float32, 4)}, {Vector: make([]float32, 4)}, {}} + vecCount, dims := computeDocProductsVectorAudit(products) + if vecCount != 2 { + t.Fatalf("vecCount = %d, want 2", vecCount) + } + if dims[4] != 2 { + t.Fatalf("dims[4] = %d, want 2", dims[4]) + } +} + +func TestComputeKNNGroupDumps(t *testing.T) { + groups := map[string]*mergeGroup{ + "g1": {existing: kccommon.Product{ID: "e1"}, candidates: []kccommon.Product{{DocID: "kb"}, {DocID: "kb"}}, score: 0.9}, + } + dumps := computeKNNGroupDumps(groups) + if len(dumps) != 1 { + t.Fatalf("dumps = %d, want 1", len(dumps)) + } + if dumps[0].Existing != "e1" || len(dumps[0].Candidates) != 2 || dumps[0].Score != 0.9 { + t.Fatalf("dump = %+v, want Existing=e1 Candidates=2 Score=0.9", dumps[0]) + } +} + +// --------------------------------------------------------------------------- +// helpers +// --------------------------------------------------------------------------- + +func itoa(i int) string { + if i == 0 { + return "0" + } + neg := i < 0 + if neg { + i = -i + } + var b [20]byte + pos := len(b) + for i > 0 { + pos-- + b[pos] = byte('0' + i%10) + i /= 10 + } + if neg { + pos-- + b[pos] = '-' + } + return string(b[pos:]) +} + +func keysOf(m map[string]kccommon.Product) []string { + out := make([]string, 0, len(m)) + for k := range m { + out = append(out, k) + } + return out +} + +var _ = types.SearchRequest{} diff --git a/internal/ingestion/knowledge_compile/pr7_exists_filter_test.go b/internal/ingestion/knowledge_compile/pr7_exists_filter_test.go new file mode 100644 index 0000000000..d96619ef3e --- /dev/null +++ b/internal/ingestion/knowledge_compile/pr7_exists_filter_test.go @@ -0,0 +1,113 @@ +// 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. +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// 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. +// See the License for the specific language governing permissions and +// limitations under the License. + +package knowledge_compile + +import ( + "context" + "fmt" + "testing" +) + +// TestLoadDocProductsByDocsExistsFilterSent guards the perf fix for PR #22241: +// the batched reader must bound its query to compiled rows via an `exists: +// compile_kwd` clause. Without it the IN-filter also pulls every source chunk +// (hundreds per document), which both defeats the performance win and overflows +// the engine's result window on a large KB. This test fails (RED) until the +// production reader emits the exists clause. +func TestLoadDocProductsByDocsExistsFilterSent(t *testing.T) { + const nDocs = 500 + const perDoc = 16 // compiled rows only + docIDs := make([]string, 0, nDocs) + chunks := make([]map[string]interface{}, 0, nDocs*perDoc) + for i := 0; i < nDocs; i++ { + id := fmt.Sprintf("d%d", i) + docIDs = append(docIDs, id) + for j := 0; j < perDoc; j++ { + chunks = append(chunks, map[string]interface{}{ + "id": fmt.Sprintf("p_%d_%d", i, j), + "doc_id": id, + "compile_kwd": "wiki_page", + "available_int": 1, + "content_with_weight": "c", + }) + } + } + eng := &fakeEngine{searchChunks: chunks, windowed: true} + r := engineReader{eng: eng} + if _, err := r.LoadDocProductsByDocs(context.Background(), "tenant", "kb", docIDs); err != nil { + t.Fatalf("LoadDocProductsByDocs: %v", err) + } + if got, ok := eng.lastSearchReq.Filter["exists"]; !ok || got != "compile_kwd" { + t.Fatalf("LoadDocProductsByDocs must send exists:compile_kwd; got Filter=%#v", eng.lastSearchReq.Filter) + } +} + +// TestLoadDocProductsByDocsBoundedByExistsAvoidsWindowOverflow is the adversarial +// guard for the same fix: a 500-document batch mixes compiled rows (which carry +// compile_kwd) with ordinary source chunks (which do not). Without the exists +// clause the engine must scan ~20000 rows, so the third page's from+size exceeds +// the result window and the read errors. With the exists clause the query is +// bounded to the ~8000 compiled rows, stays within a single window, and returns +// every compiled product exactly once. +func TestLoadDocProductsByDocsBoundedByExistsAvoidsWindowOverflow(t *testing.T) { + const nDocs = 500 + const compiledPerDoc = 16 + const sourcePerDoc = 24 // 500 * 40 = 20000 rows unfiltered (> MAX_RESULT_WINDOW) + docIDs := make([]string, 0, nDocs) + chunks := make([]map[string]interface{}, 0, nDocs*(compiledPerDoc+sourcePerDoc)) + for i := 0; i < nDocs; i++ { + id := fmt.Sprintf("d%d", i) + docIDs = append(docIDs, id) + for j := 0; j < compiledPerDoc; j++ { + chunks = append(chunks, map[string]interface{}{ + "id": fmt.Sprintf("c_%d_%d", i, j), + "doc_id": id, + "compile_kwd": "wiki_page", + "available_int": 1, + "content_with_weight": "c", + }) + } + for j := 0; j < sourcePerDoc; j++ { + chunks = append(chunks, map[string]interface{}{ + "id": fmt.Sprintf("s_%d_%d", i, j), + "doc_id": id, + "available_int": 1, + "content_with_weight": "source", + }) + } + } + eng := &fakeEngine{searchChunks: chunks, windowed: true} + r := engineReader{eng: eng} + got, err := r.LoadDocProductsByDocs(context.Background(), "tenant", "kb", docIDs) + if err != nil { + t.Fatalf("LoadDocProductsByDocs must stay within the result window via exists:compile_kwd; got err=%v", err) + } + // Every compiled row must be returned exactly once; source chunks excluded. + count := 0 + seen := map[string]int{} + for _, ps := range got { + for _, p := range ps { + count++ + seen[p.ID]++ + } + } + if count != nDocs*compiledPerDoc { + t.Fatalf("recovered products = %d, want %d (silent truncation or source-chunk leak)", count, nDocs*compiledPerDoc) + } + for id, n := range seen { + if n != 1 { + t.Fatalf("product %s recovered %d times, want 1 (duplicate across pages)", id, n) + } + } +} diff --git a/internal/ingestion/knowledge_compile/pr7_f1_window_bounds_test.go b/internal/ingestion/knowledge_compile/pr7_f1_window_bounds_test.go new file mode 100644 index 0000000000..5d2dadaf0e --- /dev/null +++ b/internal/ingestion/knowledge_compile/pr7_f1_window_bounds_test.go @@ -0,0 +1,201 @@ +// 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. +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// 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. +// See the License for the specific language governing permissions and +// limitations under the License. + +package knowledge_compile + +import ( + "context" + "fmt" + "testing" +) + +// buildCompiledOnlyChunks returns nDocs documents each with compiledPerDoc +// compiled rows (compile_kwd set, no source chunks). This isolates the F1 +// failure mode: the doc_id IN batch aggregates more compiled rows than the +// engine result window, independent of the source-chunk inflation that the +// exists:compile_kwd filter already fixes. +func buildCompiledOnlyChunks(nDocs, compiledPerDoc int) ([]string, []map[string]interface{}) { + docIDs := make([]string, 0, nDocs) + chunks := make([]map[string]interface{}, 0, nDocs*compiledPerDoc) + for i := 0; i < nDocs; i++ { + id := fmt.Sprintf("d%d", i) + docIDs = append(docIDs, id) + for j := 0; j < compiledPerDoc; j++ { + chunks = append(chunks, map[string]interface{}{ + "id": fmt.Sprintf("c_%d_%d", i, j), + "doc_id": id, + "compile_kwd": "wiki_page", + "available_int": 1, + "content_with_weight": "c", + }) + } + } + return docIDs, chunks +} + +// TestLoadDocProductsByDocsCompiledRowsExceedWindow is the adversarial guard for +// F1: a 500-document batch whose COMPILED rows alone (no source chunks) exceed +// MAX_RESULT_WINDOW (10000) must not hard-fail. The exists:compile_kwd filter +// already excludes source chunks, but a product-dense batch can still aggregate +// more than 10k compiled rows; plain from/size cannot page past the window, so +// the reader must fall back to per-document reads for that sub-batch. +func TestLoadDocProductsByDocsCompiledRowsExceedWindow(t *testing.T) { + const nDocs = 500 + const compiledPerDoc = 30 // 500*30 = 15000 compiled rows > MAX_RESULT_WINDOW(10000) + docIDs, chunks := buildCompiledOnlyChunks(nDocs, compiledPerDoc) + eng := &fakeEngine{searchChunks: chunks, windowed: true} + r := engineReader{eng: eng} + got, err := r.LoadDocProductsByDocs(context.Background(), "tenant", "kb", docIDs) + if err != nil { + t.Fatalf("compiled-only batch exceeding window must fall back to per-doc reads, not error: %v", err) + } + // Fallback must have engaged (batched path alone would error or need >2 calls). + if eng.searchCalls <= 2 { + t.Fatalf("expected fallback to per-doc reads (more than 2 Search calls), got %d", eng.searchCalls) + } + // Every compiled product returned exactly once, grouped by doc. + count := 0 + seen := map[string]int{} + for _, ps := range got { + for _, p := range ps { + count++ + seen[p.ID]++ + } + } + if count != nDocs*compiledPerDoc { + t.Fatalf("recovered products = %d, want %d (silent truncation past window)", count, nDocs*compiledPerDoc) + } + for id, n := range seen { + if n != 1 { + t.Fatalf("product %s recovered %d times, want 1 (duplicate across pages/fallback)", id, n) + } + } +} + +// TestLoadDocProductsByDocsExactlyAtWindowFallsBack is the off-by-one guard: +// a batch whose compiled rows land exactly on MAX_RESULT_WINDOW (10000) still +// triggers a 3rd page request (offset 10000) that plain from/size rejects, so +// the fallback must engage here too. +func TestLoadDocProductsByDocsExactlyAtWindowFallsBack(t *testing.T) { + const nDocs = 500 + const compiledPerDoc = 20 // 500*20 = 10000 exactly + docIDs, chunks := buildCompiledOnlyChunks(nDocs, compiledPerDoc) + eng := &fakeEngine{searchChunks: chunks, windowed: true} + r := engineReader{eng: eng} + got, err := r.LoadDocProductsByDocs(context.Background(), "tenant", "kb", docIDs) + if err != nil { + t.Fatalf("exactly-at-window batch must fall back to per-doc reads, not error: %v", err) + } + if eng.searchCalls <= 2 { + t.Fatalf("expected fallback to per-doc reads (more than 2 Search calls), got %d", eng.searchCalls) + } + count := 0 + seen := map[string]int{} + for _, ps := range got { + for _, p := range ps { + count++ + seen[p.ID]++ + } + } + if count != nDocs*compiledPerDoc { + t.Fatalf("recovered products = %d, want %d", count, nDocs*compiledPerDoc) + } + for id, n := range seen { + if n != 1 { + t.Fatalf("product %s recovered %d times, want 1", id, n) + } + } +} + +// TestLoadDocProductsByDocsUnderWindowStaysBatched guards against regressing the +// common path: a batch comfortably under the window must keep using the batched +// IN read (2 Search calls for 8000 rows paged at 5000), never silently falling +// back to one query per document. +func TestLoadDocProductsByDocsUnderWindowStaysBatched(t *testing.T) { + const nDocs = 500 + const compiledPerDoc = 16 // 500*16 = 8000 < MAX_RESULT_WINDOW(10000) + docIDs, chunks := buildCompiledOnlyChunks(nDocs, compiledPerDoc) + eng := &fakeEngine{searchChunks: chunks, windowed: true} + r := engineReader{eng: eng} + got, err := r.LoadDocProductsByDocs(context.Background(), "tenant", "kb", docIDs) + if err != nil { + t.Fatalf("under-window batch must not error: %v", err) + } + if eng.searchCalls != 2 { + t.Fatalf("expected batched path (2 Search calls), got %d (regressed to per-doc?)", eng.searchCalls) + } + count := 0 + for _, ps := range got { + count += len(ps) + } + if count != nDocs*compiledPerDoc { + t.Fatalf("recovered products = %d, want %d", count, nDocs*compiledPerDoc) + } +} + +// TestLoadDocProductsByDocsMixedOverflowSubBatches guards the per-sub-batch +// scoping of the fallback: across a single call, the first 500-doc sub-batch is +// product-dense (overflows -> falls back to per-doc) while the second 500-doc +// sub-batch is comfortably under the window (stays batched). The overflow of one +// sub-batch must not disable batching for the next, nor corrupt its results. +func TestLoadDocProductsByDocsMixedOverflowSubBatches(t *testing.T) { + const densePerDoc = 30 // first 500 docs -> 15000 compiled rows > window + const sparsePerDoc = 16 // next 500 docs -> 8000 compiled rows < window + const nDocs = 1000 + docIDs := make([]string, 0, nDocs) + chunks := make([]map[string]interface{}, 0, nDocs*densePerDoc) + for i := 0; i < nDocs; i++ { + id := fmt.Sprintf("d%d", i) + docIDs = append(docIDs, id) + perDoc := sparsePerDoc + if i < 500 { + perDoc = densePerDoc + } + for j := 0; j < perDoc; j++ { + chunks = append(chunks, map[string]interface{}{ + "id": fmt.Sprintf("c_%d_%d", i, j), + "doc_id": id, + "compile_kwd": "wiki_page", + "available_int": 1, + "content_with_weight": "c", + }) + } + } + eng := &fakeEngine{searchChunks: chunks, windowed: true} + r := engineReader{eng: eng} + got, err := r.LoadDocProductsByDocs(context.Background(), "tenant", "kb", docIDs) + if err != nil { + t.Fatalf("mixed sub-batches must not error: %v", err) + } + // 502 calls for the dense (overflow) sub-batch + 2 for the sparse sub-batch. + if eng.searchCalls != 504 { + t.Fatalf("expected 504 Search calls (502 fallback + 2 batched), got %d", eng.searchCalls) + } + count := 0 + seen := map[string]int{} + for _, ps := range got { + for _, p := range ps { + count++ + seen[p.ID]++ + } + } + want := 500*densePerDoc + 500*sparsePerDoc + if count != want { + t.Fatalf("recovered products = %d, want %d", count, want) + } + for id, n := range seen { + if n != 1 { + t.Fatalf("product %s recovered %d times, want 1", id, n) + } + } +} diff --git a/internal/ingestion/knowledge_compile/pr7_window_fix_test.go b/internal/ingestion/knowledge_compile/pr7_window_fix_test.go new file mode 100644 index 0000000000..2aef88050f --- /dev/null +++ b/internal/ingestion/knowledge_compile/pr7_window_fix_test.go @@ -0,0 +1,164 @@ +// 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. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// 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. +// See the License for the specific language governing permissions and +// limitations under the License. + +package knowledge_compile + +import ( + "context" + "fmt" + "reflect" + "testing" + + kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common" +) + +// TestLoadDocProductsByDocsSendsNoSearchAfterOrderBy guards the fix for PR +// #22241 review findings: the batched reader must NOT rely on OrderBy: Asc("id") +// to drive search_after deep pagination. On OpenSearch parseOrderByExpr drops the +// "id" field (no sort emitted), so useSearchAfter opens with a nil cursor and the +// read silently re-reads the first page (truncation + duplicates); on Elasticsearch +// 8.x the id-> _id translation is rejected by default (indices.id_field_data off), +// so the whole read fails. A 500-document batch never approaches the result window +// in practice, so plain from/size paging is correct and safe. The reader must +// therefore send NO OrderBy. +func TestLoadDocProductsByDocsSendsNoSearchAfterOrderBy(t *testing.T) { + eng := &fakeEngine{searchChunks: []map[string]interface{}{ + {"id": "p1", "doc_id": "d1", "compile_kwd": "wiki_page", "available_int": 1, "content_with_weight": "c1"}, + }} + r := engineReader{eng: eng} + if _, err := r.LoadDocProductsByDocs(context.Background(), "tenant", "kb", []string{"d1"}); err != nil { + t.Fatalf("LoadDocProductsByDocs: %v", err) + } + if eng.lastSearchReq == nil { + t.Fatalf("expected a SearchRequest to be captured") + } + if eng.lastSearchReq.OrderBy != nil && len(eng.lastSearchReq.OrderBy.Fields) > 0 { + t.Fatalf("LoadDocProductsByDocs must not send an OrderBy (would drive broken search_after); got %#v", eng.lastSearchReq.OrderBy) + } +} + +// TestLoadDocProductsByDocsReadsAllRowsWithinWindow is the happy-path guard: a +// batch whose aggregated rows stay under the backend max_result_window must return +// every product exactly once, paged with plain from/size (no search_after). +func TestLoadDocProductsByDocsReadsAllRowsWithinWindow(t *testing.T) { + const nDocs = 500 + const perDoc = 16 // 500 * 16 = 8000 rows (< MAX_RESULT_WINDOW 10000) + docIDs := make([]string, 0, nDocs) + chunks := make([]map[string]interface{}, 0, nDocs*perDoc) + for i := 0; i < nDocs; i++ { + id := fmt.Sprintf("d%d", i) + docIDs = append(docIDs, id) + for j := 0; j < perDoc; j++ { + chunks = append(chunks, map[string]interface{}{ + "id": fmt.Sprintf("p_%d_%d", i, j), + "doc_id": id, + "compile_kwd": "wiki_page", + "available_int": 1, + "content_with_weight": "c", + }) + } + } + eng := &fakeEngine{searchChunks: chunks, windowed: true} + r := engineReader{eng: eng} + got, err := r.LoadDocProductsByDocs(context.Background(), "tenant", "kb", docIDs) + if err != nil { + t.Fatalf("LoadDocProductsByDocs: %v", err) + } + count := 0 + for _, ps := range got { + count += len(ps) + } + if count != nDocs*perDoc { + t.Fatalf("recovered products = %d, want %d (silent truncation past the result window)", count, nDocs*perDoc) + } + // no duplicates: each (doc, product id) must appear once. + seen := map[string]int{} + for _, ps := range got { + for _, p := range ps { + seen[p.ID]++ + } + } + for id, n := range seen { + if n != 1 { + t.Fatalf("product %s recovered %d times, want 1 (duplicate across pages)", id, n) + } + } +} + +// TestRecoverDocVariantsBatchChunksBounded guards review finding #1: +// RebuildDataset must not hold the entire KB's compiled products (including +// vectors) in memory at once. recoverDocVariantsBatch walks docs in +// docProductsBatchSize chunks, so each LoadDocProductsByDocs call stays bounded +// and the in-flight product set is released after each chunk. +func TestRecoverDocVariantsBatchChunksBounded(t *testing.T) { + const nDocs = 1200 + docs := make([]string, 0, nDocs) + products := make([]kccommon.Product, 0, nDocs) + for i := 0; i < nDocs; i++ { + id := fmt.Sprintf("d%d", i) + docs = append(docs, id) + products = append(products, kccommon.Product{DocID: id, Kind: "structure"}) + } + fr := &fakeReader{products: products} + c := &Consumer{reader: fr} + got, err := c.recoverDocVariantsBatch(context.Background(), "t", "kb", docs) + if err != nil { + t.Fatalf("recoverDocVariantsBatch: %v", err) + } + if len(got) != nDocs { + t.Fatalf("recovered docs = %d, want %d", len(got), nDocs) + } + for _, id := range docs { + vr, ok := got[id] + if !ok { + t.Fatalf("missing recovery for %s", id) + } + if !reflect.DeepEqual(vr.variants, []string{string(kccommon.VariantStructure)}) { + t.Fatalf("%s variants = %v, want [structure]", id, vr.variants) + } + } + // Memory-bounding: every batch read must stay within docProductsBatchSize. + for i, sz := range fr.batchSizes { + if sz > docProductsBatchSize { + t.Fatalf("batch %d size = %d, exceeds docProductsBatchSize %d (unbounded read)", i, sz, docProductsBatchSize) + } + } + // 1200 docs chunked at 500 => exactly 3 calls (500 + 500 + 200). + if len(fr.batchSizes) != 3 { + t.Fatalf("batch calls = %d, want 3 (chunked into docProductsBatchSize)", len(fr.batchSizes)) + } + if fr.batchSizes[0] != docProductsBatchSize || fr.batchSizes[1] != docProductsBatchSize || fr.batchSizes[2] != 200 { + t.Fatalf("batch sizes = %v, want [500 500 200]", fr.batchSizes) + } +} + +// TestRecoverDocVariantsBatchDocWithoutProducts ensures a document with no +// compiled products yields empty variants (the caller skips it) rather than an +// error. +func TestRecoverDocVariantsBatchDocWithoutProducts(t *testing.T) { + fr := &fakeReader{products: []kccommon.Product{ + {DocID: "d1", Kind: "structure"}, + }} + c := &Consumer{reader: fr} + got, err := c.recoverDocVariantsBatch(context.Background(), "t", "kb", []string{"d1", "d2"}) + if err != nil { + t.Fatalf("recoverDocVariantsBatch: %v", err) + } + if vr, ok := got["d2"]; !ok || len(vr.variants) != 0 { + t.Fatalf("d2 must recover empty variants, got %#v", vr) + } + if !reflect.DeepEqual(got["d1"].variants, []string{string(kccommon.VariantStructure)}) { + t.Fatalf("d1 variants = %v, want [structure]", got["d1"].variants) + } +} diff --git a/internal/ingestion/knowledge_compile/reader.go b/internal/ingestion/knowledge_compile/reader.go index 940190e1e0..90ee1752a9 100644 --- a/internal/ingestion/knowledge_compile/reader.go +++ b/internal/ingestion/knowledge_compile/reader.go @@ -18,6 +18,7 @@ package knowledge_compile import ( "context" "encoding/json" + "errors" "fmt" "strconv" "strings" @@ -47,6 +48,10 @@ type Reader interface { // document (doc_id == source_doc). Bounded by one document, never the whole // KB. LoadDocProducts(ctx context.Context, tenant, kb, docID string) ([]kccommon.Product, error) + // LoadDocProductsByDocs returns the per-document compiled rows for many + // documents in one query, grouped by doc_id. It is the batch form of + // LoadDocProducts used to avoid one query per document. + LoadDocProductsByDocs(ctx context.Context, tenant, kb string, docIDs []string) (map[string][]kccommon.Product, error) // SearchSimilar runs a dense (KNN) search over the existing merged rows of // the given variant and returns the single most-similar row whose score is @@ -152,6 +157,30 @@ var wikiSelectFields = []string{ // LoadDocProducts pages until the engine returns fewer than a full page. const loadDocProductsLimit = 5000 +// pageThrough issues req against eng, paging with plain from/size, and invokes +// visit for every page's chunks. It stops after the final (short) page so a +// result set larger than one page is read completely without silent +// truncation or looping. req.Offset is advanced on each iteration; callers must +// not set it beforehand. req.Limit is left to the caller (typically +// loadDocProductsLimit). +func pageThrough(ctx context.Context, eng engine.DocEngine, req *types.SearchRequest, limit int, visit func(chunks []map[string]interface{}) error) error { + for offset := 0; ; offset += limit { + req.Offset = offset + res, err := eng.Search(ctx, req) + if err != nil { + return err + } + if res != nil { + if err := visit(res.Chunks); err != nil { + return err + } + } + if len(res.Chunks) < limit { + return nil + } + } +} + // LoadDocProducts returns the per-document compiled rows for a single document. // It is bounded to one document, so the consumer never loads the whole KB. The // results are paged so a document with more than loadDocProductsLimit rows is @@ -162,20 +191,15 @@ func (r engineReader) LoadDocProducts(ctx context.Context, tenant, kb, docID str } eng := r.eng var out []kccommon.Product - offset := 0 - for { - res, err := eng.Search(ctx, &types.SearchRequest{ - IndexNames: []string{fmt.Sprintf("ragflow_%s", tenant)}, - KbIDs: []string{kb}, - Filter: map[string]interface{}{"doc_id": docID}, - SelectFields: append(append([]string(nil), compiledSelectFields...), wikiSelectFields...), - Limit: loadDocProductsLimit, - Offset: offset, - }) - if err != nil { - return nil, err - } - for _, c := range res.Chunks { + req := &types.SearchRequest{ + IndexNames: []string{fmt.Sprintf("ragflow_%s", tenant)}, + KbIDs: []string{kb}, + Filter: map[string]interface{}{"doc_id": docID}, + SelectFields: append(append([]string(nil), compiledSelectFields...), wikiSelectFields...), + Limit: loadDocProductsLimit, + } + if err := pageThrough(ctx, eng, req, loadDocProductsLimit, func(chunks []map[string]interface{}) error { + for _, c := range chunks { // Only compiled products carry compile_kwd; skip ordinary source chunks. if _, ok := c["compile_kwd"]; !ok || types.IsNavigationRow(c) { continue @@ -190,32 +214,155 @@ func (r engineReader) LoadDocProducts(ctx context.Context, tenant, kb, docID str out = append(out, p) } } - if len(res.Chunks) < loadDocProductsLimit { - break - } - offset += loadDocProductsLimit + return nil + }); err != nil { + return nil, err } // Diagnostics: count how many loaded per-doc products carry an embedding // and report the (deduped) vector dimensions. If this shows 0/empty while // per-doc rows in ES do have q__vec, VectorFromChunkMap is failing to - // restore the vector and the merged rows will end up embedding-less. - vecCount := 0 - dims := map[int]int{} - for _, p := range out { - if dim := len(p.Vector); dim > 0 { - vecCount++ - dims[dim]++ + // restore the vector and the merged rows will end up embedding-less. Gated to + // debug: the map construction below allocates even in quiet production. + if common.IsDebugEnabled() { + vecCount, dims := computeDocProductsVectorAudit(out) + common.Info("knowledge_compile: LoadDocProducts vector audit", + zap.String("kb_id", kb), + zap.String("doc_id", docID), + zap.Int("products", len(out)), + zap.Int("with_vector", vecCount), + zap.Any("vector_dims", dims)) + } + return out, nil +} + +// docProductsBatchSize bounds how many doc_ids are sent in a single IN filter so a +// batch read over a large KB does not build an unbounded query. +const docProductsBatchSize = 500 + +// LoadDocProductsByDocs returns the per-document compiled products for many +// documents in one query (doc_id IN [...]) instead of one query per document. +// Results are grouped by doc_id; documents with no compiled rows map to an empty +// (non-nil) slice so callers can distinguish "no products" from "missing". +// +// Each doc_id IN sub-batch is read with a single bounded query. Because plain +// from/size cannot page past the engine result window (MAX_RESULT_WINDOW), a +// sub-batch whose aggregated compiled rows would force a page beyond the window +// falls back to one bounded per-document read per document (the pre-PR behavior) +// instead of erroring — a product-dense large KB must not lose or mis-read +// merged state during a rebuild. +func (r engineReader) LoadDocProductsByDocs(ctx context.Context, tenant, kb string, docIDs []string) (map[string][]kccommon.Product, error) { + if r.eng == nil { + return nil, errEngineNotConfigured + } + if len(docIDs) == 0 { + return map[string][]kccommon.Product{}, nil + } + out := make(map[string][]kccommon.Product, len(docIDs)) + for _, docID := range docIDs { + out[docID] = nil // ensure every requested doc has an entry + } + fields := append(append([]string(nil), compiledSelectFields...), wikiSelectFields...) + for start := 0; start < len(docIDs); start += docProductsBatchSize { + end := start + docProductsBatchSize + if end > len(docIDs) { + end = len(docIDs) + } + batch := docIDs[start:end] + local, overflow, err := r.loadDocProductsBatch(ctx, tenant, kb, batch, fields) + if err != nil { + return nil, err + } + if overflow { + // The sub-batch's compiled rows exceed the engine result window, so a + // single doc_id IN read cannot page past it without erroring. Fall + // back to one bounded per-document read per document: each document + // is bounded to itself and never aggregates to the window, so the + // read stays correct for product-dense batches. + for _, docID := range batch { + ps, perr := r.LoadDocProducts(ctx, tenant, kb, docID) + if perr != nil { + return nil, perr + } + out[docID] = ps + } + continue + } + for docID, ps := range local { + out[docID] = ps } } - common.Info("knowledge_compile: LoadDocProducts vector audit", - zap.String("kb_id", kb), - zap.String("doc_id", docID), - zap.Int("products", len(out)), - zap.Int("with_vector", vecCount), - zap.Any("vector_dims", dims)) return out, nil } +// errBatchOverflow is returned by loadDocProductsBatch's paging visitor to signal +// that the next page would exceed the engine result window; the caller falls +// back to per-document reads instead of letting the engine reject the query. +var errBatchOverflow = errors.New("knowledge_compile: batch read reached result window") + +// loadDocProductsBatch reads the compiled products for a slice of doc_ids with a +// single doc_id IN query, paging with plain from/size. It returns the products +// grouped by doc_id. If the aggregate result set is large enough that the next +// page would require offset+limit > MAX_RESULT_WINDOW (which the engine rejects +// for plain from/size), it stops paging early and reports overflow=true so the +// caller can fall back to per-document reads rather than erroring. +func (r engineReader) loadDocProductsBatch(ctx context.Context, tenant, kb string, batch []string, fields []string) (map[string][]kccommon.Product, bool, error) { + local := make(map[string][]kccommon.Product, len(batch)) + req := &types.SearchRequest{ + IndexNames: []string{fmt.Sprintf("ragflow_%s", tenant)}, + KbIDs: []string{kb}, + Filter: map[string]interface{}{"doc_id": docIDInterfaceSlice(batch), "exists": "compile_kwd"}, + SelectFields: fields, + Limit: loadDocProductsLimit, + } + overflow := false + err := pageThrough(ctx, r.eng, req, loadDocProductsLimit, func(chunks []map[string]interface{}) error { + // The next page would land at offset+limit, and that fetch needs + // (offset+limit)+limit <= MAX_RESULT_WINDOW. If it would not, stop now + // and let the caller fall back to per-document reads for this batch. + // A next page (offset+limit) is only fetched when this page is full, and + // only then could a deep page exceed the engine result window. Signal + // overflow only in that case, so a short final page never triggers the + // per-document fallback. + if len(chunks) == loadDocProductsLimit && req.Offset+2*loadDocProductsLimit > common.MAX_RESULT_WINDOW { + overflow = true + return errBatchOverflow + } + for _, c := range chunks { + // Only compiled products carry compile_kwd; skip ordinary source chunks. + if _, ok := c["compile_kwd"]; !ok || types.IsNavigationRow(c) { + continue + } + rowVariant, verr := KwdToVariant(asString(c["compile_kwd"])) + if verr != nil { + continue + } + if p, ok := productFromChunkMap(c, tenant, rowVariant); ok { + if docID, _ := c["doc_id"].(string); docID != "" { + local[docID] = append(local[docID], p) + } + } + } + return nil + }) + if err != nil { + if errors.Is(err, errBatchOverflow) { + return local, true, nil + } + return nil, false, err + } + return local, overflow, nil +} + +// docIDInterfaceSlice converts a string doc_id list into the []interface{} form +// the engine's IN filter expects. +func docIDInterfaceSlice(docIDs []string) []interface{} { + out := make([]interface{}, len(docIDs)) + for i, d := range docIDs { + out[i] = d + } + return out +} + func (r engineReader) LoadMergedProduct(ctx context.Context, tenant, kb, id string) (kccommon.Product, error) { if r.eng == nil { return kccommon.Product{}, errEngineNotConfigured @@ -248,30 +395,22 @@ func (r engineReader) LoadMergedWikiPages(ctx context.Context, tenant, kb string return nil, errEngineNotConfigured } eng := r.eng - filter := map[string]interface{}{"available_int": 1, "scope_kwd": "dataset", "type_kwd": compileKwdWikiPage} var pages []kccommon.Product - for offset := 0; ; offset += loadDocProductsLimit { - res, err := eng.Search(ctx, &types.SearchRequest{ - IndexNames: []string{fmt.Sprintf("ragflow_%s", tenant)}, KbIDs: []string{kb}, - Limit: loadDocProductsLimit, Offset: offset, - OrderBy: (&types.OrderByExpr{}).Asc("id"), - SelectFields: append(append([]string(nil), compiledSelectFields...), wikiSelectFields...), - Filter: filter, - }) - if err != nil { - return nil, err - } - if res == nil || len(res.Chunks) == 0 { - break - } - for _, row := range res.Chunks { + req := &types.SearchRequest{ + IndexNames: []string{fmt.Sprintf("ragflow_%s", tenant)}, KbIDs: []string{kb}, + Limit: loadDocProductsLimit, + SelectFields: append(append([]string(nil), compiledSelectFields...), wikiSelectFields...), + Filter: map[string]interface{}{"available_int": 1, "scope_kwd": "dataset", "type_kwd": compileKwdWikiPage}, + } + if err := pageThrough(ctx, eng, req, loadDocProductsLimit, func(chunks []map[string]interface{}) error { + for _, row := range chunks { if product, ok := productFromChunkMap(row, tenant, kccommon.VariantWiki); ok && kccommon.MetaString(product.Meta, "kind") == "page" { pages = append(pages, product) } } - if len(res.Chunks) < loadDocProductsLimit { - break - } + return nil + }); err != nil { + return nil, err } return pages, nil } @@ -300,28 +439,21 @@ func (r engineReader) LoadDocumentWikiPagesBySlugs(ctx context.Context, tenant, filter["doc_id"] = enabledIDs } var pages []kccommon.Product - for offset := 0; ; offset += loadDocProductsLimit { - res, err := eng.Search(ctx, &types.SearchRequest{ - IndexNames: []string{fmt.Sprintf("ragflow_%s", tenant)}, KbIDs: []string{kb}, - Limit: loadDocProductsLimit, Offset: offset, - OrderBy: (&types.OrderByExpr{}).Asc("id"), - SelectFields: append(append([]string(nil), compiledSelectFields...), wikiSelectFields...), - Filter: filter, - }) - if err != nil { - return nil, err - } - if res == nil || len(res.Chunks) == 0 { - break - } - for _, row := range res.Chunks { + req := &types.SearchRequest{ + IndexNames: []string{fmt.Sprintf("ragflow_%s", tenant)}, KbIDs: []string{kb}, + Limit: loadDocProductsLimit, + SelectFields: append(append([]string(nil), compiledSelectFields...), wikiSelectFields...), + Filter: filter, + } + if err := pageThrough(ctx, eng, req, loadDocProductsLimit, func(chunks []map[string]interface{}) error { + for _, row := range chunks { if product, ok := productFromChunkMap(row, tenant, kccommon.VariantWiki); ok && kccommon.MetaString(product.Meta, "kind") == "page" { pages = append(pages, product) } } - if len(res.Chunks) < loadDocProductsLimit { - break - } + return nil + }); err != nil { + return nil, err } return pages, nil } diff --git a/internal/ingestion/knowledge_compile/rebuild_variants_test.go b/internal/ingestion/knowledge_compile/rebuild_variants_test.go index 639719e055..26d656cdc6 100644 --- a/internal/ingestion/knowledge_compile/rebuild_variants_test.go +++ b/internal/ingestion/knowledge_compile/rebuild_variants_test.go @@ -8,70 +8,50 @@ import ( kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common" ) -// fakeReader is a Reader double that returns a canned product set for -// recoverDocTypes tests. +// fakeReader is a Reader double that returns a canned product set for the +// recoverDocVariantsBatch tests. type fakeReader struct { products []kccommon.Product + // batchSizes records the doc_id slice length passed to each + // LoadDocProductsByDocs call, so chunking tests can assert bounded batches. + batchSizes []int } func (f *fakeReader) LoadDocProducts(context.Context, string, string, string) ([]kccommon.Product, error) { return f.products, nil } +func (f *fakeReader) LoadDocProductsByDocs(_ context.Context, _, _ string, docIDs []string) (map[string][]kccommon.Product, error) { + f.batchSizes = append(f.batchSizes, len(docIDs)) + out := make(map[string][]kccommon.Product, len(docIDs)) + for _, docID := range docIDs { + out[docID] = nil + } + for _, p := range f.products { + out[p.DocID] = append(out[p.DocID], p) + } + return out, nil +} + func (f *fakeReader) SearchSimilar(context.Context, string, string, kccommon.Variant, []float64, int, float64) (kccommon.Product, float64, error) { return kccommon.Product{}, 0, nil } -// TestRecoverDocTypes_AuthoritativeKind covers B1a/O2a: the authoritative -// Product.Kind (compilation_template_kind_kwd) is mapped through KindToVariant; -// results are sorted/deduped. -func TestRecoverDocVariants_AuthoritativeKind(t *testing.T) { - c := &Consumer{reader: &fakeReader{products: []kccommon.Product{ - {DocID: "d1", Kind: "structure"}, - {DocID: "d1", Kind: "tree"}, - // duplicate variant, deduped - {DocID: "d1", Kind: "structure"}, - }}} - got, taskTypes, err := c.recoverDocTypes(context.Background(), "t1", "kb1", "d1") - if err != nil { - t.Fatalf("recoverDocTypes error: %v", err) - } - want := []string{string(kccommon.VariantStructure), string(kccommon.VariantTree)} - if !reflect.DeepEqual(got, want) { - t.Fatalf("recoverDocTypes variants = %v, want %v", got, want) - } - if !reflect.DeepEqual(taskTypes, []string{kccommon.TaskTypeGraph, kccommon.TaskTypeTree}) { - t.Fatalf("recoverDocTypes task types = %v, want [Graph Tree]", taskTypes) - } -} - -// TestRecoverDocTypes_UnknownKindHardFails covers O2a: a whitelist-out -// authoritative kind aborts recovery (returns an error) so the rebuild does not -// proceed with an incomplete variant set. -func TestRecoverDocVariants_UnknownKindHardFails(t *testing.T) { - c := &Consumer{reader: &fakeReader{products: []kccommon.Product{ - {DocID: "d1", Kind: "structure"}, - {DocID: "d1", Kind: "garbage"}, - }}} - if _, _, err := c.recoverDocTypes(context.Background(), "t1", "kb1", "d1"); err == nil { - t.Fatal("recoverDocTypes must hard-fail on an unknown authoritative kind (O2a)") - } -} - -// TestRecoverDocTypes_FallbackVariant covers B1a: a product without an -// authoritative kind falls back to its reverse-mapped variant. -func TestRecoverDocVariants_FallbackVariant(t *testing.T) { - c := &Consumer{reader: &fakeReader{products: []kccommon.Product{ +// TestRecoverDocTypesFromFallbackVariant covers B1a: a product without an +// authoritative kind falls back to its reverse-mapped variant. This is the +// branch the deleted recoverDocTypes wrapper used to exercise; it now lives on +// the shared recoverDocTypesFrom helper used by the rebuild batch path. +func TestRecoverDocTypesFromFallbackVariant(t *testing.T) { + got, taskTypes, err := recoverDocTypesFrom([]kccommon.Product{ {DocID: "d1", Variant: kccommon.VariantWiki}, - }}} - got, taskTypes, err := c.recoverDocTypes(context.Background(), "t1", "kb1", "d1") + }) if err != nil { - t.Fatalf("recoverDocTypes error: %v", err) + t.Fatalf("recoverDocTypesFrom error: %v", err) } if len(got) != 1 || got[0] != string(kccommon.VariantWiki) { - t.Fatalf("recoverDocTypes fallback = %v, want [wiki]", got) + t.Fatalf("recoverDocTypesFrom fallback = %v, want [wiki]", got) } if !reflect.DeepEqual(taskTypes, []string{kccommon.TaskTypeWiki}) { - t.Fatalf("recoverDocTypes fallback task types = %v, want [Wiki]", taskTypes) + t.Fatalf("recoverDocTypesFrom fallback task types = %v, want [Wiki]", taskTypes) } } diff --git a/internal/ingestion/knowledge_compile/wiki_contribution.go b/internal/ingestion/knowledge_compile/wiki_contribution.go index 2f29a53e49..6e1fc78d24 100644 --- a/internal/ingestion/knowledge_compile/wiki_contribution.go +++ b/internal/ingestion/knowledge_compile/wiki_contribution.go @@ -305,18 +305,28 @@ func (c *Consumer) prepareWikiContributionDiff(ctx context.Context, tenantID, da affectedSlugs: map[string]struct{}{}, } completedByDoc := make(map[string]BacklogEntry, len(completed)) + completedDocIDs := make([]string, 0, len(completed)) + seenDoc := make(map[string]struct{}) for _, entry := range completed { - completedByDoc[entry.DocID] = entry - products, err := c.reader.LoadDocProducts(ctx, tenantID, datasetID, entry.DocID) - if err != nil { - return wikiContributionDiffResult{}, err + completedByDoc[entry.DocID] = entry // last entry per doc wins, as before + if _, ok := seenDoc[entry.DocID]; !ok { + seenDoc[entry.DocID] = struct{}{} + completedDocIDs = append(completedDocIDs, entry.DocID) } - result.productsByDoc[entry.DocID] = products + } + // One query for every completed document instead of one query per document. + productsByDocAll, err := c.reader.LoadDocProductsByDocs(ctx, tenantID, datasetID, completedDocIDs) + if err != nil { + return wikiContributionDiffResult{}, err + } + for docID, entry := range completedByDoc { + products := productsByDocAll[docID] + result.productsByDoc[docID] = products wikiProducts := products if len(entry.Variants) > 0 && !variantsContain(entry.Variants, "wiki") { wikiProducts = nil } - result.currentByDoc[entry.DocID] = buildWikiDocumentContribution(entry.DocID, wikiProducts) + result.currentByDoc[docID] = buildWikiDocumentContribution(docID, wikiProducts) } allDocIDs := make([]string, 0, len(completed)+len(deleted))