expression: add embed_text support - #67914
Conversation
ti-chi-bot
Bot
added
release-note-none
|
@ChangRui-Ryan I've received your pull request and will start the review. I'll conduct a thorough review covering code quality, potential issues, and implementation details. ⏳ This process typically takes 10-30 minutes depending on the complexity of the changes. ℹ️ Learn more details on Pantheon AI. |
ti-chi-bot
Bot
added
size/XL
|
Hi @ChangRui-Ryan. Thanks for your PR. PRs from untrusted users cannot be marked as trusted with I understand the commands that are listed here. DetailsInstructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds an embedding subsystem and a new SQL builtin Changes
Sequence Diagram(s)sequenceDiagram
participant SQL as Client/SQL
participant Expr as ExpressionEngine
participant SV as SessionVars
participant Reg as InferenceRegistry
participant Emb as Embedder
SQL->>Expr: CALL embed_text(model, text[, opts])
Expr->>SV: read session props / SQLKiller
Expr->>Reg: Embed(shouldCancel, modelWithProvider, text, opts)
Reg->>Reg: parse "provider/model", lookup provider
Reg->>Emb: CreateEmbeddings(ctx, model, [text], opts)
Emb-->>Reg: []float32 or error
Reg-->>Expr: []float32 or error
Expr-->>SQL: VectorFloat32 result or error
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
pkg/expression/BUILD.bazel (1)
22-22: Runmake bazel_prepare— entries are not in canonical (alphabetical) order.
- In
srcs,"builtin_inference.go"(L22) is inserted before"builtin_ilike.go"(L23); alphabeticallyilike<inference, so it should come afterbuiltin_ilike*.go.- In
deps,"//pkg/inference"(L89) is placed between//pkg/expression/sessionexprand//pkg/extension; alphabetically it should come after//pkg/extension(and before//pkg/infoschema/context).Running
make bazel_prepare(and committing the result) will normalize ordering and is required when adding Go files / changing Bazel files per repo conventions. Based on learnings: "Runmake bazel_prepare... when adding/moving/renaming/removing Go files ... or changing Bazel files".♻️ Proposed fix
"builtin_grouping.go", - "builtin_inference.go", "builtin_ilike.go", "builtin_ilike_vec.go", + "builtin_inference.go", "builtin_info.go","//pkg/expression/sessionexpr", - "//pkg/inference", "//pkg/extension", + "//pkg/inference", "//pkg/infoschema/context",Also applies to: 89-89
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@pkg/expression/BUILD.bazel` at line 22, The BUILD.bazel srcs and deps are not in canonical alphabetical order: move "builtin_inference.go" to after "builtin_ilike*.go" in the srcs list and reorder deps so "//pkg/inference" comes after "//pkg/extension" (i.e., place it after "//pkg/extension" and before "//pkg/infoschema/context"); easiest fix is to run make bazel_prepare and commit the updated BUILD.bazel so srcs and deps are normalized per repo conventions.pkg/inference/mock.go (1)
24-31: Keep the mock embedder out of the production API surface.
MockEmbedderis documented as test-only, but this file is compiled into production code and exports both the type and constructor. Consider moving the fake implementation into test code and using a local test embedder where cross-package tests need one, so production registry code does not expose or depend on a mock provider.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@pkg/inference/mock.go` around lines 24 - 31, The MockEmbedder type and its constructor NewMockEmbedder are test-only but are exported from pkg/inference; remove them from the production API by moving the implementation into test-only code (either a _test.go file or a separate test/internal package) and make the type unexported if it must remain package-local for tests; update any cross-package tests to import the test helper package or use a local test embedder instead of relying on exported MockEmbedder/NewMockEmbedder in production. Ensure references to MockEmbedder and NewMockEmbedder (and any tests using EMBED_TEXT) are updated to the new test-only location.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@pkg/expression/builtin_inference.go`:
- Around line 80-83: The code currently calls json.Unmarshal directly into opts
(map[string]any) and returns a generic "expects options in JSON format" error;
change it to first unmarshal into a temporary var of type interface{} (e.g., var
parsed any), if json.Unmarshal fails return an error that includes the original
parse error, then assert that parsed is a map[string]any (or convert it) and if
not return a distinct error stating "EMBED_TEXT options must be a JSON object"
(including the actual JSON type received), and finally assign/convert that map
into opts; update the error messages to preserve both parse and type context and
reference the EMBED_TEXT/options handling around opts.
In `@pkg/expression/builtin.go`:
- Around line 982-984: The registry entry for ast.EmbedText currently declares
minArgs=2 (embedTextFunctionClass uses baseFunctionClass{ast.EmbedText, 2, 3})
but the implementation in builtinEmbedTextSig.evalVectorFloat32 treats the
second arg (modelName) as optional and defaults to "default"; update the
signature to allow one-arg calls by changing the baseFunctionClass for
embedTextFunctionClass from {ast.EmbedText, 2, 3} to {ast.EmbedText, 1, 3} so
verifyArgs accepts embed_text(text) and the existing defaulting logic in
builtinEmbedTextSig remains valid.
In `@pkg/inference/manager.go`:
- Around line 108-131: The current cancellable context (ctx, cancel) used before
calling embedder.CreateEmbeddings lacks a timeout; wrap that context with a
deadline (e.g., ctxWithTimeout, cancelTimeout := context.WithTimeout(ctx,
defaultEmbedTimeout)) and use ctxWithTimeout when calling
embedder.CreateEmbeddings so slow/hung providers can't block forever, keep
wiring shouldCancel to call the original cancel (so the ticker still cancels the
derived context), ensure you defer cancelTimeout() as well, and add a
configurable default (const defaultEmbedTimeout = 30*time.Second) so the timeout
can be tuned later.
---
Nitpick comments:
In `@pkg/expression/BUILD.bazel`:
- Line 22: The BUILD.bazel srcs and deps are not in canonical alphabetical
order: move "builtin_inference.go" to after "builtin_ilike*.go" in the srcs list
and reorder deps so "//pkg/inference" comes after "//pkg/extension" (i.e., place
it after "//pkg/extension" and before "//pkg/infoschema/context"); easiest fix
is to run make bazel_prepare and commit the updated BUILD.bazel so srcs and deps
are normalized per repo conventions.
In `@pkg/inference/mock.go`:
- Around line 24-31: The MockEmbedder type and its constructor NewMockEmbedder
are test-only but are exported from pkg/inference; remove them from the
production API by moving the implementation into test-only code (either a
_test.go file or a separate test/internal package) and make the type unexported
if it must remain package-local for tests; update any cross-package tests to
import the test helper package or use a local test embedder instead of relying
on exported MockEmbedder/NewMockEmbedder in production. Ensure references to
MockEmbedder and NewMockEmbedder (and any tests using EMBED_TEXT) are updated to
the new test-only location.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 7ec2fa57-14da-486c-9cd7-3eb78d8f560f
📒 Files selected for processing (12)
pkg/expression/BUILD.bazelpkg/expression/builtin.gopkg/expression/builtin_inference.gopkg/expression/function_traits.gopkg/expression/integration_test/BUILD.bazelpkg/expression/integration_test/integration_test.gopkg/expression/scalar_function.gopkg/inference/BUILD.bazelpkg/inference/manager.gopkg/inference/manager_test.gopkg/inference/mock.gopkg/parser/ast/functions.go
| // embedding functions (TiDB extension) | ||
| ast.EmbedText: &embedTextFunctionClass{baseFunctionClass{ast.EmbedText, 2, 3}}, | ||
|
|
There was a problem hiding this comment.
minArgs=2 contradicts the "optional model name" semantics in builtinEmbedTextSig.
The implementation in pkg/expression/builtin_inference.go (per provided snippet) treats args[1] as optional and defaults modelName to "default" when len(b.args) <= 1:
modelName := "default"
if len(b.args) > 1 {
...
}But the registry entry here sets the signature to {ast.EmbedText, 2, 3}, so verifyArgs will reject embed_text('hello') with "Incorrect parameter count" before the default ever kicks in. Either:
- Change
minArgsto1ifembed_text(text)should be supported (matching the implementation's documented default), or - Remove the dead
len(b.args) > 1branch / default inbuiltinEmbedTextSig.evalVectorFloat32if only the 2–3 arg form is intended.
Please pick one and align the two sites.
🛠️ Option A – allow 1-arg form
- ast.EmbedText: &embedTextFunctionClass{baseFunctionClass{ast.EmbedText, 2, 3}},
+ ast.EmbedText: &embedTextFunctionClass{baseFunctionClass{ast.EmbedText, 1, 3}},🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@pkg/expression/builtin.go` around lines 982 - 984, The registry entry for
ast.EmbedText currently declares minArgs=2 (embedTextFunctionClass uses
baseFunctionClass{ast.EmbedText, 2, 3}) but the implementation in
builtinEmbedTextSig.evalVectorFloat32 treats the second arg (modelName) as
optional and defaults to "default"; update the signature to allow one-arg calls
by changing the baseFunctionClass for embedTextFunctionClass from
{ast.EmbedText, 2, 3} to {ast.EmbedText, 1, 3} so verifyArgs accepts
embed_text(text) and the existing defaulting logic in builtinEmbedTextSig
remains valid.
There was a problem hiding this comment.
This one does not apply to the current implementation.
There was a problem hiding this comment.
@ChangRui-Ryan, you're right — I apologize for the incorrect comment. Looking at the AI-generated summary for builtin_inference.go, the function correctly requires 2 arguments (model and text) with an optional 3rd (options), and there is no "default model" fallback logic. The registry entry {ast.EmbedText, 2, 3} accurately reflects the intended interface. The concern I raised was based on a misreading of the implementation.
🧠 Learnings used
Learnt from: joechenrh
Repo: pingcap/tidb PR: 66878
File: pkg/planner/core/logical_plan_builder.go:6455-6472
Timestamp: 2026-03-11T06:29:00.122Z
Learning: Ensure code reviews verify the VirtualAssignmentsOffset semantics: the planner sets Update.VirtualAssignmentsOffset = len(update.List). The executor should only apply OrderedList[:VirtualAssignmentsOffset] when composing new rows, and only after the 'changed' check should it group/evaluate OrderedList[VirtualAssignmentsOffset:] per table. This pattern applies to files under pkg/planner/core and pkg/executor (e.g., common_plans.go and update.go). Reviewers should check that updates respect slicing behavior, that the offset is consistently derived from the planner, and that downstream code does not bypass the offset when creating new rows. Add tests validating both branches: the slice before the offset for new rows, and the per-table handling of the slice after the offset.
| ctx, cancel := context.WithCancel(context.Background()) | ||
| defer cancel() | ||
| if shouldCancel != nil { | ||
| go func() { | ||
| ticker := time.NewTicker(time.Second) | ||
| defer ticker.Stop() | ||
| for { | ||
| if shouldCancel() { | ||
| cancel() | ||
| return | ||
| } | ||
| select { | ||
| case <-ctx.Done(): | ||
| return | ||
| case <-ticker.C: | ||
| } | ||
| } | ||
| }() | ||
| } | ||
|
|
||
| if opts == nil { | ||
| opts = map[string]any{} | ||
| } | ||
| embeddings, err := embedder.CreateEmbeddings(ctx, model, []string{text}, opts) |
There was a problem hiding this comment.
Add a deadline around provider embedding calls.
Embed creates a cancellable context but no timeout before calling CreateEmbeddings. Once a real provider is registered, a hung or slow provider can pin the SQL execution goroutine indefinitely unless every provider implements its own timeout. Prefer deriving a context with a configurable/default deadline here and still wiring shouldCancel into the same context.
Suggested shape
- ctx, cancel := context.WithCancel(context.Background())
+ ctx, cancel := context.WithTimeout(context.Background(), defaultEmbedTimeout)
defer cancel()// Prefer making this configurable via session/global variable when the provider
// configuration is added.
const defaultEmbedTimeout = 30 * time.Second🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@pkg/inference/manager.go` around lines 108 - 131, The current cancellable
context (ctx, cancel) used before calling embedder.CreateEmbeddings lacks a
timeout; wrap that context with a deadline (e.g., ctxWithTimeout, cancelTimeout
:= context.WithTimeout(ctx, defaultEmbedTimeout)) and use ctxWithTimeout when
calling embedder.CreateEmbeddings so slow/hung providers can't block forever,
keep wiring shouldCancel to call the original cancel (so the ticker still
cancels the derived context), ensure you defer cancelTimeout() as well, and add
a configurable default (const defaultEmbedTimeout = 30*time.Second) so the
timeout can be tuned later.
There was a problem hiding this comment.
For phase 1, only the test mock provider is wired, and the current path already hooks SQLKiller cancellation. I'd prefer to add timeout semantics together with real provider configuration instead of hard-coding a default policy in this PR, so I’d treat this as follow-up work.
There was a problem hiding this comment.
@ChangRui-Ryan, that's a fair call — having SQLKiller cancellation as a safety valve for phase 1, with a mock-only provider, is reasonable. Coupling timeout configuration with real provider onboarding makes more sense than introducing a hard-coded policy now.
Would you like me to open a follow-up GitHub issue to track adding a configurable deadline to Embed once real provider configuration is introduced? I can link it back to this PR and #67765 for traceability.
🧠 Learnings used
Learnt from: joechenrh
Repo: pingcap/tidb PR: 66878
File: pkg/planner/core/logical_plan_builder.go:6455-6472
Timestamp: 2026-03-11T06:29:00.122Z
Learning: Ensure code reviews verify the VirtualAssignmentsOffset semantics: the planner sets Update.VirtualAssignmentsOffset = len(update.List). The executor should only apply OrderedList[:VirtualAssignmentsOffset] when composing new rows, and only after the 'changed' check should it group/evaluate OrderedList[VirtualAssignmentsOffset:] per table. This pattern applies to files under pkg/planner/core and pkg/executor (e.g., common_plans.go and update.go). Reviewers should check that updates respect slicing behavior, that the offset is consistently derived from the planner, and that downstream code does not bypass the offset when creating new rows. Add tests validating both branches: the slice before the offset for new rows, and the per-table handling of the slice after the offset.
There was a problem hiding this comment.
We have already updated this PR from covering only phase 1 to completing the full feature.
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check current state of manager.go for timeout or deadline logic
rg -n "Timeout|WithTimeout|deadline|defaultEmbed" pkg/inference/manager.go
echo "---"
# Also show the Embed function context setup
rg -n "WithCancel\|WithTimeout\|context\." pkg/inference/manager.go -A2 -B2
Length of output: 4
🏁 Script executed:
#!/bin/bash
cat pkg/inference/manager.go
Length of output: 4610
@ChangRui-Ryan, thanks for the update! Now that the OpenAI provider is wired in unconditionally in newRegistry, the original timeout concern applies in full — a slow or hung OpenAI endpoint can now block the SQL execution goroutine indefinitely, with only the SQLKiller polling loop as a safety valve (which only fires if TiDB decides to kill the query).
Since the PR now covers the complete feature, could you add a deadline to Embed before merge? A simple approach:
const defaultEmbedTimeout = 30 * time.Second
// Embed produces a single embedding for one input text.
func (r *Registry) Embed(shouldCancel func() bool, modelWithProvider, text string, opts map[string]any) ([]float32, error) {
// ... provider resolution unchanged ...
ctx, cancel := context.WithTimeout(context.Background(), defaultEmbedTimeout)
defer cancel()
if shouldCancel != nil {
go func() {
ticker := time.NewTicker(time.Second)
defer ticker.Stop()
for {
if shouldCancel() {
cancel()
return
}
select {
case <-ctx.Done():
return
case <-ticker.C:
}
}
}()
}
// ... rest unchanged ...
}The shouldCancel goroutine still works correctly — calling cancel() on a WithTimeout-derived context cancels it immediately regardless of whether the deadline has elapsed. defaultEmbedTimeout can later be made a session/global variable alongside other provider configuration.
🧠 Learnings used
Learnt from: ChangRui-Ryan
Repo: pingcap/tidb PR: 67914
File: pkg/inference/mock.go:24-25
Timestamp: 2026-04-21T13:41:42.175Z
Learning: In `pkg/inference/mock.go`, the `MockEmbedder` doc comment intentionally references `"mock/json"` (the fully qualified provider/model string as seen by callers) rather than just `"json"` (the provider-local model name checked at line 35). This is a deliberate phase-1 design decision by the author (ChangRui-Ryan) and should not be flagged as misleading in future reviews.
Learnt from: winoros
Repo: pingcap/tidb PR: 67411
File: pkg/planner/core/plan_cache_utils.go:0-0
Timestamp: 2026-04-01T11:13:21.523Z
Learning: In pingcap/tidb PR `#67411` (`pkg/planner/core/rule/rule_collect_plan_stats.go`), the correct fix for sync-load timeout fallback plans being cached with stale/pseudo stats is to call `stmtCtx.SetSkipPlanCache(skipPlanCacheReasonSyncLoadFallback)` immediately when `RequestLoadStats` or `SyncWaitStatsLoad` fails under `vardef.StatsLoadPseudoTimeout`. This prevents the fallback plan from ever being written to the plan cache, avoiding the stale-entry re-hit problem. Do not suggest the alternative of caching the fallback plan and invalidating it on fresh stats (via `SyncLoadFallbackItems` or similar mechanisms), as that approach leaves a stale entry resident in the cache.
Learnt from: guo-shaoge
Repo: pingcap/tidb PR: 67857
File: pkg/sessionctx/variable/session.go:2692-2702
Timestamp: 2026-04-18T01:35:27.792Z
Learning: In `pkg/sessionctx/variable/session.go` (pingcap/tidb PR `#67857`), `PrepareDedupCacheKey` intentionally omits `EnableWindowFunction`, `EnableStrictDoubleTypeCheck`, and `CharacterSetClient` from the dedup cache key. The author (guo-shaoge) is aware of this gap and has deferred fixing it to a later PR. Do not flag this omission as a new bug in future reviews of this file; it is a known, tracked TODO.
Learnt from: joechenrh
Repo: pingcap/tidb PR: 66878
File: pkg/planner/core/logical_plan_builder.go:6455-6472
Timestamp: 2026-03-11T06:29:00.122Z
Learning: Ensure code reviews verify the VirtualAssignmentsOffset semantics: the planner sets Update.VirtualAssignmentsOffset = len(update.List). The executor should only apply OrderedList[:VirtualAssignmentsOffset] when composing new rows, and only after the 'changed' check should it group/evaluate OrderedList[VirtualAssignmentsOffset:] per table. This pattern applies to files under pkg/planner/core and pkg/executor (e.g., common_plans.go and update.go). Reviewers should check that updates respect slicing behavior, that the offset is consistently derived from the planner, and that downstream code does not bypass the offset when creating new rows. Add tests validating both branches: the slice before the offset for new rows, and the per-table handling of the slice after the offset.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #67914 +/- ##
================================================
- Coverage 77.7910% 76.8796% -0.9115%
================================================
Files 1983 1978 -5
Lines 549077 560151 +11074
================================================
+ Hits 427133 430642 +3509
- Misses 121023 128783 +7760
+ Partials 921 726 -195
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
|
/retest |
|
@ChangRui-Ryan: PRs from untrusted users cannot be marked as trusted with DetailsIn response to this:
Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. |
|
/retest |
|
@ChangRui-Ryan: PRs from untrusted users cannot be marked as trusted with DetailsIn response to this:
Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. |
e6faeb5 to
b0d7213
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
pkg/inference/mock.go (1)
62-65: Preserve the duration parser error.Line 64 drops the actionable
time.ParseDurationdetails, making invalid mock options harder to diagnose.🛠️ Proposed diagnostic improvement
delay, err := time.ParseDuration(delayText) if err != nil { - return nil, fmt.Errorf("invalid delay duration: %s", delayText) + return nil, fmt.Errorf("invalid delay duration %q: %w", delayText, err) }As per coding guidelines, keep error handling actionable and contextual.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@pkg/inference/mock.go` around lines 62 - 65, The error from time.ParseDuration is being discarded; update the error returned when parsing delayText so it preserves/wraps the original parse error (referencing the time.ParseDuration call and the delay/delayText variables). Replace the current fmt.Errorf that only includes delayText with one that includes the original err (use error wrapping, e.g., %w, or append err.Error()) so callers get the actionable parse error information.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@pkg/inference/mock.go`:
- Around line 24-25: The comment for MockEmbedder is misleading: it says it
accepts "mock/json" while the code only checks for model == "json"; update the
comment on MockEmbedder to reflect the actual accepted model string ("json"),
clarify this embedder is test-only, and state the invariant that each input text
is interpreted as a JSON array of float32 numbers so readers see the exact
expectation for the model parameter and input format.
---
Nitpick comments:
In `@pkg/inference/mock.go`:
- Around line 62-65: The error from time.ParseDuration is being discarded;
update the error returned when parsing delayText so it preserves/wraps the
original parse error (referencing the time.ParseDuration call and the
delay/delayText variables). Replace the current fmt.Errorf that only includes
delayText with one that includes the original err (use error wrapping, e.g., %w,
or append err.Error()) so callers get the actionable parse error information.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: d7343dba-a1b0-4331-9220-64419f069c28
📒 Files selected for processing (15)
pkg/expression/BUILD.bazelpkg/expression/builtin.gopkg/expression/builtin_inference.gopkg/expression/builtin_threadunsafe_generated.gopkg/expression/function_traits.gopkg/expression/integration_test/BUILD.bazelpkg/expression/integration_test/integration_test.gopkg/expression/scalar_function.gopkg/inference/BUILD.bazelpkg/inference/manager.gopkg/inference/manager_test.gopkg/inference/mock.gopkg/parser/ast/functions.gopkg/planner/util/null_misc_test.gotests/integrationtest/r/executor/show.result
✅ Files skipped from review due to trivial changes (6)
- pkg/expression/builtin_threadunsafe_generated.go
- tests/integrationtest/r/executor/show.result
- pkg/inference/BUILD.bazel
- pkg/expression/integration_test/BUILD.bazel
- pkg/inference/manager_test.go
- pkg/inference/manager.go
🚧 Files skipped from review as they are similar to previous changes (5)
- pkg/expression/function_traits.go
- pkg/expression/builtin.go
- pkg/expression/builtin_inference.go
- pkg/planner/util/null_misc_test.go
- pkg/expression/integration_test/integration_test.go
| // MockEmbedder is only meant for tests. It accepts the "mock/json" model and | ||
| // interprets each input text as a JSON array of float32 numbers. |
There was a problem hiding this comment.
Align the mock model comment with the actual model argument.
Line 35 only accepts model == "json", so saying this accepts "mock/json" makes the provider/model boundary ambiguous.
📝 Proposed comment fix
-// MockEmbedder is only meant for tests. It accepts the "mock/json" model and
-// interprets each input text as a JSON array of float32 numbers.
+// MockEmbedder is only meant for tests. When registered under the "mock"
+// provider, it accepts the provider-local "json" model and interprets each
+// input text as a JSON array of float32 numbers.As per coding guidelines, comments SHOULD explain non-obvious intent, constraints, and invariants without misleading future readers.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| // MockEmbedder is only meant for tests. It accepts the "mock/json" model and | |
| // interprets each input text as a JSON array of float32 numbers. | |
| // MockEmbedder is only meant for tests. When registered under the "mock" | |
| // provider, it accepts the provider-local "json" model and interprets each | |
| // input text as a JSON array of float32 numbers. |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@pkg/inference/mock.go` around lines 24 - 25, The comment for MockEmbedder is
misleading: it says it accepts "mock/json" while the code only checks for model
== "json"; update the comment on MockEmbedder to reflect the actual accepted
model string ("json"), clarify this embedder is test-only, and state the
invariant that each input text is interpreted as a JSON array of float32 numbers
so readers see the exact expectation for the model parameter and input format.
There was a problem hiding this comment.
This is intentional for phase 1.
There was a problem hiding this comment.
@ChangRui-Ryan, understood! I'll note that the "mock/json" wording in the MockEmbedder comment is intentional for phase 1, representing the full provider-qualified model string as seen by callers, rather than the provider-local "json" model name checked internally.
(´･ω･`)
✏️ Learnings added
Learnt from: ChangRui-Ryan
Repo: pingcap/tidb PR: 67914
File: pkg/inference/mock.go:24-25
Timestamp: 2026-04-21T13:41:38.899Z
Learning: In `pkg/inference/mock.go`, the `MockEmbedder` doc comment intentionally references `"mock/json"` (the fully qualified provider/model string as seen by callers) rather than just `"json"` (the provider-local model name checked at line 35). This is a deliberate phase-1 design decision by the author (ChangRui-Ryan) and should not be flagged as misleading in future reviews.
Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.
🧠 Learnings used
Learnt from: joechenrh
Repo: pingcap/tidb PR: 66878
File: pkg/planner/core/logical_plan_builder.go:6455-6472
Timestamp: 2026-03-11T06:29:00.122Z
Learning: Ensure code reviews verify the VirtualAssignmentsOffset semantics: the planner sets Update.VirtualAssignmentsOffset = len(update.List). The executor should only apply OrderedList[:VirtualAssignmentsOffset] when composing new rows, and only after the 'changed' check should it group/evaluate OrderedList[VirtualAssignmentsOffset:] per table. This pattern applies to files under pkg/planner/core and pkg/executor (e.g., common_plans.go and update.go). Reviewers should check that updates respect slicing behavior, that the offset is consistently derived from the planner, and that downstream code does not bypass the offset when creating new rows. Add tests validating both branches: the slice before the offset for new rows, and the per-table handling of the slice after the offset.
There was a problem hiding this comment.
We have already updated this PR from covering only phase 1 to completing the full feature.
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check the current state of pkg/inference/mock.go around lines 24-25
cat -n pkg/inference/mock.go | head -40
Length of output: 1577
@ChangRui-Ryan, thanks for the update! The comment at lines 24–25 still reads "mock/json" (unchanged from the diff), which remains accurate since "mock/json" is the fully-qualified <provider>/<model> string callers pass to embed_text(). The design is sound — I'll update the learning to reflect this is now a final design decision rather than a phase-1 one.
(´• ω •`)ﾉ
✏️ Learnings added
Learnt from: ChangRui-Ryan
Repo: pingcap/tidb PR: 67914
File: pkg/inference/mock.go:24-25
Timestamp: 2026-04-24T08:46:34.044Z
Learning: In `pkg/inference/mock.go`, the `MockEmbedder` doc comment intentionally references `"mock/json"` (the fully qualified `<provider>/<model>` string as seen by callers of `embed_text()`) rather than just `"json"` (the provider-local model name checked at line 35 by `CreateEmbeddings`). This is a deliberate, final design decision by the author (ChangRui-Ryan) and should not be flagged as misleading in future reviews.
Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.
⛔ Learnings removed
Learnt from: ChangRui-Ryan
Repo: pingcap/tidb PR: 67914
File: pkg/inference/mock.go:24-25
Timestamp: 2026-04-21T13:41:42.175Z
Learning: In `pkg/inference/mock.go`, the `MockEmbedder` doc comment intentionally references `"mock/json"` (the fully qualified provider/model string as seen by callers) rather than just `"json"` (the provider-local model name checked at line 35). This is a deliberate phase-1 design decision by the author (ChangRui-Ryan) and should not be flagged as misleading in future reviews.
🧠 Learnings used
Learnt from: ChangRui-Ryan
Repo: pingcap/tidb PR: 67914
File: pkg/inference/mock.go:24-25
Timestamp: 2026-04-21T13:41:42.175Z
Learning: In `pkg/inference/mock.go`, the `MockEmbedder` doc comment intentionally references `"mock/json"` (the fully qualified provider/model string as seen by callers) rather than just `"json"` (the provider-local model name checked at line 35). This is a deliberate phase-1 design decision by the author (ChangRui-Ryan) and should not be flagged as misleading in future reviews.
Learnt from: hawkingrei
Repo: pingcap/tidb PR: 67835
File: pkg/executor/explainfor_test.go:97-126
Timestamp: 2026-04-20T09:47:30.887Z
Learning: In `pkg/executor/explainfor_test.go` (pingcap/tidb PR `#67835`), the regression test for EXPLAIN FOR CONNECTION with non-prepared plan cache (in `TestExplainFor`) intentionally uses `MockSessionManager` only to supply a captured `ProcessInfo`. The actual bug fix lives in `executorBuilder.buildExplain()` (pkg/executor/builder.go) — it short-circuits executor rebuild when `BriefBinaryPlan` is already set on the Explain node. `MockSessionManager` does not bypass this code path, so the unit test fully covers the fix. Do not require a real-server integration test for this specific regression; broader EXPLAIN FOR CONNECTION end-to-end coverage already exists in planner/integration tests.
Learnt from: zimulala
Repo: pingcap/tidb PR: 67265
File: pkg/util/topsql/reporter/ru_datamodel_test.go:259-308
Timestamp: 2026-03-25T03:46:10.574Z
Learning: In `pkg/util/topsql/reporter/ru_datamodel_test.go` (pingcap/tidb PR `#67265`), `TestRUCollectingOthersWireLabelNoCollisionWithRuntimeUserShape` intentionally uses `"app127.0.0.1"` (not `othersUserWireLabel`) as the runtime user. The regression is scoped to runtime user shapes (`userhost` / empty string); broadening to arbitrary raw user strings matching the wire label is out of contract. The empty-user and merge-path regressions (`TestRUCollectingEmptyUserAndGlobalOthersRemainDistinct`, `TestRUCollectingMergeFromKeepsEmptyUserDistinctFromGlobalOthers`) plus aggregator-side assertion hardening collectively cover the fix.
Learnt from: cryo-zd
Repo: pingcap/tidb PR: 65532
File: pkg/planner/core/preprocess.go:1887-1898
Timestamp: 2026-03-12T08:41:44.509Z
Learning: In pkg/planner/core/preprocess.go (pingcap/tidb PR `#65532`), the `checkLockClauseTables` function intentionally accepts `FOR UPDATE OF <base_table>` even when the FROM item was introduced with an alias (`matched.alias.L != "" && !matchedByAlias`). Rather than returning an error (strict MySQL behavior), TiDB emits a warning to guide users toward the alias form. This backward-compatibility fallback is a deliberate acceptance criterion for the change — do not flag it as a MySQL-incompatibility bug in future reviews.
Learnt from: guo-shaoge
Repo: pingcap/tidb PR: 67857
File: pkg/sessionctx/variable/session.go:2692-2702
Timestamp: 2026-04-18T01:35:27.792Z
Learning: In `pkg/sessionctx/variable/session.go` (pingcap/tidb PR `#67857`), `PrepareDedupCacheKey` intentionally omits `EnableWindowFunction`, `EnableStrictDoubleTypeCheck`, and `CharacterSetClient` from the dedup cache key. The author (guo-shaoge) is aware of this gap and has deferred fixing it to a later PR. Do not flag this omission as a new bug in future reviews of this file; it is a known, tracked TODO.
Learnt from: joechenrh
Repo: pingcap/tidb PR: 67054
File: pkg/executor/importer/kv_encode.go:89-96
Timestamp: 2026-03-18T01:55:56.861Z
Learning: In `pkg/executor/importer/kv_encode.go` (pingcap/tidb), `TableKVEncoder.initInsertColFileMapping` uses `intest.Assert` (not a returned error) to guard the invariant that `fieldMappings` column IDs match `insertColumns` in order. This invariant is established at construction time and is covered by the broader test suite; it cannot be violated in production. Suggesting a production `error` return for this check is a false positive — the `intest.Assert` pattern is intentional and preferred here.
Learnt from: joechenrh
Repo: pingcap/tidb PR: 66878
File: pkg/planner/core/logical_plan_builder.go:6455-6472
Timestamp: 2026-03-11T06:29:00.122Z
Learning: Ensure code reviews verify the VirtualAssignmentsOffset semantics: the planner sets Update.VirtualAssignmentsOffset = len(update.List). The executor should only apply OrderedList[:VirtualAssignmentsOffset] when composing new rows, and only after the 'changed' check should it group/evaluate OrderedList[VirtualAssignmentsOffset:] per table. This pattern applies to files under pkg/planner/core and pkg/executor (e.g., common_plans.go and update.go). Reviewers should check that updates respect slicing behavior, that the offset is consistently derived from the planner, and that downstream code does not bypass the offset when creating new rows. Add tests validating both branches: the slice before the offset for new rows, and the per-table handling of the slice after the offset.
|
/retest |
|
@ChangRui-Ryan: PRs from untrusted users cannot be marked as trusted with DetailsIn response to this:
Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. |
winoros
left a comment
winoros
left a comment
There was a problem hiding this comment.
lgtm for the planner part.
The newly added function will not return null if given args are (not null, not null, null).
[LGTM Timeline notifier]Timeline:
|
b0d7213 to
73af86f
Compare
ti-chi-bot
Bot
added
size/XXL
ChangRui-Ryan
changed the title
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@pkg/sessionctx/variable/sysvar_test.go`:
- Around line 1310-1321: The test assumes the process-global embed base is empty
but only restores it later; after saving originalBase from
vardef.EmbedOpenAIAPIBase in the subtest, explicitly store an empty string
(vardef.EmbedOpenAIAPIBase.Store("")) to force the default fallback, then
proceed to call GetSysVar(vardef.TiDBExpEmbedOpenAIAPIBase) and assert it equals
vardef.DefTiDBExpEmbedOpenAIAPIBase; keep the existing t.Cleanup that restores
originalBase so the global state is returned after the test (reference:
vardef.EmbedOpenAIAPIBase, GetSysVar, TiDBExpEmbedOpenAIAPIBase,
DefTiDBExpEmbedOpenAIAPIBase, vars).
In `@pkg/sessionctx/variable/sysvar.go`:
- Around line 3974-3980: The current whitelist lets any Azure host pass by
checking strings.HasSuffix(host, ".azure.com"); narrow this to only Azure
OpenAI-compatible endpoints by replacing that suffix check with explicit allowed
Azure OpenAI patterns (for example check strings.HasSuffix(host,
".openai.azure.com") and, if supported, strings.HasSuffix(host,
".cognitiveservices.azure.com")), using the same host variable derived from
parsedURL.Hostname() and keeping the existing error return using
openAIEmbeddingAPIBaseWhitelistErrMsg when none match.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 78dbbcbf-5901-489c-bee2-53b8970f851b
📒 Files selected for processing (19)
pkg/expression/BUILD.bazelpkg/expression/builtin.gopkg/expression/builtin_inference.gopkg/expression/builtin_threadunsafe_generated.gopkg/expression/function_traits.gopkg/expression/integration_test/BUILD.bazelpkg/expression/integration_test/integration_test.gopkg/expression/scalar_function.gopkg/inference/BUILD.bazelpkg/inference/manager.gopkg/inference/manager_test.gopkg/inference/mock.gopkg/inference/openai.gopkg/parser/ast/functions.gopkg/planner/util/null_misc_test.gopkg/sessionctx/vardef/tidb_vars.gopkg/sessionctx/variable/sysvar.gopkg/sessionctx/variable/sysvar_test.gotests/integrationtest/r/executor/show.result
✅ Files skipped from review due to trivial changes (10)
- pkg/expression/BUILD.bazel
- pkg/planner/util/null_misc_test.go
- tests/integrationtest/r/executor/show.result
- pkg/expression/integration_test/BUILD.bazel
- pkg/expression/scalar_function.go
- pkg/expression/builtin_threadunsafe_generated.go
- pkg/expression/function_traits.go
- pkg/parser/ast/functions.go
- pkg/inference/BUILD.bazel
- pkg/expression/builtin.go
73af86f to
43aaf0d
Compare
|
/retest |
|
@ChangRui-Ryan: PRs from untrusted users cannot be marked as trusted with DetailsIn response to this:
Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. |
43aaf0d to
d1aa399
Compare
d1aa399 to
fd97b56
Compare
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: winoros The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:
Approvers can indicate their approval by writing |
|
/retest |
|
@ChangRui-Ryan: PRs from untrusted users cannot be marked as trusted with DetailsIn response to this:
Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. |
|
@windtalker @bb7133 @D3Hunter PTAL, thanks |
2 participants
What problem does this PR solve?
Issue Number: ref #67765
Problem Summary:
TiDB already supports vector data types and vector-related workflows, but it does not provide a built-in SQL function to generate embeddings from text directly inside SQL. Users have to call external embedding services outside TiDB and then write vectors back manually, which makes SQL-native vector workflows less convenient and harder to integrate.
This PR adds end-to-end support for
embed_text, so TiDB can generate aVECTOR FLOAT32value directly from SQL by calling an OpenAI-compatible embedding service.What changed and how does it work?
This PR adds a new builtin function
embed_text(model, text[, options]).embed_textaccepts a provider-qualified model name in the<provider>/<model>format, for exampleopenai/text-embedding-3-small. The builtin parses the provider/model identifier, dispatches the request through the inference registry, calls the corresponding embedding provider, and converts the returned embedding intoVECTOR FLOAT32.The implementation includes:
pkg/inference;embed_text;tidb_exp_embed_openai_api_base;For the OpenAI-compatible provider, the current implementation supports provider-specific options such as
dimensionsanduser. The API base can be resolved from the sysvar first, and falls back to the environment configuration when needed.Check List
Tests
Side effects
Documentation
Release note
Please refer to Release Notes Language Style Guide to write a quality release note.
Summary by CodeRabbit
New Features
Behavioral Changes
Tests