feat(function): add onnx_run builtin to evaluate ONNX models from SQL

[fengttt]

What

Adds a builtin scalar function onnx_run that evaluates an ONNX model over each row and returns the output as JSON.

select onnx_run(model, input, input_shape, output_shape) from T;

• model — varbinary model bytes, or a datalink to a model file in a stage.
• input — JSON flat array for the input tensor.
• input_shape / output_shape — JSON '{"dim":[1,1,4],"dtype":"float32"}'. A NULL output_shape means the output is not a plain tensor: the model's outputs (sequence/map, e.g. sklearn ZipMap probabilities) are returned as a JSON object keyed by output name.

Example:

select onnx_run(
    cast('stage://models/mnist.onnx' as datalink),
    cast(json_array(...) as varchar),
    '{"dim":[1,1,28,28],"dtype":"float32"}',
    '{"dim":[1,10],"dtype":"float32"}');

How

• pkg/mlai/onnx isolates all onnxruntime_go usage. It is cgo-free — the onnxruntime shared library is dlopen'd at runtime. The environment is initialized lazily and at most once; if the library cannot be loaded, every onnx_run call fails cleanly while the database still starts normally.
• The onnxruntime shared library (CPU, v1.26.0) is fetched per-platform (linux x64/arm64, macOS arm64) by the thirdparties/Makefile onnxruntime target into thirdparties/install/lib, and resolved at runtime via $MO_ONNXRUNTIME_LIB, beside mo-service, or under thirdparties/install/lib. GPU is intentionally out of scope for now.
• The operator (func_builtin_onnx.go) caches one onnxruntime session per expression instance and reuses it across rows, closing it when the expression is released — the same newOpWithFree lifecycle as serial/serial_full.
• Supported dtypes: int8/16/32/64, uint variants, float32/64, bool, and float16 (via CustomDataTensor + IEEE-754 half conversion).

Tests

• Unit tests in pkg/mlai/onnx (shape/dtype parsing, half conversion, and real inference on the checked-in models).
• BVT test/distributed/cases/function/onnx_run — passes 13/13 — covering:
  • sum_and_difference (both the datalink and varbinary overloads),
  • non_tensor_outputs (sklearn random forest, NULL output_shape),
  • per-row inference over the checked-in iris dataset (session reuse + json_array; predictions match labels),
  • mnist and mnist_float16 (bonus),
  • error paths (bad dtype, shape/length mismatch, NULL model).
• Test models and the iris dataset are checked into test/distributed/resources/onnx/ (small files; the generating Python scripts are intentionally not included).

Notes

• CPU-only for now; the pkg/mlai/onnx.NewSession seam leaves room for a later GPU (CUDA execution provider) path behind a config flag.
• ONNX model outputs are numerically deterministic for the shipped x86_64 library; a different-architecture CI runner could differ in the last digits of float outputs — worth watching if BVT runs on arm64.

🤖 Generated with Claude Code

[qodo-code-review]

Qodo reviews are paused for this user.

Troubleshooting steps vary by plan Learn more →

On a Teams plan?
Reviews resume once this user has a paid seat and their Git account is linked in Qodo.
Link Git account →

Using GitHub Enterprise Server, GitLab Self-Managed, or Bitbucket Data Center?
These require an Enterprise plan - Contact us
Contact us →

[fengttt]

Review — focus: error handling & efficiency

Overall: sound architecture (lazy fail-safe init, serial-style session lifecycle, clean layer separation). Two real error-handling bugs, two worthwhile efficiency items, and one test-coverage gap that hides bug 2.

🔴 High — error handling

1. Pre-allocated output tensor leaks when Run fails — pkg/mlai/onnx/session.go

The cleanup defer is registered after the error return:

outputs[0] = outTensor                            // C-allocated
if err := s.s.Run(...); err != nil {
    return nil, moerr...                          // <- outputs[0] leaks
}
defer func() { for _, o := range outputs { o.Destroy() } }()  // too late

