⚡️ perf(query-completion): 限制大元数据候选构建

2026-08-16 11:54:08 +08:00 · 2026-07-22 01:47:14 +08:00
parent f9b685f5e7
commit ea7c05f28f
4 changed files with 1104 additions and 320 deletions
--- a/frontend/src/components/queryEditor/QueryEditorHelpers.test.ts
+++ b/frontend/src/components/queryEditor/QueryEditorHelpers.test.ts
@@ -1,9 +1,15 @@
 import { describe, expect, it } from 'vitest';

 import {
+    buildBoundedQueryEditorCompletionSuggestions,
    buildQueryEditorAliasMap,
    collectQueryEditorReferencedDatabaseNames,
+    createBoundedQueryEditorCompletionCandidateBatch,
+    findCompletionTablesByDatabase,
+    getCompletionTableSchemaCounts,
    isOracleBaseTableReference,
+    materializeBoundedQueryEditorCompletionBatches,
+    rankQueryEditorCompletionCandidate,
    resolveOracleLikeDefaultSchemaName,
    resolveOracleLikeExecutionSchemaName,
    resolveOracleLikeLookupSchemaCandidates,
@@ -13,6 +19,161 @@ import {
    shouldHandleQueryEditorRunShortcutFallback,
 } from './QueryEditorHelpers';

+describe('QueryEditor completion candidate budget', () => {
+    it('builds at most the budget after ranking exact, prefix, and substring matches', () => {
+        const candidates = [
+            ...Array.from({ length: 5_000 }, (_, index) => `archive_entity_${String(index).padStart(4, '0')}`),
+            'entity_primary',
+            'entity',
+        ];
+        let materialized = 0;
+
+        const suggestions = buildBoundedQueryEditorCompletionSuggestions({
+            candidates,
+            prefix: 'entity',
+            getMatchRank: (candidate, prefix) => rankQueryEditorCompletionCandidate(prefix, [candidate]),
+            getSelectionKey: (candidate, _prefix, rank) => `${rank}${candidate}`,
+            buildSuggestion: (candidate) => {
+                materialized += 1;
+                const rank = rankQueryEditorCompletionCandidate('entity', [candidate]);
+                return { label: candidate, sortText: `${rank}${candidate}` };
+            },
+        });
+
+        expect(suggestions).toHaveLength(200);
+        expect(materialized).toBe(200);
+        expect(suggestions.slice(0, 2).map((item) => item.label)).toEqual(['entity', 'entity_primary']);
+    });
+
+    it('stops scanning an empty-prefix source once the budget is full', () => {
+        let inspected = 0;
+        const candidates = Array.from({ length: 10_000 }, (_, index) => `table_${String(index).padStart(5, '0')}`);
+
+        const suggestions = buildBoundedQueryEditorCompletionSuggestions({
+            candidates,
+            prefix: '',
+            getMatchRank: () => {
+                inspected += 1;
+                return 0;
+            },
+            getSelectionKey: (candidate) => candidate,
+            buildSuggestion: (candidate) => ({ label: candidate }),
+            sourceAlreadySortedBySelection: true,
+        });
+
+        expect(suggestions).toHaveLength(200);
+        expect(inspected).toBe(200);
+    });
+
+    it('keeps a late same-rank candidate when its final sort key is better', () => {
+        const candidates = [
+            ...Array.from({ length: 200 }, (_, index) => ({
+                label: `other_${index}`,
+                sortText: `10${String(index).padStart(3, '0')}`,
+            })),
+            { label: 'current_late', sortText: '00current_late' },
+        ];
+        let materialized = 0;
+
+        const suggestions = buildBoundedQueryEditorCompletionSuggestions({
+            candidates,
+            prefix: '',
+            getMatchRank: () => 0,
+            getSelectionKey: (candidate) => candidate.sortText,
+            buildSuggestion: (candidate) => {
+                materialized += 1;
+                return candidate;
+            },
+        });
+
+        expect(suggestions).toHaveLength(200);
+        expect(materialized).toBe(200);
+        expect(suggestions.map((item) => item.label)).toContain('current_late');
+    });
+
+    it('uses final sortText semantics when a late exact match ranks after current-database prefixes', () => {
+        const candidates = [
+            ...Array.from({ length: 200 }, (_, index) => ({
+                label: `current_prefix_${String(index).padStart(3, '0')}`,
+                matchRank: 1 as const,
+                sortText: `00current_${String(index).padStart(3, '0')}`,
+            })),
+            { label: 'other_exact', matchRank: 0 as const, sortText: '01other_exact' },
+        ];
+
+        const suggestions = buildBoundedQueryEditorCompletionSuggestions({
+            candidates,
+            prefix: 'target',
+            getMatchRank: (candidate) => candidate.matchRank,
+            getSelectionKey: (candidate) => candidate.sortText,
+            buildSuggestion: (candidate) => candidate,
+        });
+
+        expect(suggestions).toHaveLength(200);
+        expect(suggestions.map((item) => item.label)).not.toContain('other_exact');
+        expect(suggestions[199]?.label).toBe('current_prefix_199');
+    });
+
+    it('materializes at most one global budget across nine completion categories', () => {
+        let materialized = 0;
+        const batches = Array.from({ length: 9 }, (_, groupIndex) => (
+            createBoundedQueryEditorCompletionCandidateBatch({
+                candidates: Array.from({ length: 500 }, (_, candidateIndex) => ({
+                    label: `group_${groupIndex}_${candidateIndex}`,
+                    sortText: `${String(groupIndex).padStart(2, '0')}${String(candidateIndex).padStart(3, '0')}`,
+                })),
+                prefix: '',
+                getMatchRank: () => 0,
+                getSelectionKey: (candidate) => candidate.sortText,
+                buildSuggestion: (candidate) => {
+                    materialized += 1;
+                    return candidate;
+                },
+            })
+        ));
+
+        expect(materialized).toBe(0);
+        const suggestions = materializeBoundedQueryEditorCompletionBatches(batches);
+
+        expect(suggestions).toHaveLength(200);
+        expect(materialized).toBe(200);
+        expect(suggestions[0]?.label).toBe('group_0_0');
+        expect(suggestions[199]?.label).toBe('group_0_199');
+    });
+
+    it('caches schema counts on the current-database partition without reading other table names', () => {
+        let currentTableNameReads = 0;
+        let otherTableNameReads = 0;
+        const currentTables = Array.from({ length: 3 }, (_, index) => ({
+            dbName: 'main',
+            get tableName() {
+                currentTableNameReads += 1;
+                return index < 2 ? `schema_${index}.users` : 'orders';
+            },
+        }));
+        const otherTables = Array.from({ length: 5_000 }, (_, index) => ({
+            dbName: 'archive',
+            get tableName() {
+                otherTableNameReads += 1;
+                return `archive_${index}`;
+            },
+        }));
+        const allTables = [...otherTables, ...currentTables];
+
+        const firstPartition = findCompletionTablesByDatabase(allTables, 'main');
+        const firstCounts = getCompletionTableSchemaCounts(firstPartition);
+        const secondPartition = findCompletionTablesByDatabase(allTables, 'main');
+        const secondCounts = getCompletionTableSchemaCounts(secondPartition);
+
+        expect(firstPartition).toBe(secondPartition);
+        expect(firstCounts).toBe(secondCounts);
+        expect(firstCounts.get('users')).toBe(2);
+        expect(firstCounts.get('orders')).toBe(1);
+        expect(currentTableNameReads).toBe(3);
+        expect(otherTableNameReads).toBe(0);
+    });
+});
+
 describe('QueryEditor Monaco SQL grammar', () => {
    it.each([
        [{ config: { type: 'mysql' } }, 'mysql'],
--- a/frontend/src/components/queryEditor/QueryEditorHelpers.ts
+++ b/frontend/src/components/queryEditor/QueryEditorHelpers.ts
@@ -28,6 +28,271 @@ export type CompletionRoutineMeta = {dbName: string, routineName: string, routin
 export type CompletionSequenceMeta = {dbName: string, sequenceName: string, schemaName?: string};
 export type CompletionPackageMeta = {dbName: string, packageName: string, schemaName?: string};

+// Metadata refreshes replace the source array, so identity-keyed partitions stay correct and let
+// repeated completion requests avoid allocating/scanning another full-database filter result.
+const completionTablesByDatabaseCache = new WeakMap<CompletionTableMeta[], Map<string, CompletionTableMeta[]>>();
+
+export const findCompletionTablesByDatabase = (
+    tables: CompletionTableMeta[],
+    dbName: string,
+): CompletionTableMeta[] => {
+    let index = completionTablesByDatabaseCache.get(tables);
+    if (!index) {
+        index = new Map<string, CompletionTableMeta[]>();
+        tables.forEach((table) => {
+            const key = String(table.dbName || '').trim().toLowerCase();
+            const matches = index!.get(key);
+            if (matches) {
+                matches.push(table);
+            } else {
+                index!.set(key, [table]);
+            }
+        });
+        completionTablesByDatabaseCache.set(tables, index);
+    }
+    return index.get(String(dbName || '').trim().toLowerCase()) || [];
+};
+
+export const QUERY_EDITOR_COMPLETION_SUGGESTION_LIMIT = 200;
+
+export type QueryEditorCompletionMatchRank = 0 | 1 | 2 | null;
+
+export const rankQueryEditorCompletionCandidate = (
+    prefix: string,
+    candidates: readonly string[],
+    includeSubstring = true,
+): QueryEditorCompletionMatchRank => {
+    const normalizedPrefix = String(prefix || '').trim().toLowerCase();
+    if (!normalizedPrefix) return 0;
+
+    let hasPrefixMatch = false;
+    let hasSubstringMatch = false;
+    for (const candidate of candidates) {
+        const normalizedCandidate = String(candidate || '').trim().toLowerCase();
+        if (!normalizedCandidate) continue;
+        if (normalizedCandidate === normalizedPrefix) return 0;
+        if (normalizedCandidate.startsWith(normalizedPrefix)) {
+            hasPrefixMatch = true;
+        } else if (includeSubstring && normalizedCandidate.includes(normalizedPrefix)) {
+            hasSubstringMatch = true;
+        }
+    }
+    if (hasPrefixMatch) return 1;
+    if (hasSubstringMatch) return 2;
+    return null;
+};
+
+type RankedQueryEditorCompletionCandidate<Candidate> = {
+    candidate: Candidate;
+    selectionKey: string;
+    sourceIndex: number;
+};
+
+export type QueryEditorCompletionCandidateBatch<Candidate, Suggestion> = {
+    rankedCandidates: RankedQueryEditorCompletionCandidate<Candidate>[];
+    buildSuggestion: (candidate: Candidate) => Suggestion;
+};
+
+const compareRankedQueryEditorCompletionCandidates = <Candidate,>(
+    left: RankedQueryEditorCompletionCandidate<Candidate>,
+    right: RankedQueryEditorCompletionCandidate<Candidate>,
+): number => {
+    if (left.selectionKey < right.selectionKey) return -1;
+    if (left.selectionKey > right.selectionKey) return 1;
+    return left.sourceIndex - right.sourceIndex;
+};
+
+export const createBoundedQueryEditorCompletionCandidateBatch = <Candidate, Suggestion>({
+    candidates,
+    prefix,
+    getMatchRank,
+    getSelectionKey,
+    buildSuggestion,
+    limit = QUERY_EDITOR_COMPLETION_SUGGESTION_LIMIT,
+    sourceAlreadySortedBySelection = false,
+}: {
+    candidates: readonly Candidate[];
+    prefix: string;
+    getMatchRank: (candidate: Candidate, normalizedPrefix: string) => QueryEditorCompletionMatchRank;
+    getSelectionKey: (
+        candidate: Candidate,
+        normalizedPrefix: string,
+        matchRank: Exclude<QueryEditorCompletionMatchRank, null>,
+    ) => string;
+    buildSuggestion: (candidate: Candidate) => Suggestion;
+    limit?: number;
+    sourceAlreadySortedBySelection?: boolean;
+}): QueryEditorCompletionCandidateBatch<Candidate, Suggestion> => {
+    const normalizedLimit = Math.max(0, Math.floor(Number(limit) || 0));
+    if (normalizedLimit === 0 || candidates.length === 0) {
+        return { rankedCandidates: [], buildSuggestion };
+    }
+
+    const normalizedPrefix = String(prefix || '').trim().toLowerCase();
+    const compareCandidateToRanked = (
+        selectionKey: string,
+        sourceIndex: number,
+        right: RankedQueryEditorCompletionCandidate<Candidate>,
+    ): number => {
+        if (selectionKey < right.selectionKey) return -1;
+        if (selectionKey > right.selectionKey) return 1;
+        return sourceIndex - right.sourceIndex;
+    };
+    // Max-heap: the worst retained candidate stays at index 0 and can be replaced in O(log limit).
+    const selectedHeap: RankedQueryEditorCompletionCandidate<Candidate>[] = [];
+    const siftUpWorst = (startIndex: number) => {
+        let childIndex = startIndex;
+        while (childIndex > 0) {
+            const parentIndex = Math.floor((childIndex - 1) / 2);
+            if (compareRankedQueryEditorCompletionCandidates(selectedHeap[parentIndex], selectedHeap[childIndex]) >= 0) break;
+            [selectedHeap[parentIndex], selectedHeap[childIndex]] = [selectedHeap[childIndex], selectedHeap[parentIndex]];
+            childIndex = parentIndex;
+        }
+    };
+    const siftDownWorst = (startIndex: number) => {
+        let parentIndex = startIndex;
+        while (true) {
+            const leftIndex = parentIndex * 2 + 1;
+            if (leftIndex >= selectedHeap.length) break;
+            const rightIndex = leftIndex + 1;
+            let worseChildIndex = leftIndex;
+            if (
+                rightIndex < selectedHeap.length
+                && compareRankedQueryEditorCompletionCandidates(selectedHeap[rightIndex], selectedHeap[leftIndex]) > 0
+            ) {
+                worseChildIndex = rightIndex;
+            }
+            if (compareRankedQueryEditorCompletionCandidates(selectedHeap[parentIndex], selectedHeap[worseChildIndex]) >= 0) break;
+            [selectedHeap[parentIndex], selectedHeap[worseChildIndex]] = [selectedHeap[worseChildIndex], selectedHeap[parentIndex]];
+            parentIndex = worseChildIndex;
+        }
+    };
+
+    for (let sourceIndex = 0; sourceIndex < candidates.length; sourceIndex += 1) {
+        const candidate = candidates[sourceIndex];
+        const rank = getMatchRank(candidate, normalizedPrefix);
+        if (rank === null) continue;
+        const selectionKey = String(getSelectionKey(candidate, normalizedPrefix, rank) || '');
+        if (selectedHeap.length < normalizedLimit) {
+            selectedHeap.push({ candidate, selectionKey, sourceIndex });
+            siftUpWorst(selectedHeap.length - 1);
+        } else {
+            const worst = selectedHeap[0];
+            if (compareCandidateToRanked(selectionKey, sourceIndex, worst) < 0) {
+                // Reuse the root entry so descending input cannot allocate one retained wrapper per source row.
+                worst.candidate = candidate;
+                worst.selectionKey = selectionKey;
+                worst.sourceIndex = sourceIndex;
+                siftDownWorst(0);
+            }
+        }
+        // Early-stop is safe only when the caller guarantees the source already follows the same
+        // final selection-key + stable-input-order tuple used by this bounded top-k.
+        if (sourceAlreadySortedBySelection && selectedHeap.length >= normalizedLimit) {
+            break;
+        }
+    }
+
+    selectedHeap.sort(compareRankedQueryEditorCompletionCandidates);
+    return { rankedCandidates: selectedHeap, buildSuggestion };
+};
+
+export const materializeBoundedQueryEditorCompletionBatches = <Suggestion,>(
+    batches: readonly QueryEditorCompletionCandidateBatch<any, Suggestion>[],
+    limit = QUERY_EDITOR_COMPLETION_SUGGESTION_LIMIT,
+): Suggestion[] => {
+    const normalizedLimit = Math.max(0, Math.floor(Number(limit) || 0));
+    if (normalizedLimit === 0 || batches.length === 0) return [];
+
+    type GlobalRankedCandidate = {
+        batchIndex: number;
+        candidateIndex: number;
+        selectionKey: string;
+        sourceIndex: number;
+    };
+    const compareGlobalCandidates = (left: GlobalRankedCandidate, right: GlobalRankedCandidate): number => {
+        if (left.selectionKey < right.selectionKey) return -1;
+        if (left.selectionKey > right.selectionKey) return 1;
+        if (left.batchIndex !== right.batchIndex) return left.batchIndex - right.batchIndex;
+        return left.sourceIndex - right.sourceIndex;
+    };
+    const compareCandidateToGlobal = (
+        batchIndex: number,
+        selectionKey: string,
+        sourceIndex: number,
+        right: GlobalRankedCandidate,
+    ): number => {
+        if (selectionKey < right.selectionKey) return -1;
+        if (selectionKey > right.selectionKey) return 1;
+        if (batchIndex !== right.batchIndex) return batchIndex - right.batchIndex;
+        return sourceIndex - right.sourceIndex;
+    };
+    const selectedHeap: GlobalRankedCandidate[] = [];
+    const siftUpWorst = (startIndex: number) => {
+        let childIndex = startIndex;
+        while (childIndex > 0) {
+            const parentIndex = Math.floor((childIndex - 1) / 2);
+            if (compareGlobalCandidates(selectedHeap[parentIndex], selectedHeap[childIndex]) >= 0) break;
+            [selectedHeap[parentIndex], selectedHeap[childIndex]] = [selectedHeap[childIndex], selectedHeap[parentIndex]];
+            childIndex = parentIndex;
+        }
+    };
+    const siftDownWorst = (startIndex: number) => {
+        let parentIndex = startIndex;
+        while (true) {
+            const leftIndex = parentIndex * 2 + 1;
+            if (leftIndex >= selectedHeap.length) break;
+            const rightIndex = leftIndex + 1;
+            let worseChildIndex = leftIndex;
+            if (
+                rightIndex < selectedHeap.length
+                && compareGlobalCandidates(selectedHeap[rightIndex], selectedHeap[leftIndex]) > 0
+            ) {
+                worseChildIndex = rightIndex;
+            }
+            if (compareGlobalCandidates(selectedHeap[parentIndex], selectedHeap[worseChildIndex]) >= 0) break;
+            [selectedHeap[parentIndex], selectedHeap[worseChildIndex]] = [selectedHeap[worseChildIndex], selectedHeap[parentIndex]];
+            parentIndex = worseChildIndex;
+        }
+    };
+
+    batches.forEach((batch, batchIndex) => {
+        batch.rankedCandidates.forEach((candidate, candidateIndex) => {
+            if (selectedHeap.length < normalizedLimit) {
+                selectedHeap.push({
+                    batchIndex,
+                    candidateIndex,
+                    selectionKey: candidate.selectionKey,
+                    sourceIndex: candidate.sourceIndex,
+                });
+                siftUpWorst(selectedHeap.length - 1);
+                return;
+            }
+            const worst = selectedHeap[0];
+            if (compareCandidateToGlobal(batchIndex, candidate.selectionKey, candidate.sourceIndex, worst) < 0) {
+                worst.batchIndex = batchIndex;
+                worst.candidateIndex = candidateIndex;
+                worst.selectionKey = candidate.selectionKey;
+                worst.sourceIndex = candidate.sourceIndex;
+                siftDownWorst(0);
+            }
+        });
+    });
+
+    selectedHeap.sort(compareGlobalCandidates);
+    return selectedHeap.map(({ batchIndex, candidateIndex }) => {
+        const batch = batches[batchIndex];
+        return batch.buildSuggestion(batch.rankedCandidates[candidateIndex].candidate);
+    });
+};
+
+export const buildBoundedQueryEditorCompletionSuggestions = <Candidate, Suggestion>(
+    options: Parameters<typeof createBoundedQueryEditorCompletionCandidateBatch<Candidate, Suggestion>>[0],
+): Suggestion[] => {
+    const batch = createBoundedQueryEditorCompletionCandidateBatch(options);
+    return materializeBoundedQueryEditorCompletionBatches([batch], options.limit);
+};
+
 export const selectUnqualifiedCompletionSynonyms = (
    synonyms: CompletionSynonymMeta[],
    loginOwnerName: string,
@@ -743,6 +1008,27 @@ export const splitCompletionSchemaAndTable = (qualified: string): { schema: stri
    return { schema: '', table: parts[0] || '' };
 };

+// The caller passes the cached current-database partition. Schema duplicate detection therefore
+// reads only that partition once instead of walking every visible database on each keystroke.
+const completionTableSchemaCountCache = new WeakMap<CompletionTableMeta[], Map<string, number>>();
+
+export const getCompletionTableSchemaCounts = (
+    currentDatabaseTables: CompletionTableMeta[],
+): Map<string, number> => {
+    const cached = completionTableSchemaCountCache.get(currentDatabaseTables);
+    if (cached) return cached;
+
+    const counts = new Map<string, number>();
+    currentDatabaseTables.forEach((table) => {
+        const parsed = splitCompletionSchemaAndTable(table.tableName || '');
+        const pureTable = String(parsed.table || table.tableName || '').toLowerCase();
+        if (!pureTable) return;
+        counts.set(pureTable, (counts.get(pureTable) || 0) + 1);
+    });
+    completionTableSchemaCountCache.set(currentDatabaseTables, counts);
+    return counts;
+};
+
 export const DEFAULT_QUERY_TEMPLATE = 'SELECT * FROM ';

 export const resolveNewQueryDefaultTemplate = (