Every failed inference with a declared output_shape leaks one C-side tensor. Fix: register the defer immediately after outputs[0] = outTensor, before calling Run.

2. Integer input conversion silently swallows errors and truncates — pkg/mlai/onnx/tensor.go convInts/convUints

v, _ := strconv.ParseInt(string(num), 10, 64)   // error discarded
out[i] = T(v)                                    // unchecked narrowing

Two silent-wrong-answer paths for all 8 integer dtypes:

• [1.5, 2.5] with "dtype":"int32" → ParseInt fails → tensor becomes [0, 0]. No error, wrong inference.
• [300] with "dtype":"int8" → parses as 300, int8(300) wraps to 44. No error, wrong inference.

The float path correctly propagates via badNumber(); the int paths should too, with a bitSize per element type.

3. (Low) Trailing garbage after the input array is ignored — decodeNumbers

json.Decoder.Decode stops after the first value, so '[1,2,3,4]junk' is accepted.

🟡 Medium — efficiency

4. Triple JSON serialization per row

json.Marshal(result) (session.go) → bytejson.ParseFromByteSlice → bj.Marshal() (func_builtin_onnx.go). appendBinaryJSON accepts only nil/bool/int64/uint64/float64/json.Number/string/[]any/map[string]any, so the current []any (holding float32, int8, …) can't feed CreateByteJSON directly — the text round-trip is currently load-bearing. Fix: normalize scalars to float64/int64/uint64, return the value tree from Run, and CreateByteJSON it directly — one serialization pass instead of three. Meaningful for large outputs.

5. Per-row constant re-work in the eval loop — func_builtin_onnx.go

• onnx.ParseShape runs on every row for what is almost always a constant literal. Cache keyed by raw bytes.
• ensureSession's bytes.Equal(op.lastArg, rawArg) is O(model size) per row — trivial for datalink URLs, ~ms per row for a 22 MB varbinary model. Checking params[0].IsConst() once short-circuits both the compare and the copy.

6. (Nit) Model parsed twice per session build — NewSession

GetInputOutputInfoWithONNXData builds a temp ORT session, then NewDynamicAdvancedSessionWithONNXData builds the real one. One-time per distinct model, acceptable — worth a comment.

Test coverage

• Good: both overloads, non-tensor outputs, per-row session reuse, float16 both directions, error paths, deterministic BVT.
• Gap: no test exercises any integer dtype input — precisely where bug 2 lives.

Verdict

Fix 1 and 2 before merge (both are silent-failure bugs: a C-memory leak and wrong query results), plus the integer-dtype test. 4/5 are worthwhile in-PR follow-ups.

[fengttt]

Pushed 25de48f addressing all review findings:

• 1 (leak): cleanup defer now registered before Run — pre-allocated output tensor is destroyed on the error path.
• 2 (silent int truncation): integer parsing uses the declared bit width and propagates errors — 1.5 as int32 and 300 as int8 are now clean errors (new unit + BVT cases).
• 3 (trailing garbage): rejected.
• 4 (triple serialization): Session.Run returns a ByteJson-native value tree (float32 widened via shortest-repr so rendering is unchanged); operator uses CreateByteJSON directly — one encoding pass. NaN/Inf still error as before.
• 5 (per-row constant re-work): parsed shapes cached across rows; per-row bytes.Equal/copy skipped when the model vector is const.
• 6: double-parse comment added in NewSession.

Verification: all unit tests pass (8/8, incl. 3 new); BVT regenerated and passes 15/15; diff of the .result shows only the two added error cases — all pre-existing output is byte-identical, confirming the serialization change preserves rendering exactly.

[gouhongshen]

Codex automated review

The prior review findings are addressed, but four blocking resource/supply-chain issues remain.

P1 - Make inference observe query cancellation (pkg/mlai/onnx/session.go:124)

Run calls the synchronous s.s.Run API without RunOptions and Session.Run receives no context, so proc.Ctx cancellation never reaches an in-flight model. A model containing a very large Loop or otherwise expensive graph continues occupying the executor and ONNX worker threads after KILL, timeout, or client disconnect. Use RunWithOptions and terminate the run when the query context is canceled, with synchronization that cannot race session destruction.

P1 - Bound peak conversion memory rather than only raw tensor bytes (pkg/mlai/onnx/helpers.go:56)

MaxTensorBytes limits only elementCount*dtypeWidth. An accepted 64 MiB int8/bool tensor contains 67,108,864 elements, after which intsToAny/boolsToAny allocates a []any of that length—at least 1 GiB on 64-bit Go—before ByteJSON adds another large allocation. Input decoding has similar amplification through []json.Number. These allocations are outside the query mpool, so one valid call can still OOM the CN despite the new cap. The limit must account for the complete peak representation or the conversion must stream/encode directly with query-memory accounting.

P1 - Apply the tensor bound to auto-allocated outputs (pkg/mlai/onnx/session.go:124)

When output_shape is NULL, outputs consists entirely of nil slots. DynamicAdvancedSession.Run therefore asks ONNX Runtime to allocate every model output with no size validation or query-memory accounting; ParseShape/MaxTensorBytes is never involved. A small model can generate an arbitrarily large tensor or sequence and exhaust the CN before valueToJSON can inspect it. NULL output_shape is required for the advertised sequence/map path, so this needs an enforceable runtime/output-memory bound rather than a post-allocation length check.

P1 - Pin and verify the downloaded ONNX Runtime binary (thirdparties/Makefile:169)

Every clean build downloads a native shared library from the mutable matrixorigin/thirdparty main branch and installs it without checking a digest. The same MatrixOne commit can therefore silently produce different binaries, and a changed or compromised mirror artifact becomes native code loaded into every CN. Pin an immutable source revision and verify an architecture-specific SHA-256 before installing the library.

[fengttt]

@gouhongshen response to the second review round (fixes in c0172fd):

P1 — pin and verify the downloaded runtime: fixed. The Makefile now downloads from an immutable matrixorigin/thirdparty commit (6553650) instead of main, and verifies a per-platform SHA-256 before installing (download to temp → checksum → move; a tampered hash aborts the build with nothing installed — tested both directions). Same MatrixOne commit ⇒ same native library, byte for byte.

P1 — apply the tensor bound to auto-allocated outputs: fixed. anyTensorFlat (the NULL-output_shape conversion path) now applies the raw-tensor-byte cap to model-produced outputs before any conversion, since those never pass through ParseShape (single conservative check at MaxTensorBytes/8 elements; regression test added). Note the C-side allocation itself happens inside onnxruntime and cannot be pre-bounded through the Go binding — the guard stops the Go-side conversion and returns a clean SQL error.

P1 — query cancellation for in-flight inference: declined (maintainer decision). Not necessary for the small, fast CPU models this feature targets; the plumbing (RunOptions/Terminate watcher per row) isn't warranted at this point. Can be revisited if long-running models become a real use case.

P1 — bound peak conversion memory: declined (maintainer decision). The raw tensor byte caps (64 MB per tensor, on declared shapes, model outputs, and the 1 GiB datalink model cap) are the intended bound; conversion-layer accounting beyond that is not needed for tensors this size.

[gouhongshen]

Codex automated review

History assessment at c0172fdc5b: the runtime pin/checksum, volatility, and earlier error-handling findings are fixed. The new auto-output check limits Go-side conversion after inference, but it does not bound ONNX Runtime allocation or intermediate tensors. The author explicitly declined cancellation and conversion-memory changes as maintainer decisions; both remain blocking because the SQL API accepts arbitrary model bytes and the intended small-model usage is not enforced by a permission, model constraint, execution limit, or memory boundary.

P1 - Enforce a memory limit during ONNX execution (pkg/mlai/onnx/session.go:124)

Author response: The new anyTensorFlat element check is intended to cap model-produced outputs before conversion; the response also acknowledges that the C-side allocation occurs inside ONNX Runtime and cannot be pre-bounded through the current Go binding.

Why this remains blocking: That check runs only after DynamicAdvancedSession.Run returns. By then ONNX Runtime has already allocated the output, and the binding has copied automatic tensor data into Go memory. Model intermediates never pass through anyTensorFlat at all. An arbitrary uploaded model can therefore use Expand, ConstantOfShape, or a large intermediate followed by a small reduction to exhaust CN memory before the new check executes. The intended small-model scope is not enforced by the API.

To close this finding, enforce or query-account a memory limit during ORT execution, or introduce an enforced trust/model contract that rejects graphs capable of unbounded allocation before Run. A post-run output check alone cannot close the allocation path.

P1 - Propagate query cancellation into inference (pkg/sql/plan/function/func_builtin_onnx.go:221)

Author response: Cancellation plumbing was declined because this feature targets small, fast CPU models and a per-row RunOptions/termination watcher was considered unwarranted.

Why this remains blocking: The implemented SQL function accepts arbitrary model bytes and does not enforce a complexity, operator, duration, or privilege boundary corresponding to that intended usage. proc.Ctx is discarded before synchronous ORT execution, so a model containing a large Loop can continue consuming executor and ORT worker resources after query timeout, KILL, or client disconnect.

To close this finding, propagate cancellation through RunWithOptions/RunOptions.Terminate with lifecycle synchronization, or enforce a concrete model/execution boundary that makes long-running inference unreachable.

P1 - Bound conversion peak rather than only raw tensor bytes (pkg/mlai/onnx/helpers.go:56)

Author response: Conversion-memory work was declined because the 64 MiB raw-tensor cap is the intended bound and was considered sufficient for supported tensor sizes.

Why this remains blocking: Raw bytes are not the peak representation. A valid 64 MiB bool or int8 output can contain 67,108,864 elements; converting it to []any allocates roughly 1 GiB for the interface slice alone on a 64-bit CN, before reshape structures, ByteJSON construction, serialization copies, and the final mpool-backed result. The new MaxTensorBytes/8 check protects only the NULL-output_shape path; declared narrow outputs still reach this amplification path.

To close this finding, cap element count according to peak conversion/serialization cost, encode without the []any expansion, or account the complete conversion against query memory.

[fengttt]

UT failure on the last run was a data race in pkg/frontend (Session.reset writing through a shared *time.Location at session.go:2017, racing CDC log timestamp formatting) — untouched by this PR; filed as #26122 with full analysis. Branch updated to current main; CI re-running.

[XuPeng-SH]

Deep review of exact head 325242048126e1f93914bdc243956d69251ba0bf.

The normal execution/lifecycle paths look solid, and I verified the ONNX package, its race suite, the focused SQL-function race suite, the full SQL-function package, build, and vet with the pinned ORT 1.26 runtime. However, three P1 resource-control paths remain open, so I cannot approve this SQL surface yet.

1. Bound memory during ORT execution, not after it (pkg/mlai/onnx/session.go:124, pkg/mlai/onnx/helpers.go:135). The NULL-shape guard runs only after DynamicAdvancedSession.Run returns. In the exact onnxruntime_go v1.31.0 dependency, auto-allocated outputs are also copied from ORT memory into Go tensors before RunWithOptions returns. A model can therefore allocate an oversized output, or a large intermediate followed by a small output, before this check is reachable. Since arbitrary model bytes are accepted and there is no trust/privilege/operator-complexity boundary, this remains a CN OOM path. Please enforce/account a runtime memory limit, or add an enforced trust/model contract that makes unbounded graphs unreachable.

2. Propagate query cancellation into inference (pkg/sql/plan/function/func_builtin_onnx.go:221). proc.Ctx is not passed to Session.Run; a model with a large Loop can keep the executor and ORT workers busy after query timeout, KILL, or disconnect. The pinned binding already exposes RunWithOptions plus RunOptions.Terminate. Please add lifecycle-safe cancellation, or enforce a concrete execution/model boundary that makes long-running inference unreachable.

3. Cap peak conversion cost for declared narrow outputs (pkg/mlai/onnx/output.go:25, pkg/mlai/onnx/helpers.go:55). ParseShape permits 64 MiB of int8/bool, i.e. 67,108,864 elements. flatData expands that to []any, whose interface backing array alone is roughly 1 GiB on a 64-bit CN, before nested reshape slices, ByteJSON construction, marshal copying, and the mpool-backed result. The MaxTensorBytes/8 element guard applies only to NULL output_shape, so the declared-output path still allows this amplification. Please cap by complete conversion/serialization peak, stream/encode without []any, or account it against query memory.

The intended “small, fast CPU models” scope is reasonable, but it must be enforced in code before it can serve as the safety boundary for a generally callable SQL function.

[gouhongshen]

Codex automated review

The clean merge since c017 adds no authored ONNX delta. Fixed: tensor cleanup, integer/trailing-JSON validation, declared-shape and datalink pre-read bounds, volatility, and runtime pin/checksum. Withdrawn: none. No previously blocking risk is demonstrably non-blocking: the author declined execution-memory, cancellation, and conversion-peak controls, and all remain reachable.

P1 - ONNX Runtime allocations remain unbounded during Run (pkg/mlai/onnx/session.go:124)

Author response: the post-run anyTensorFlat check was presented as the auto-output bound, while acknowledging that ORT allocates before that check; the intended small, fast CPU-model scope was used to decline further controls.

Why this remains blocking: with NULL output_shape, outputs contains nil slots and line 124 calls DynamicAdvancedSession.Run. The pinned binding first lets ORT allocate output values, then copies non-string tensors into a Go slice before Run returns; only afterwards does this code reach valueToJSON/anyTensorFlat. Thus an Expand/ConstantOfShape output can exhaust native and Go memory before the new check executes. A large intermediate followed by a small result never reaches that check at all. Small arbitrary ONNX models are accepted through the varbinary overload, with no enforced model-complexity, privilege, or execution-memory boundary. Enforce/account a memory limit during ORT execution, or enforce a model trust/validation contract that rules out such allocations.

P1 - Query cancellation never reaches ONNX execution (pkg/sql/plan/function/func_builtin_onnx.go:221)

Author response: cancellation plumbing was declined because the feature is intended for small, fast CPU models.

Why this remains blocking: line 221 invokes Session.Run without proc.Ctx; Session.Run then invokes the synchronous s.s.Run API. A query timeout, KILL, or disconnect therefore cannot interrupt a valid model containing a large Loop, leaving the executor and ORT workers occupied. The API accepts arbitrary model bytes and enforces no duration, operator, or permission boundary corresponding to the intended usage. The pinned binding exposes RunWithOptions and RunOptions.Terminate. Propagate proc.Ctx cancellation through those options, with synchronization against Session.Close/Reset, or add an enforced execution/model boundary.

P1 - 64 MiB narrow outputs expand to multi-GiB Go allocations (pkg/mlai/onnx/helpers.go:55)

Author response: the author declined conversion-memory accounting, treating the 64 MiB raw-tensor cap as sufficient.

Why this remains blocking: ParseShape deliberately accepts 67,108,864 int8 elements at the cap. For a declared int8/bool output of that shape, tensorToNested reaches intsToAny/boolsToAny, whose make([]any, len(s)) allocates at least 1 GiB for the interface backing array on this 64-bit target before reshaping, ByteJSON construction, marshaling, and the result vector. These allocations are on the Go heap, not proc.Mp; ByteJSON then allocates its own per-element representation. A small model can Expand a scalar to this output, so the input need not be large. Cap element count by complete conversion/result cost, encode directly with query-memory accounting, or otherwise enforce a peak-memory bound.