diff --git a/docs/design/sirius-streamed-input.md b/docs/design/sirius-streamed-input.md new file mode 100644 index 0000000000000..2b6117faec354 --- /dev/null +++ b/docs/design/sirius-streamed-input.md @@ -0,0 +1,670 @@ +# Sirius streamed MatrixOne input protocol + +Status: proposed for design approval + +Design version: 3 + +Approval: pending distinct reviewer approval + +Owner: MatrixOne query execution + +Owning issue: [#27586](https://github.com/matrixorigin/matrixone/issues/27586) + +Implementation: MatrixOne [#27599](https://github.com/matrixorigin/matrixone/pull/27599), Sirius [#6](https://github.com/matrixorigin/sirius/pull/6), sidecar [#14](https://github.com/matrixorigin/mo-sirius-sidecar/pull/14) + +## 1. Decision + +MatrixOne may execute an eligible table scan at its transaction snapshot and +stream the resulting MatrixOne-native batches to a separately running Sirius +sidecar. Sirius consumes those batches through a GPU-native scan operator and +returns MatrixOne-native result batches. Arrow Flight supplies authenticated +RPC framing only; Arrow record batches are not the data representation in +either direction. + +This is an explicit, single-CN compatibility path selected by +`/*+ SIDECAR STREAM */`. It does not replace or alter direct `TaeRead`, and it +does not silently fall back to MatrixOne execution after the explicit stream +mode has been selected. + +The protocol is version 6. It requires an exact capability-document match +across MatrixOne, the sidecar, and Sirius. Mixed protocol revisions fail before +result rows are exposed. + +After approval, any semantic change to this document increments the design +version and requires fresh approval. Editorial corrections may retain the +version only when they do not change a protocol, ownership, resource, rollout, +or acceptance contract. + +## 2. Problem and invariants + +Direct `TaeRead` is efficient when the sidecar can read every object needed for +an admitted snapshot. It cannot cover all MatrixOne-visible states, including +unflushed committed rows, visible tombstones, or storage not reachable by the +sidecar. The streamed path keeps snapshot and storage semantics in MatrixOne +while retaining Sirius execution for the remainder of the plan. + +The first implementation acknowledged each staged batch and immediately +self-scheduled another `GPU_MO_SCAN` task. That bounded the Flight slot but not +the downstream Sirius repositories or task queue. Version 2 removed that eager +source continuation, but SF10 disproved its claimed end-to-end bound: a `FULL` +partition barrier legitimately asks its producer to finish and retained 2.71 +GiB of source-derived host data for Q9. Once the barrier opened, two configured +GPU workers launched `cudf::hash_partition` concurrently on one GPU. Repeated +runs either returned different Q9 values or failed at that overlap with +`cudaErrorInvalidDevice` followed by `cudaErrorIllegalAddress`; the same plan +with one GPU worker was byte-identical to native MatrixOne. Q1 and Q6 remained +byte-identical with two workers, and reduced Q9 fingerprints remained exact +through `part`, `lineitem`, `partsupp`, and `orders`, isolating the failure to +the large concurrent partition phase rather than Flight or the native codec. + +Version 3 keeps ordinary GPU task parallelism but admits only one `PARTITION` +execution per GPU until that operator's CUDA stream is synchronized. It also +replaces the false one-published-batch memory claim with a process-global, +fail-closed 8 GiB StreamRead host-memory budget shared by every execution. +Full barriers may retain multiple source-derived batches, but every live input +frame, staging copy, and retained native host representation holds a lease from +that budget. + +The primary correctness invariant is: + +> For every `StreamRead`, Sirius consumes exactly the projected rows and schema +> produced by the MatrixOne scan at the statement snapshot, once and only once, +> without cross-query reuse or resources surviving terminal execution. + +The supporting invariants are: + +1. MatrixOne is the sole owner of transaction visibility, native scan, + MVCC/tombstone application, scan filtering, and scan projection. +2. A stream identity is bound to one account, query, snapshot, schema, + capability set, and expiry. It is single-use within one execution ticket. +3. Every input frame is acknowledged only after Sirius has copied the bytes it + needs and the sidecar no longer retains that frame as the current input. +4. Each `StreamRead` has at most one active Sirius source task; the source never + self-schedules and advances only when downstream asks for more input. A + blocking downstream barrier may retain multiple completed source batches, + each charged to the global StreamRead host budget. +5. Input, result, plan, ticket, and execution counts have hard bounds. +6. Cancellation can interrupt a blocked input acknowledgement, blocked result + receive, and Sirius worker independently of the data path. +7. MatrixOne does not release snapshot/query resources until local producers, + the sidecar execution, and all sidecar input handlers are quiescent. +8. At most one GPU `PARTITION` operator executes per physical GPU. The permit is + held through stream synchronization and is released on every terminal path; + non-partition GPU operators retain their configured concurrency. +9. Direct `TaeRead` keeps its existing schema, physical-type, lease, and + fallback contract. It shares the core per-GPU partition safety invariant. + +The negation of the contract includes lost or duplicate batches, using a stream +from another account/query/snapshot, acknowledging a batch before ownership is +transferred, unbounded buffering under a slow consumer, releasing resources +before quiescence, eagerly scheduling the next Sirius source batch without +downstream demand, executing on DuckDB CPU, or accepting a different wire ABI. + +## 3. Scope + +### 3.1 Included + +- Explicit `SIDECAR STREAM` execution for read-only `SELECT` statements. +- One MatrixOne CN and its paired sidecar per execution. +- Up to 16 independently named `StreamRead` inputs in one Substrait plan. +- Native MatrixOne scan filters, projections, offsets, and limits below the + stream boundary. +- Sirius GPU execution of admitted Substrait joins, filters, projections, + aggregates, sorts, fetches, and references. +- MatrixOne-native input and result batches over mutually authenticated Flight. +- Success, prepare failure, producer failure, consumer failure, cancellation, + timeout, disconnect, sidecar shutdown, and result-side early completion. +- Bounded TPC-H execution where all live StreamRead input memory fits the + configured process-global budget. Budget exhaustion is an explicit terminal + resource error, never fallback or unbounded waiting. + +### 3.2 Excluded + +- Multi-CN producer fan-in or a remote/distributed scan scheduler. +- Current-transaction writes or a transaction workspace with prior writes. +- Replaying one physical scan node through multiple `ReferenceRel` consumers. +- Transparent replacement of direct `TaeRead`. +- A general remote-execution framework. +- Unbounded blocking plans, spill-to-disk for streamed native input, or a claim + that one source batch bounds all downstream operator state. +- A stable public MatrixOne batch ABI across arbitrary MatrixOne releases. +- Production enablement without the rollout and acceptance gates in this + document. + +## 4. Alternatives + +### 4.1 Direct `TaeRead` + +The sidecar reads flushed TAE objects directly. This avoids CN scan and network +serialization and remains the preferred analytical fast path. It cannot +represent every MatrixOne-visible state and requires storage accessibility plus +GC-safe object leases. The streamed path therefore complements rather than +replaces it. + +### 4.2 Arrow record batches over Flight + +This was the original experiment in #27586. Arrow is interoperable and already +fits Flight, but MatrixOne would need an additional conversion and dependency at +both input and result boundaries. Sirius would then convert Arrow/DuckDB chunks +again for its GPU-native path. That duplicates type mapping, allocations, and +copies in the hot path. + +### 4.3 MatrixOne-native batches over Flight (selected) + +MatrixOne already owns the source batches and result consumer. Reusing +`Batch.MarshalBinary` avoids Arrow-Go in MatrixOne and lets Sirius reuse its +GPU-native TAE vector decoder. The cost is a deliberately strict, same-release +wire ABI. Exact capability negotiation, codec versioning, endian/size markers, +canonical decoding, and coordinated rollout contain that cost. + +### 4.4 One GPU worker for the whole query + +Setting `executor.pipeline.num_threads` to one made SF10 Q9 deterministic, but +it serializes joins, aggregates, projections, and scans that did not violate the +contract. It also changes deployment behavior for direct `TaeRead`. Version 3 +therefore models only `PARTITION` as an exclusive per-GPU resource and preserves +ordinary task concurrency. + +### 4.5 CUDA retry after partition failure + +Retrying launch-resource errors can be appropriate before any kernel mutates +output, but the observed failures included invalid-device and illegal-address +errors and successful concurrent runs produced nondeterministic values. Those +states are not safe to replay. Version 3 prevents the unsupported overlap and +treats every CUDA execution error as terminal. + +### 4.6 Immediate migration to current upstream Sirius + +Current upstream Sirius has newer partial-barrier, adaptive-join, and task +admission machinery, but the MatrixOne TAE/Substrait stack is based on a fork +hundreds of upstream commits behind it. Porting that stack is the preferred +long-term route to general streaming, but it is a separate migration with a +larger compatibility and validation surface. Version 3 deliberately stabilizes +bounded TPC-H on the current fork and records unbounded execution as a non-goal. + +## 5. Architecture and ownership + +The data path is: + +```text +MO planner + -> export Substrait with StreamRead leaves + -> prepare one sidecar execution and validate result schema + -> attach every DoPut input + -> start local-CN snapshot scan producers + -> MOB1 / MO Batch.MarshalBinary frames + -> sidecar one-slot StreamRead sources + -> Sirius GPU_MO_SCAN and GPU physical plan + -> MOB1 / canonical flat MO result frames over DoGet + -> MO result decoder and existing result writer +``` + +MatrixOne owns the statement transaction, scan pipeline, source batches, +producer goroutine, Flight client, decoded result batches, and final MySQL +write. The sidecar owns Flight admission, ticket/idempotency registry, +query-local input registry, Sirius connection/transaction, result frame slot, +execution worker, and input-handler join. Sirius owns the physical GPU plan, +task scheduling, GPU-native source conversion, pipeline data, and result packer. + +Ownership transfer for one input frame is: + +```text +MO scan batch + -> MO marshalled payload + -> Flight DoPut frame + -> sidecar one-slot native_batch_view + -> Sirius staged host representation + -> consumed acknowledgement + -> MO releases payload and scan batch +``` + +The acknowledgement is the linearization point. Before it, MatrixOne retains +the frame's source lifetime. After it, Sirius owns a copy and MatrixOne may +produce the next frame into the sidecar's one-slot input. The acknowledgement +does not admit another Sirius source task: `GPU_MO_SCAN` advances only after the +current published batch has driven downstream demand for more input. + +## 6. Admission and plan contract + +Stream mode is admitted only when all of these conditions hold: + +- the statement is an explicit streamed `SELECT`, not internal or prepared + execution; +- the transaction workspace is read-only and has no current or snapshot write + offset; +- every exported scan has one occurrence and an admitted physical input type; +- the plan has 1 through 16 stream inputs; +- all scan scopes remain on the current CN; +- every `StreamRead` schema is the deterministic post-scan native schema; +- the Substrait plan and result schema pass the exact capability contract. + +MatrixOne retains semantic operators above the scan in Substrait. It clears +native scan aggregation because the semantic aggregate remains in Sirius. It +keeps native scan filtering, projection, offset, and limit at the source. + +`CHAR` and `VARCHAR` are represented as Substrait `VarChar` while the physical +input vector keeps its original MatrixOne OID, width, charset, and nullability. +Sirius accepts both physical OIDs as the same string family. This is a protocol +mapping, not a catalog or planner rewrite. + +## 7. Wire protocol + +### 7.1 Transport and authentication + +All calls use Arrow Flight over gRPC with TLS 1.2 or newer. MatrixOne verifies +the sidecar server certificate and presents the configured CN client +certificate. The sidecar requires a trusted client certificate. Endpoint +redirection is rejected, so the authenticated connection is the only data +channel. + +The streamed path does not call the direct-read HTTPS resolver and exposes no +object path or storage credential. The direct `TaeRead` resolver and its +separate sidecar client identity remain unchanged. + +### 7.2 Capability negotiation + +At connection initialization, MatrixOne calls `GetCapabilities`. The returned +document must be byte-for-byte equal to MatrixOne's document. Its SHA-256 hash +is repeated in `ExecuteSubstraitRequest`, every `StreamRead`, and +`FlightInfo.app_metadata`. + +The version-6 capability fixes these values: + +- Substrait 0.78.0; +- `StreamRead` version 1 with feature bits 0; +- native batch frame/codec version 1; +- native result schema version 1; +- little endian, 16-byte MatrixOne type records, and 24-byte varlena records; +- at most 16 stream inputs and one buffered input slot per read; +- at most 4 MiB per input batch payload; +- StreamRead host accounting contract `global-fail-closed-v1`; +- at most 16 MiB per Substrait plan; +- the exact operator, expression, function, join, and type allow-lists. + +Any change to these semantics requires a new protocol or feature bit and a new +capability document. Unknown fields, enums, flags, or feature bits are rejected. + +### 7.3 Prepare + +MatrixOne sends `ExecuteSubstraitRequest` through `GetFlightInfo`: + +| Field | Contract | +| --- | --- | +| `protocol_version` | exactly 6 | +| `substrait_version` | exactly `0.78.0` | +| `capability_hash` | exact negotiated SHA-256 | +| `max_batch_bytes` | non-zero and no more than the sidecar result limit | +| `max_input_batch_bytes` | non-zero and no more than 4 MiB or the sidecar limit | +| `deadline_unix_ms` | future time, capped by the sidecar ticket TTL | +| `plan` | non-empty, at most 16 MiB | +| `query_id` | 16-byte statement identity | +| `account_id` | MatrixOne tenant identity | +| `idempotency_key` | SHA-256 of little-endian account ID followed by query ID | +| `result_schema` | canonical native result schema v1 | + +The sidecar fingerprints the whole request. Reusing an idempotency key with a +different fingerprint is terminal. A matching unclaimed prepare may return the +same ticket. Tickets are 32 random bytes and single-use. `FlightInfo` must have +one local endpoint, no redirection, the exact result schema, and the capability +hash. + +The prepare response is the last point at which the non-stream direct mode may +choose native fallback. Explicit `SIDECAR STREAM` is strict: prepare or +admission failure is returned to the statement and is never hidden by native +execution. No result row or schema is exposed before prepare validation. + +### 7.4 `StreamRead` + +Each Substrait extension read contains: + +| Field | Contract | +| --- | --- | +| protocol/feature bits | version 1, bits 0 | +| `stream_ref` | 32 cryptographically random bytes, unique in the plan | +| `query_id` | the 16-byte prepare identity | +| `account_id` | equal to the prepare account | +| `snapshot_ts` | 12-byte MatrixOne statement snapshot | +| `schema_digest` | SHA-256 of the deterministic `NamedStruct` | +| `capability_hash` | exact negotiated hash | +| expiry | future and within MatrixOne's signed timestamp range | + +Sirius and the sidecar reject identity, schema, capability, expiry, duplicate, +and unknown-field mismatches before binding the query-local `mo_stream_scan` +view. + +### 7.5 Native input `DoPut` + +Every input is attached before `DoGet` starts. The command descriptor contains +the ticket and `stream_ref`. The server first returns a `ready` acknowledgement +with zero counters. Data frames use `FlightData.app_metadata`; Arrow data bodies +and schemas are forbidden on this stream. + +The native envelope is: + +```text +offset size value +0 4 "MOB1" +4 2 little-endian codec version 1 +6 2 zero flags/reserved +8 8 non-zero contiguous sequence +16 8 payload byte length +24 N canonical MatrixOne Batch.MarshalBinary payload +``` + +One payload is at most the negotiated input limit and never more than 4 MiB. +MatrixOne splits a larger scan batch at row boundaries; a single row larger +than the limit is rejected. The sidecar validates vector count, physical types, +row lengths, attributes, null maps, areas, metadata, sequence, and trailing +bytes before publishing the frame. + +The server has one input slot per read. It acknowledges a frame only after the +Sirius scan copies all required bytes into bounded staging. Sirius has at most +one active source task per read and does not call `next_batch` again until its +existing downstream task-demand recursion reaches that source. A downstream +`FULL` barrier may retain completed source batches after their source tasks +finish; those batches remain charged to the global StreamRead host budget until +conversion or destruction. A slow consumer may leave one unacknowledged frame +in the sidecar slot and retained batches downstream; the following `DoPut` +blocks or fails with resource exhaustion before memory can exceed the budget. +The cumulative acknowledged batch, row, and byte counters must exactly match +both endpoints. + +Producer EOF receives a final `complete` acknowledgement only after consumer +EOF. If the plan prunes a read or result completion makes further input +unnecessary, the sidecar returns `complete + not_needed` with exact counters. + +### 7.6 Native result `DoGet` + +`DoGet` claims the ticket once. Flight emits its mandatory empty Arrow transport +schema first. Every later `FlightData.data_header` contains one `MOB1` result +frame; data bodies, metadata, descriptors, and redirection are forbidden. + +The sidecar splits GPU result batches at row boundaries so each payload is at +most `max_batch_bytes`. Result vectors are canonical flat MatrixOne vectors. +MatrixOne remarshal-checks the batch, rejects trailing or non-canonical bytes, +rejects non-flat/constant/dictionary result vectors, validates every physical +type and nullability field against the negotiated result schema, then gives the +batch to the existing result writer. It does not request the next frame until +the writer returns and the current batch is released. + +## 8. Type compatibility + +The first version deliberately supports the TPC-H family: + +| Family | Input | Substrait/Sirius | Result | +| --- | --- | --- | --- | +| boolean | MatrixOne `bool` | `bool` | exact `bool` | +| signed integers | `int8/16/32/64` | `i8/16/32/64` | exact signed OID | +| unsigned extract | not admitted as a streamed physical input | semantic `i64` | checked conversion to MO `uint32` | +| floating point | `float32/64` | `fp32/64` | exact floating OID | +| strings | physical `char` or `varchar` | `VarChar` | negotiated `char` or `varchar` | +| decimal | MO decimal64 for precision up to 18; decimal128 for 19 through 38 | Substrait decimal | exact width/scale and checked result conversion | +| date | MO date epoch | Substrait Unix-day date | checked conversion to MO date | + +Timestamp, binary/blob, JSON, arrays, UUID, enum, row-id, and any unlisted +physical type are rejected. Width, scale, charset, nullability, vector class, +and fixed physical size are part of validation, not advisory metadata. + +## 9. State machines and terminal ownership + +### 9.1 MatrixOne execution + +| From | Event / linearization | To | Required side effects | +| --- | --- | --- | --- | +| native | stream admission succeeds | admitted | bind snapshot and create stream identities | +| admitted | exact prepare/schema succeeds | prepared | own ticket and cancellation identity | +| prepared | every input returns `ready` | attached | no result worker has started yet | +| attached | `DoGet` claims ticket | running | start result worker and local scan producer | +| running | result EOF | result-complete | retire/interrupt all producers, join producer | +| any prepared/running | error, cancel, timeout, panic, shutdown | cancelling | abort inputs, interrupt `DoGet`, send `CancelExecution` | +| result-complete/cancelling | sidecar says `quiesced` or `not-found` | quiesced | sidecar worker and handlers are no longer owners | +| quiesced | local producers joined and cleanup succeeds | terminal | release query/snapshot resources | + +The result EOF is authoritative success. It retires an input even if the sidecar +pruned it before its first data frame. `Retire` has a cancellation path +independent of the input mutex so it can interrupt a blocked attachment or +acknowledgement. + +If prepare may have succeeded but its ticket is unavailable, MatrixOne cancels +by idempotency key. If quiescence or release cannot be proved, the execution is +retained by the reconciliation owner, which retries with bounded exponential +delay. Snapshot resources are not released before quiescence. + +### 9.2 Sidecar execution + +| State | Accepted events | Terminal transition | +| --- | --- | --- | +| preparing | same-idempotency replay or cancellation identity | publish one ticket or cancel preparation | +| prepared/unclaimed | attach inputs, `DoGet` claim, cancel, deadline | cancellation releases resolutions without starting work | +| claimed/running | input frames/EOF, result reads, cancel, deadline | worker records success/failure/cancel and becomes quiescent | +| quiescent | input-handler detach | remove ticket/idempotency record after handler count reaches zero | + +`CancelExecution` first publishes cancellation to every input, then interrupts +Sirius and the DuckDB connection, then waits for both worker quiescence and zero +active `DoPut` handlers. Multiple cancellation callers share one serialized +worker join. Deadline reaping and server shutdown use the same terminal path. + +### 9.3 Sirius execution + +Sirius permits `PREPARED -> RUNNING -> SUCCEEDED|FAILED|CANCELLED` and +`PREPARED -> CANCELLED`. The transition to `RUNNING` is single-use. All terminal +paths release query-local resolved views and stream-source references. The +actual backend is marked `SIRIUS_GPU` immediately before execution; successful +completion without a backend mark is invalid. + +For each `StreamRead`, Sirius permits `idle -> active -> published -> idle` and +`idle|active|published -> exhausted|cancelled|failed`. An atomic claim grants +one source task the `active` state. Successful publication releases the claim +without scheduling a continuation. Only downstream task demand may claim the +next generation. EOF marks `exhausted`; failure and cancellation admit no new +generation. + +Each physical GPU owns one partition permit. A task whose remaining operator +chain contains `PARTITION` must acquire that permit before reserving GPU memory +or launching work. The permit remains owned until the partition operator's CUDA +stream is synchronized. Waiting tasks are FIFO and hold no GPU reservation; +query cancellation removes them from the wait queue. Release admits exactly +one waiter and occurs through one RAII owner on success, exception, +cancellation, drain, and executor shutdown. Other GPU task classes do not use +this permit. + +Pipeline input claim and `tasks_created` publication are one atomic transition +under the pipeline status mutex. Completion checks use the same mutex and +notify parent pipelines only after unlocking, so an empty input repository +cannot be mistaken for completion while a task is being constructed. + +## 10. Q1-Q3 resource closure + +### Q1: destruction ownership + +| Resource | Creator | Effective terminal owner | +| --- | --- | --- | +| MO stream identity and scan scope | MO compile | compile release after producer join | +| Flight prepare/ticket/idempotency | sidecar registry | terminal callback after worker and handlers quiesce | +| input frame slot | sidecar `stream_input` | consume/not-needed/cancel path | +| StreamRead host-budget lease | sidecar frame / Sirius scan task | transferred with bytes; frame, staging, or representation destruction releases it | +| Sirius staged/pinned host data | GPU scan task | task/representation destruction releases data and its lease | +| per-GPU partition permit | GPU executor | one task-scoped RAII permit; release or query drain admits next generation | +| GPU pipeline data | Sirius repositories | existing pipeline/memory manager cleanup | +| result frame slot | sidecar execution entry | `DoGet` read or cancellation | +| decoded MO result batch | MO result loop | per-frame deferred clean | +| query-local DuckDB views/transaction | sidecar execution entry | resolution destruction and connection rollback | + +### Q2: wait-for closure + +The possible waits are partition-permit admission, input-slot publication, +input consumption, producer EOF, result-frame publication, result receive, +worker join, handler join, and cleanup reconciliation. Every data wait observes +cancellation/not-needed/deadline. A partition waiter owns no GPU reservation and +is removed by query cancellation or executor drain. +Cancellation does not take an input's data mutex before cancelling its gRPC +context. Sidecar cancellation wakes input and result condition variables and +interrupts Sirius/DuckDB. All RPCs are bounded by the minimum of caller, +request, lease-safe, and sidecar ticket deadlines. + +### Q3: accumulation bounds + +| Accumulation | Hard bound | +| --- | --- | +| plan | 16 MiB | +| stream inputs per execution | 16 | +| sidecar slot per input | one frame, at most 4 MiB payload, charged to the global budget | +| Sirius expanded input batch | 64 MiB | +| Sirius staged input per active read task | 96 MiB encoded plus at most 96 MiB pinned copy during transfer, both charged while live | +| Sirius source output | 32 MiB target and 96 MiB hard maximum per batch; any number retained by a barrier remains charged | +| all live StreamRead input memory in one sidecar process | configurable global budget, default 8 GiB; non-blocking admission fails closed | +| concurrent partition execution | one task per physical GPU; FIFO waiters retain no GPU reservation | +| result slot | one frame, at most negotiated result limit (default 64 MiB) | +| result schema | 1 MiB and 4096 columns | +| active tickets | configured limit, default 128 | +| reconciliation retry | one worker per retained execution, delay capped at 5 seconds, stopped by runtime close | + +The 4 MiB wire limit is not a memory-reservation estimate. The global budget +charges each actual live allocation before it is created: the sidecar frame, +per-column staging strings, the contiguous host representation, and any overlap +during ownership transfer. Leases are move-only and transfer with ownership; +replacement or destruction releases the corresponding bytes exactly once. +The budget is process-global, so concurrent executions compete for the same +fixed capacity rather than multiplying a per-query allowance. Admission never +waits because a `FULL` barrier could otherwise deadlock while waiting for memory +that only barrier completion can release. It terminates the execution with a +resource-exhausted error and normal cancellation/quiescence cleanup. + +The packaged SF10 configuration uses a 10 GiB Sirius host capacity and an 8 +GiB StreamRead budget, leaving 2 GiB for non-stream host work. Measured SF10 Q9 +input is 2,708,678,611 bytes and fits with headroom; SF10 Q1 transfers +5,046,665,324 bytes cumulatively but continuously releases batches, proving why +the bound applies to retained live memory rather than total traffic. Acceptance +records configured capacity, peak charged bytes, rejected admissions, and a +zero terminal balance. + +This is the stream-specific host envelope, not the complete Sirius GPU budget. +GPU operators remain subject to Sirius's independent reservation and usage +limits. The current allocator may eagerly create a pool at the configured GPU +usage limit during startup, so deployment must choose an absolute or fractional +limit that fits alongside the driver, sidecar, and concurrent query working +sets. Acceptance records configured, startup-reserved, peak-reserved, and +peak-used GPU bytes; a successful small query is not evidence that an SF10 join +or aggregate fits. + +## 11. Security and failure containment + +- Flight mTLS authenticates the paired CN; certificates and private keys never + appear in plans, tickets, logs, or artifacts. +- Account and query identities must agree in prepare and every `StreamRead`. +- Stream/ticket identities are random and single-use; idempotency identities are + deterministic only within the authenticated account/query pair. +- Capability and schema hashes prevent cross-version or cross-schema reuse. +- Unknown fields, malformed frames, overflow, non-canonical batches, oversized + rows, unsupported vector classes, and endpoint redirects are terminal. +- Stream mode exposes no TAE path, manifest, object credential, or resolver + endpoint. +- One sidecar is paired with one local CN. A sidecar is not shared across CNs. +- Failure is contained to the statement and its query-local sidecar entry. + +## 12. Compatibility, rollout, and rollback + +The MatrixOne native batch representation is an internal ABI. Protocol v6 is +therefore supported only for the exact capability document and pinned +MatrixOne/Sirius/sidecar revisions validated together. It is not a promise that +an arbitrary older or newer MatrixOne batch codec is compatible. + +Rollout order is: + +1. merge and publish merge-ready Sirius and sidecar revisions; +2. deploy the paired sidecar and verify capability negotiation while stream mode + remains unused; +3. deploy MatrixOne with stream mode disabled by default; +4. run the acceptance matrix on one local CN/sidecar pair; +5. permit explicit `SIDECAR STREAM` use only after approval. + +Mixed CN revisions are allowed only because each CN negotiates with its own +sidecar and the feature is explicit. A mismatched pair rejects negotiation. A +rolling upgrade must update a pair as one unit before enabling stream mode. + +Rollback disables explicit stream use and restores the previously pinned +sidecar image. Direct `TaeRead` and native MatrixOne execution remain available +and unchanged. Protocol or schema mismatch never triggers post-visibility +fallback. + +## 13. Observability + +Acceptance and production diagnostics must identify one query without exposing +secrets: + +- protocol/capability revision and outcome; +- prepare, first-input, first-result, quiescence, and cleanup durations; +- input batches/rows/payload bytes and result batches/rows/payload bytes; +- GPU backend evidence showing `GPU_MO_SCAN` and `SIRIUS_GPU` actually started; +- cancellation source and terminal outcome; +- active tickets, active input handlers, retained reconciliation owners; +- StreamRead host-budget capacity, current charge, peak charge, rejected bytes, + and terminal zero-balance result; +- per-GPU active and queued partition tasks, including the observed maximum + concurrent partition count; +- MatrixOne allocation-account terminal snapshot; +- sidecar host peak, GPU peak/utilization, and storage/network byte counters. + +Query, stream, and ticket values must be logged only as bounded opaque hashes. +Certificate material, SQL values, batch contents, object paths, and credentials +must not be logged. + +## 14. Verification and acceptance gates + +The feature is not merge-ready until evidence names the exact three revisions +and closes every row below. + +| Contract | Required evidence | +| --- | --- | +| protocol and schema | cross-repository codec fixtures plus malformed/version/hash/schema/type/sequence/size controls | +| one-pass ownership | duplicate claim/ref/attachment rejection and exact cumulative acknowledgements | +| success lifecycle | producer join, sidecar worker quiescence, zero input handlers, zero retained execution after result EOF | +| early/pruned input | result EOF before first batch and `not_needed` after current/previous acknowledgement | +| cancellation | cancel while input ack is blocked and while result receive/write is blocked; bounded termination | +| injected failure | MO producer failure, sidecar input failure, Sirius consumer failure, disconnect, timeout, and retryable cleanup | +| slow consumer and full barrier | deterministic barriers prove one active source task per read, all retained batches hold budget leases, resource exhaustion fails without waiting, and terminal charge returns to zero | +| GPU execution | query-scoped evidence records `SIRIUS_GPU` and `GPU_MO_SCAN`; two configured workers never exceed one active `PARTITION` per GPU while non-partition tasks remain concurrent | +| correctness | typed native-MO equality for all 22 TPC-H SF1 queries on one reused process | +| SF10 correctness and decision data | typed equality for all 22 queries on one reused process; Q9 repeats ten times; record storage bytes, rows/bytes before serialization, transferred bytes, CN CPU/peak memory, sidecar host/GPU peak and utilization, time to first row, and total latency | +| partition safety control | repeated direct-TAE Q9 plus deterministic concurrent partition-content fingerprints prove the core gate is not StreamRead-specific | +| snapshot advantage | unflushed committed tail and visible tombstone cases equal native MatrixOne while direct `TaeRead` rejects them | +| static/build quality | MatrixOne SCA/UT/BVT/coverage; Sirius build matrix and tests; sidecar CUDA build/tests and review | + +Functional lifecycle tests use deterministic barriers rather than sleeps. +Performance and capacity measurements run in the performance harness, not as +wall-clock assertions in ordinary unit tests. + +## 15. Delivery pins + +The final evidence record must replace `pending` with immutable merge-ready +commits: + +| Component | PR | Approved commit | CI/evidence | +| --- | --- | --- | --- | +| MatrixOne | #27599 | pending | pending | +| Sirius | #6 | pending | pending | +| sidecar | #14 | pending | pending | + +The sidecar submodule must point to the approved Sirius commit. The MatrixOne PR +body must link this design at its approved commit and the final evidence record. + +## 16. Decision log + +| Decision | Rationale | +| --- | --- | +| retain direct `TaeRead` | it remains the lower-copy fast path for eligible flushed tables | +| select MO native instead of Arrow record batches | avoids Arrow-Go and duplicate conversion while accepting an exact same-release ABI | +| strict explicit stream mode | never hides a protocol, security, or GPU failure behind native execution | +| one local CN per sidecar | keeps snapshot and cancellation ownership local; multi-CN fan-in is a separate design | +| attach all inputs before `DoGet` | prevents a pruned plan from retiring a ticket before its handler can attach | +| one acknowledged slot per input | gives deterministic backpressure and ownership transfer | +| one active Sirius source task per input | removes eager self-scheduling; blocking barriers may retain completed batches only while budgeted | +| one partition task per physical GPU | the old executor produced nondeterministic Q9 values and fatal CUDA errors when two partition launches overlapped; other GPU operators retain configured parallelism | +| process-global 8 GiB StreamRead host budget | bounds all concurrent input slots, copies, and retained native representations together; fail-closed admission avoids barrier deadlock | +| bounded TPC-H stabilization on the current fork | delivers the required workload without claiming general unbounded streaming; migration to upstream partial-barrier scheduling is a separate effort | +| flat-only result vectors | prevents tiny compressed frames from expanding into unbounded MatrixOne result work | +| exact capability equality | rejects mixed ABI revisions before execution rather than attempting unsafe compatibility | + +There are no deferred correctness or lifecycle decisions. Production enablement +remains blocked on design approval and the acceptance evidence in sections 14 +and 15. diff --git a/pkg/container/vector/vector.go b/pkg/container/vector/vector.go index e8379b535cd86..31244ee6c1312 100644 --- a/pkg/container/vector/vector.go +++ b/pkg/container/vector/vector.go @@ -3243,6 +3243,11 @@ func (v *Vector) IsConst() bool { return v.class == CONSTANT } +// IsFlat reports whether the vector stores one physical value per logical row. +func (v *Vector) IsFlat() bool { + return v.class == FLAT +} + func (v *Vector) IsGrouping() bool { return v.length > 0 && v.length == v.gsp.Count() && diff --git a/pkg/frontend/computation_wrapper.go b/pkg/frontend/computation_wrapper.go index 14cd39572f06d..cbb9de1c23875 100644 --- a/pkg/frontend/computation_wrapper.go +++ b/pkg/frontend/computation_wrapper.go @@ -2191,6 +2191,9 @@ func compileStatementContexts( ) (requestCtx, compileCtx context.Context) { requestCtx = perfcounter.AttachCompilePlanMarkKey(ctx, crs) if siriusStatementSelected(sql, stmt) { + if strings.HasPrefix(strings.ToUpper(strings.TrimSpace(sql)), sidecarStreamHintPrefix) { + return requestCtx, compile.WithSiriusStreamOffload(requestCtx) + } return requestCtx, compile.WithSiriusOffload(requestCtx) } return requestCtx, requestCtx diff --git a/pkg/frontend/computation_wrapper_test.go b/pkg/frontend/computation_wrapper_test.go index e0f732af54077..71e609927731d 100644 --- a/pkg/frontend/computation_wrapper_test.go +++ b/pkg/frontend/computation_wrapper_test.go @@ -2331,6 +2331,7 @@ func TestCompileStatementContextsPreserveCounterWithoutLeakingSelection(t *testi separateChild bool }{ {name: "selected", sql: "/*+ SIDECAR */ select 1", separateChild: true}, + {name: "stream selected", sql: "/*+ SIDECAR STREAM */ select 1", separateChild: true}, {name: "unselected", sql: "select 1"}, } { t.Run(test.name, func(t *testing.T) { diff --git a/pkg/frontend/sidecar_offload.go b/pkg/frontend/sidecar_offload.go index e374371eb1c9e..51d85f242c8b2 100644 --- a/pkg/frontend/sidecar_offload.go +++ b/pkg/frontend/sidecar_offload.go @@ -36,9 +36,10 @@ import ( ) const ( - sidecarHintPrefix = "/*+ SIDECAR */" - sidecarGPUHintPrefix = "/*+ SIDECAR GPU */" - sidecarMaxResponseSize = 512 << 20 // 512 MB + sidecarHintPrefix = "/*+ SIDECAR */" + sidecarGPUHintPrefix = "/*+ SIDECAR GPU */" + sidecarStreamHintPrefix = "/*+ SIDECAR STREAM */" + sidecarMaxResponseSize = 512 << 20 // 512 MB ) // errSidecarNotConfigured is a sentinel indicating sidecar offload should @@ -94,6 +95,9 @@ func getManifestBaseURL() string { func isSidecarQuery(sql string) (bool, bool) { trimmed := strings.TrimSpace(sql) upper := strings.ToUpper(trimmed) + if strings.HasPrefix(upper, sidecarStreamHintPrefix) { + return true, true + } if strings.HasPrefix(upper, sidecarGPUHintPrefix) { return true, true } @@ -108,6 +112,9 @@ func isSidecarQuery(sql string) (bool, bool) { func stripSidecarHint(sql string) string { trimmed := strings.TrimSpace(sql) upper := strings.ToUpper(trimmed) + if strings.HasPrefix(upper, sidecarStreamHintPrefix) { + return strings.TrimSpace(trimmed[len(sidecarStreamHintPrefix):]) + } if strings.HasPrefix(upper, sidecarGPUHintPrefix) { return strings.TrimSpace(trimmed[len(sidecarGPUHintPrefix):]) } diff --git a/pkg/frontend/sidecar_offload_test.go b/pkg/frontend/sidecar_offload_test.go index f515baa9c0934..93ee7836d32ec 100644 --- a/pkg/frontend/sidecar_offload_test.go +++ b/pkg/frontend/sidecar_offload_test.go @@ -57,6 +57,10 @@ func TestIsSidecarQuery(t *testing.T) { assert.True(t, isSidecar) assert.True(t, useGPU) + isSidecar, useGPU = isSidecarQuery("/*+ SIDECAR STREAM */ SELECT * FROM t") + assert.True(t, isSidecar) + assert.True(t, useGPU) + // Non-sidecar queries isSidecar, _ = isSidecarQuery("SELECT * FROM t") assert.False(t, isSidecar) @@ -73,6 +77,7 @@ func TestStripSidecarHint(t *testing.T) { assert.Equal(t, "SELECT * FROM t", stripSidecarHint(" /*+ SIDECAR */ SELECT * FROM t")) assert.Equal(t, "SELECT * FROM t", stripSidecarHint("/*+ SIDECAR GPU */ SELECT * FROM t")) assert.Equal(t, "SELECT * FROM t", stripSidecarHint(" /*+ sidecar gpu */ SELECT * FROM t")) + assert.Equal(t, "SELECT * FROM t", stripSidecarHint("/*+ SIDECAR STREAM */ SELECT * FROM t")) assert.Equal(t, "SELECT * FROM t", stripSidecarHint("SELECT * FROM t")) } diff --git a/pkg/sql/colexec/output/output.go b/pkg/sql/colexec/output/output.go index 7783ad15a519d..8a1db4128772c 100644 --- a/pkg/sql/colexec/output/output.go +++ b/pkg/sql/colexec/output/output.go @@ -89,6 +89,9 @@ func (output *Output) Call(proc *process.Process) (vm.CallResult, error) { result.Status = vm.ExecStop return result, err } + if output.stop != nil && output.stop() { + result.Status = vm.ExecStop + } // TODO: analyzer.Output(result.Batch) return result, nil @@ -140,6 +143,9 @@ func (output *Output) Call(proc *process.Process) (vm.CallResult, error) { result.Status = vm.ExecStop return result, err } + if output.stop != nil && output.stop() { + result.Status = vm.ExecStop + } result.Batch = bat // same as nonBlock diff --git a/pkg/sql/colexec/output/output_test.go b/pkg/sql/colexec/output/output_test.go index 310daa46a9d19..ecd2ab7148167 100644 --- a/pkg/sql/colexec/output/output_test.go +++ b/pkg/sql/colexec/output/output_test.go @@ -156,6 +156,46 @@ func TestOutputCallbackCPUIsNotOutputWait(t *testing.T) { require.Zero(t, proc.Mp().CurrNB()) } +func TestOutputCanStopOneProducerAfterSuccessfulCallback(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + calls := 0 + arg := NewArgument().WithFunc(func(_ *batch.Batch, _ *perfcounter.CounterSet) error { + calls++ + return nil + }).WithShouldStop(func() bool { return calls == 1 }) + require.NoError(t, arg.Prepare(proc)) + first := newBatch([]types.Type{types.T_int8.ToType()}, proc, 1) + second := newBatch([]types.Type{types.T_int8.ToType()}, proc, 1) + resetChildren(arg, []*batch.Batch{first, second}) + result, err := vm.Exec(arg, proc) + require.NoError(t, err) + require.Equal(t, vm.ExecStop, result.Status) + require.Equal(t, 1, calls) + arg.GetChildren(0).Free(proc, false, nil) + arg.Free(proc, false, nil) + proc.Free() +} + +func TestBlockingOutputCanStopAfterFirstReleasedBatch(t *testing.T) { + proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero()) + calls := 0 + arg := NewArgument().WithBlock(true).WithFunc(func(_ *batch.Batch, _ *perfcounter.CounterSet) error { + calls++ + return nil + }).WithShouldStop(func() bool { return calls == 1 }) + require.NoError(t, arg.Prepare(proc)) + first := newBatch([]types.Type{types.T_int8.ToType()}, proc, 1) + second := newBatch([]types.Type{types.T_int8.ToType()}, proc, 1) + resetChildren(arg, []*batch.Batch{first, second}) + result, err := vm.Exec(arg, proc) + require.NoError(t, err) + require.Equal(t, vm.ExecStop, result.Status) + require.Equal(t, 1, calls) + arg.GetChildren(0).Free(proc, false, nil) + arg.Free(proc, false, nil) + proc.Free() +} + // create a new block based on the type information func newBatch(ts []types.Type, proc *process.Process, rows int64) *batch.Batch { return testutil.NewBatch(ts, false, int(rows), proc.Mp()) diff --git a/pkg/sql/colexec/output/types.go b/pkg/sql/colexec/output/types.go index c75755b9cdce8..5b16020e93557 100644 --- a/pkg/sql/colexec/output/types.go +++ b/pkg/sql/colexec/output/types.go @@ -46,6 +46,7 @@ type Output struct { Data interface{} Func func(*batch.Batch, *perfcounter.CounterSet) error + stop func() bool // IsAdaptive enables the adaptive vector search fallback mechanism. // When set to true and the query completes with zero results (rowCount == 0), @@ -99,6 +100,14 @@ func (output *Output) WithFunc(Func func(*batch.Batch, *perfcounter.CounterSet) return output } +// WithShouldStop lets a sink end only its own producer pipeline after a +// successful callback. It is used when a remote consumer declares that this +// input relation is no longer needed; nil preserves the ordinary output path. +func (output *Output) WithShouldStop(stop func() bool) *Output { + output.stop = stop + return output +} + // WithBlocck set the output is blocked. If true output will block the current pipeline, and cache // all input batches. And wait for all the input's batch to be locked before outputting the cached batch // to the downstream operator. diff --git a/pkg/sql/compile/compile2.go b/pkg/sql/compile/compile2.go index ffd3beacbf882..5410cd42f5c1f 100644 --- a/pkg/sql/compile/compile2.go +++ b/pkg/sql/compile/compile2.go @@ -340,7 +340,9 @@ func (c *Compile) Run(_ uint64) (queryResult *util2.RunResult, err error) { queryResult = &util2.RunResult{} v2.TxnStatementTotalCounter.Inc() if c.siriusRead != nil { - err = c.runSiriusRead(execTopContext) + err = c.runSiriusRead(execTopContext, func(snapshot mpool.AllocationAccountTerminalSnapshot) { + resourceRecorder.recordAllocationAccountTerminal(snapshot) + }) return queryResult, err } attemptStart := time.Now() diff --git a/pkg/sql/compile/sidecarflight/arrow_ipc.go b/pkg/sql/compile/sidecarflight/arrow_ipc.go deleted file mode 100644 index e56de24c1f6ea..0000000000000 --- a/pkg/sql/compile/sidecarflight/arrow_ipc.go +++ /dev/null @@ -1,763 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// 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 sidecarflight - -import ( - "encoding/binary" - "math" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/batch" - "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" - planpb "github.com/matrixorigin/matrixone/pkg/pb/plan" -) - -const ( - arrowHeaderSchema = byte(1) - arrowHeaderRecordBatch = byte(3) - maxArrowMetadataBytes = 1 << 20 - - arrowTypeInt = byte(2) - arrowTypeFloatingPoint = byte(3) - arrowTypeUTF8 = byte(5) - arrowTypeBool = byte(6) - arrowTypeDecimal = byte(7) - arrowTypeDate = byte(8) -) - -type arrowField struct { - name string - nullable bool - arrowType byte - bitWidth int32 - isSigned bool - precision int32 - scale int32 - floatKind int16 - dateUnit int16 - expected planpb.Type - bufferCount int -} - -// Schema is the validated flat result shape shared by FlightInfo and every -// FlightData message in one execution. -type Schema struct { - fields []arrowField - headings []string -} - -// ParseSchema decodes the Arrow IPC schema without bringing an Arrow runtime -// into MatrixOne. Only the flat, dictionary-free types in the negotiated TPC-H -// capability are accepted. -func ParseSchema(wire []byte, expected []planpb.Type, headings []string) (*Schema, error) { - if len(wire) == 0 || len(wire) > maxArrowMetadataBytes { - return nil, internalErrorf("Arrow schema metadata exceeds the supported bound") - } - if len(expected) == 0 || len(headings) != len(expected) { - return nil, internalErrorf("MatrixOne result schema is empty or inconsistent") - } - metadata, err := ipcMetadata(wire) - if err != nil { - return nil, err - } - message, err := rootTable(metadata) - if err != nil { - return nil, internalErrorf("Arrow schema message: %w", err) - } - version, err := message.byteField(0, 0) - if err != nil || version != 4 { - return nil, internalErrorf("Arrow schema message has unsupported metadata version %d", version) - } - headerType, err := message.byteField(1, 0) - if err != nil || headerType != arrowHeaderSchema { - return nil, internalErrorf("Arrow schema message has header type %d", headerType) - } - if bodyLength, bodyErr := message.int64Field(3, 0); bodyErr != nil || bodyLength != 0 { - return nil, internalErrorf("Arrow schema message has an invalid body length") - } - schemaTable, ok, err := message.tableField(2) - if err != nil { - return nil, internalErrorf("Arrow schema message is missing its schema: %w", err) - } - if !ok { - return nil, internalErrorf("Arrow schema message is missing its schema") - } - endianness, err := schemaTable.int16Field(0, 0) - if err != nil || endianness != 0 { - return nil, internalErrorf("Arrow schema is not little-endian") - } - _, featureCount, _, err := schemaTable.vector(3, 8) - if err != nil || featureCount != 0 { - return nil, internalErrorf("Arrow schema uses unsupported features") - } - _, fieldCount, fieldsPresent, err := schemaTable.vector(1, 4) - if err != nil { - return nil, internalErrorf("Arrow schema fields: %w", err) - } - if !fieldsPresent { - fieldCount = 0 - } - if fieldCount != len(expected) || len(headings) != len(expected) { - return nil, internalErrorf("Arrow schema has %d fields; MatrixOne expects %d", fieldCount, len(expected)) - } - fieldTables, err := schemaTable.tableVector(1) - if err != nil { - return nil, internalErrorf("Arrow schema fields: %w", err) - } - result := &Schema{fields: make([]arrowField, len(fieldTables)), headings: append([]string(nil), headings...)} - for i, table := range fieldTables { - field, parseErr := parseArrowField(table, expected[i]) - if parseErr != nil { - return nil, internalErrorf("Arrow field %d: %w", i, parseErr) - } - if field.name != headings[i] { - return nil, internalErrorf("Arrow field %d is named %q; MatrixOne expects %q", i, field.name, headings[i]) - } - result.fields[i] = field - } - return result, nil -} - -func parseArrowField(table flatTable, expected planpb.Type) (arrowField, error) { - name, ok, err := table.stringField(0) - if err != nil { - return arrowField{}, internalErrorf("missing name: %w", err) - } - if !ok { - return arrowField{}, internalErrorf("missing name") - } - if dictionary, present, dictionaryErr := table.tableField(4); dictionaryErr != nil { - return arrowField{}, dictionaryErr - } else if present || dictionary.data != nil { - return arrowField{}, internalErrorf("dictionary encoding is not supported") - } - _, childCount, _, err := table.vector(5, 4) - if err != nil { - return arrowField{}, err - } - if childCount != 0 { - return arrowField{}, internalErrorf("nested Arrow fields are not supported") - } - nullable, err := table.boolField(1, false) - if err != nil { - return arrowField{}, err - } - typeID, err := table.byteField(2, 0) - if err != nil { - return arrowField{}, err - } - typeTable, ok, err := table.tableField(3) - if err != nil { - return arrowField{}, internalErrorf("missing Arrow type metadata: %w", err) - } - if !ok { - return arrowField{}, internalErrorf("missing Arrow type metadata") - } - field := arrowField{name: name, nullable: nullable, arrowType: typeID, expected: expected, bufferCount: 2} - moType := types.T(expected.Id) - switch typeID { - case arrowTypeBool: - if moType != types.T_bool { - return arrowField{}, typeMismatch(typeID, moType) - } - field.bitWidth = 1 - case arrowTypeInt: - field.bitWidth, err = typeTable.int32Field(0, 0) - if err != nil { - return arrowField{}, err - } - field.isSigned, err = typeTable.boolField(1, false) - if err != nil { - return arrowField{}, err - } - wantWidth, wantSigned := expectedIntegerShape(moType) - if moType == types.T_uint32 { - // Sirius transports EXTRACT's unsigned MO result as signed i64. - wantWidth, wantSigned = 64, true - } - if field.bitWidth != wantWidth || field.isSigned != wantSigned { - return arrowField{}, typeMismatch(typeID, moType) - } - case arrowTypeFloatingPoint: - field.floatKind, err = typeTable.int16Field(0, 0) - if err != nil { - return arrowField{}, err - } - if (moType == types.T_float32 && field.floatKind != 1) || (moType == types.T_float64 && field.floatKind != 2) || - (moType != types.T_float32 && moType != types.T_float64) { - return arrowField{}, typeMismatch(typeID, moType) - } - field.bitWidth = map[types.T]int32{types.T_float32: 32, types.T_float64: 64}[moType] - case arrowTypeUTF8: - if moType != types.T_char && moType != types.T_varchar { - return arrowField{}, typeMismatch(typeID, moType) - } - field.bufferCount = 3 - case arrowTypeDecimal: - field.precision, err = typeTable.int32Field(0, 0) - if err != nil { - return arrowField{}, err - } - field.scale, err = typeTable.int32Field(1, 0) - if err != nil { - return arrowField{}, err - } - field.bitWidth, err = typeTable.int32Field(2, 128) - if err != nil { - return arrowField{}, err - } - if (moType != types.T_decimal64 && moType != types.T_decimal128) || field.precision != expected.Width || - field.scale != expected.Scale || field.bitWidth != 128 { - return arrowField{}, typeMismatch(typeID, moType) - } - case arrowTypeDate: - field.dateUnit, err = typeTable.int16Field(0, 0) - if err != nil { - return arrowField{}, err - } - if moType != types.T_date || field.dateUnit != 0 { - return arrowField{}, typeMismatch(typeID, moType) - } - field.bitWidth = 32 - default: - return arrowField{}, internalErrorf("unsupported Arrow type id %d", typeID) - } - return field, nil -} - -func expectedIntegerShape(t types.T) (int32, bool) { - switch t { - case types.T_int8: - return 8, true - case types.T_int16: - return 16, true - case types.T_int32: - return 32, true - case types.T_int64: - return 64, true - default: - return 0, false - } -} - -func typeMismatch(arrowType byte, moType types.T) error { - return internalErrorf("Arrow type %d does not match MatrixOne type %s", arrowType, moType.String()) -} - -func (s *Schema) matches(other *Schema) bool { - if s == nil || other == nil || len(s.fields) != len(other.fields) { - return false - } - for i := range s.fields { - left, right := s.fields[i], other.fields[i] - if left.name != right.name || left.nullable != right.nullable || left.arrowType != right.arrowType || left.bitWidth != right.bitWidth || - left.isSigned != right.isSigned || left.precision != right.precision || left.scale != right.scale || - left.floatKind != right.floatKind || left.dateUnit != right.dateUnit || - left.expected.Id != right.expected.Id || left.expected.Width != right.expected.Width || - left.expected.Scale != right.expected.Scale || left.expected.NotNullable != right.expected.NotNullable || - left.expected.Charset != right.expected.Charset { - return false - } - } - return true -} - -func (s *Schema) validateStreamSchema(header []byte) error { - expected := make([]planpb.Type, len(s.fields)) - for i := range s.fields { - expected[i] = s.fields[i].expected - } - streamSchema, err := ParseSchema(header, expected, s.headings) - if err != nil { - return err - } - if !s.matches(streamSchema) { - return internalErrorf("Arrow stream schema differs from FlightInfo schema") - } - return nil -} - -type arrowNode struct { - length int64 - nullCount int64 -} - -type arrowBuffer struct { - offset int64 - length int64 -} - -// decodeRecordBatch converts exactly one flat Arrow record batch into MO -// vectors. The returned batch owns its memory and must be cleaned by the -// synchronous consumer before the next Flight message is requested. -func (s *Schema) decodeRecordBatch(header, body []byte, maxDecodedBytes uint64, mp *mpool.MPool) (result *batch.Batch, err error) { - if s == nil || mp == nil || maxDecodedBytes == 0 { - return nil, internalErrorf("sidecar flight: missing schema or memory pool") - } - if len(header) == 0 || len(header) > maxArrowMetadataBytes { - return nil, internalErrorf("Arrow record batch metadata exceeds the supported bound") - } - metadata, err := ipcMetadata(header) - if err != nil { - return nil, err - } - message, err := rootTable(metadata) - if err != nil { - return nil, err - } - version, err := message.byteField(0, 0) - if err != nil || version != 4 { - return nil, internalErrorf("Arrow record batch has unsupported metadata version %d", version) - } - headerType, err := message.byteField(1, 0) - if err != nil || headerType != arrowHeaderRecordBatch { - return nil, internalErrorf("Arrow message has unsupported header type %d", headerType) - } - bodyLength, err := message.int64Field(3, 0) - if err != nil || bodyLength < 0 || bodyLength != int64(len(body)) { - return nil, internalErrorf("Arrow record batch body length mismatch") - } - record, ok, err := message.tableField(2) - if err != nil { - return nil, internalErrorf("Arrow record batch metadata is missing: %w", err) - } - if !ok { - return nil, internalErrorf("Arrow record batch metadata is missing") - } - rows, err := record.int64Field(0, 0) - if err != nil || rows < 0 || rows > int64(maxInt()) { - return nil, internalErrorf("Arrow record batch row count is invalid") - } - if compression, present, compressionErr := record.tableField(3); compressionErr != nil { - return nil, compressionErr - } else if present || compression.data != nil { - return nil, internalErrorf("compressed Arrow batches are not supported") - } - _, variadicCount, _, err := record.vector(4, 8) - if err != nil || variadicCount != 0 { - return nil, internalErrorf("variadic Arrow buffers are not supported") - } - nodeBytes, nodeCount, err := record.structVector(1, 16) - if err != nil || nodeCount != len(s.fields) { - return nil, internalErrorf("Arrow record batch has %d field nodes; expected %d", nodeCount, len(s.fields)) - } - nodes := make([]arrowNode, nodeCount) - for i := range nodes { - nodes[i] = arrowNode{length: int64(binary.LittleEndian.Uint64(nodeBytes[i*16:])), nullCount: int64(binary.LittleEndian.Uint64(nodeBytes[i*16+8:]))} - if nodes[i].length != rows || nodes[i].nullCount < 0 || nodes[i].nullCount > rows { - return nil, internalErrorf("Arrow field %d has invalid node metadata", i) - } - } - bufferBytes, bufferCount, err := record.structVector(2, 16) - wantBuffers := 0 - for _, field := range s.fields { - wantBuffers += field.bufferCount - } - if err != nil || bufferCount != wantBuffers { - return nil, internalErrorf("Arrow record batch has %d buffers; expected %d", bufferCount, wantBuffers) - } - buffers := make([]arrowBuffer, bufferCount) - for i := range buffers { - buffers[i] = arrowBuffer{offset: int64(binary.LittleEndian.Uint64(bufferBytes[i*16:])), length: int64(binary.LittleEndian.Uint64(bufferBytes[i*16+8:]))} - if buffers[i].offset < 0 || buffers[i].length < 0 || buffers[i].offset > int64(len(body)) || buffers[i].length > int64(len(body))-buffers[i].offset { - return nil, internalErrorf("Arrow buffer %d is outside the record body", i) - } - } - decodedBytes := uint64(0) - bufferIndex := 0 - for _, field := range s.fields { - rowBytes := uint64(types.New(types.T(field.expected.Id), field.expected.Width, field.expected.Scale).TypeSize()) - if rowBytes == 0 || uint64(rows) > (maxDecodedBytes-decodedBytes)/rowBytes { - return nil, internalErrorf("Arrow record batch exceeds the decoded-memory budget") - } - decodedBytes += uint64(rows) * rowBytes - nullBytes := uint64(bitmapBytes(rows)) - if nullBytes > maxDecodedBytes-decodedBytes { - return nil, internalErrorf("Arrow record batch exceeds the decoded-memory budget") - } - decodedBytes += nullBytes - if field.arrowType == arrowTypeUTF8 { - dataBytes := uint64(buffers[bufferIndex+2].length) - if dataBytes > maxDecodedBytes-decodedBytes { - return nil, internalErrorf("Arrow record batch exceeds the decoded-memory budget") - } - decodedBytes += dataBytes - } - bufferIndex += field.bufferCount - } - result = batch.NewWithSize(len(s.fields)) - result.Attrs = append([]string(nil), s.headings...) - defer func() { - if err != nil && result != nil { - result.Clean(mp) - result = nil - } - }() - bufferIndex = 0 - for i, field := range s.fields { - moType := types.New(types.T(field.expected.Id), field.expected.Width, field.expected.Scale) - moType.SetNotNull(field.expected.NotNullable) - result.Vecs[i] = vector.NewVec(moType) - columnBuffers := buffers[bufferIndex : bufferIndex+field.bufferCount] - bufferIndex += field.bufferCount - if err = decodeColumn(result.Vecs[i], field, nodes[i], columnBuffers, body, mp); err != nil { - return nil, internalErrorf("Arrow field %d: %w", i, err) - } - } - result.SetRowCount(int(rows)) - return result, nil -} - -func decodeColumn(vec *vector.Vector, field arrowField, node arrowNode, buffers []arrowBuffer, body []byte, mp *mpool.MPool) error { - if (!field.nullable || field.expected.NotNullable) && node.nullCount != 0 { - return internalErrorf("required field contains nulls") - } - validity := sliceBuffer(body, buffers[0]) - if node.nullCount == 0 { - if len(validity) != 0 && int64(len(validity)) < bitmapBytes(node.length) { - return internalErrorf("validity buffer is too short") - } - } else if int64(len(validity)) < bitmapBytes(node.length) { - return internalErrorf("validity buffer is too short") - } - isNull := func(row int64) bool { - return node.nullCount != 0 && validity[row>>3]&(1< math.MaxInt64/4-1 || int64(len(offsetsBytes)) != (node.length+1)*4 { - return internalErrorf("UTF8 offsets have an invalid length") - } - previous := int32(0) - for row := int64(0); row < node.length; row++ { - start := int32(binary.LittleEndian.Uint32(offsetsBytes[row*4:])) - end := int32(binary.LittleEndian.Uint32(offsetsBytes[(row+1)*4:])) - if start != previous || end < start || end < 0 || int64(end) > int64(len(data)) { - return internalErrorf("UTF8 offsets are invalid") - } - if err := vector.AppendBytes(vec, data[start:end], isNull(row), mp); err != nil { - return err - } - previous = end - } - if int(previous) != len(data) { - return internalErrorf("UTF8 data has trailing bytes") - } - return nil - } - width := int64(field.bitWidth / 8) - if field.arrowType == arrowTypeBool { - width = 0 - if int64(len(values)) < bitmapBytes(node.length) { - return internalErrorf("boolean values buffer is too short") - } - } else if width <= 0 || node.length > math.MaxInt64/width || int64(len(values)) != node.length*width { - return internalErrorf("fixed-width values have an invalid length") - } - for row := int64(0); row < node.length; row++ { - null := isNull(row) - offset := row * width - var err error - switch field.arrowType { - case arrowTypeBool: - err = vector.AppendFixed(vec, values[row>>3]&(1< math.MaxUint32) { - return internalErrorf("signed i64 value %d overflows MatrixOne uint32", value) - } - err = vector.AppendFixed(vec, uint32(value), null, mp) - default: - return internalErrorf("unsupported MatrixOne integer type") - } - case arrowTypeFloatingPoint: - if types.T(field.expected.Id) == types.T_float32 { - err = vector.AppendFixed(vec, math.Float32frombits(binary.LittleEndian.Uint32(values[offset:])), null, mp) - } else { - err = vector.AppendFixed(vec, math.Float64frombits(binary.LittleEndian.Uint64(values[offset:])), null, mp) - } - case arrowTypeDate: - err = vector.AppendFixed(vec, types.DaysFromUnixEpochToDate(int32(binary.LittleEndian.Uint32(values[offset:]))), null, mp) - case arrowTypeDecimal: - low := binary.LittleEndian.Uint64(values[offset:]) - high := binary.LittleEndian.Uint64(values[offset+8:]) - if types.T(field.expected.Id) == types.T_decimal64 { - signExtension := uint64(0) - if low>>63 != 0 { - signExtension = math.MaxUint64 - } - if !null && high != signExtension { - return internalErrorf("decimal128 value does not fit MatrixOne decimal64") - } - err = vector.AppendFixed(vec, types.Decimal64(low), null, mp) - } else { - err = vector.AppendFixed(vec, types.Decimal128{B0_63: low, B64_127: high}, null, mp) - } - default: - return internalErrorf("unsupported Arrow field type") - } - if err != nil { - return err - } - } - return nil -} - -func bitmapBytes(rows int64) int64 { - result := rows / 8 - if rows%8 != 0 { - result++ - } - return result -} - -func sliceBuffer(body []byte, buffer arrowBuffer) []byte { - return body[int(buffer.offset):int(buffer.offset+buffer.length)] -} - -// ipcMetadata accepts both raw Message flatbuffers and stream-framed IPC -// metadata (continuation marker plus size, or the legacy size prefix). -func ipcMetadata(wire []byte) ([]byte, error) { - if len(wire) < 4 { - return nil, internalErrorf("Arrow IPC metadata is truncated") - } - if binary.LittleEndian.Uint32(wire[:4]) == math.MaxUint32 { - if len(wire) < 8 { - return nil, internalErrorf("Arrow IPC continuation header is truncated") - } - length := uint64(binary.LittleEndian.Uint32(wire[4:8])) - if length == 0 || length > uint64(len(wire)-8) { - return nil, internalErrorf("Arrow IPC metadata length is invalid") - } - return wire[8 : 8+length], nil - } - length := uint64(binary.LittleEndian.Uint32(wire[:4])) - if length != 0 && length == uint64(len(wire)-4) { - return wire[4:], nil - } - return wire, nil -} - -type flatTable struct { - data []byte - start uint64 -} - -func rootTable(data []byte) (flatTable, error) { - if len(data) < 4 { - return flatTable{}, internalErrorf("flatbuffer root is truncated") - } - start := uint64(binary.LittleEndian.Uint32(data[:4])) - table := flatTable{data: data, start: start} - if start < 4 { - return flatTable{}, internalErrorf("flatbuffer root offset is invalid") - } - if _, _, err := table.vtable(); err != nil { - return flatTable{}, err - } - return table, nil -} - -func (t flatTable) vtable() (uint64, uint16, error) { - if t.start+4 > uint64(len(t.data)) { - return 0, 0, internalErrorf("flatbuffer table is outside metadata") - } - back := int64(int32(binary.LittleEndian.Uint32(t.data[t.start:]))) - vtablePosition := int64(t.start) - back - if vtablePosition < 0 || uint64(vtablePosition) > uint64(len(t.data))-4 { - return 0, 0, internalErrorf("flatbuffer vtable offset is invalid") - } - start := uint64(vtablePosition) - length := binary.LittleEndian.Uint16(t.data[start:]) - objectLength := binary.LittleEndian.Uint16(t.data[start+2:]) - if length < 4 || start+uint64(length) > uint64(len(t.data)) || objectLength < 4 || t.start+uint64(objectLength) > uint64(len(t.data)) { - return 0, 0, internalErrorf("flatbuffer vtable has invalid bounds") - } - return start, length, nil -} - -func (t flatTable) field(index int, width uint64) (uint64, bool, error) { - vtable, length, err := t.vtable() - if err != nil { - return 0, false, err - } - entry := uint64(4 + index*2) - if entry+2 > uint64(length) { - return 0, false, nil - } - offset := uint64(binary.LittleEndian.Uint16(t.data[vtable+entry:])) - if offset == 0 { - return 0, false, nil - } - objectLength := uint64(binary.LittleEndian.Uint16(t.data[vtable+2:])) - if offset < 4 || offset > objectLength || width > objectLength-offset { - return 0, false, internalErrorf("flatbuffer field is outside its table") - } - position := t.start + offset - if position+width > uint64(len(t.data)) { - return 0, false, internalErrorf("flatbuffer field is truncated") - } - return position, true, nil -} - -func (t flatTable) byteField(index int, defaultValue byte) (byte, error) { - position, ok, err := t.field(index, 1) - if err != nil || !ok { - return defaultValue, err - } - return t.data[position], nil -} - -func (t flatTable) boolField(index int, defaultValue bool) (bool, error) { - value, err := t.byteField(index, map[bool]byte{false: 0, true: 1}[defaultValue]) - if err != nil || value > 1 { - return false, internalErrorf("flatbuffer boolean is invalid") - } - return value != 0, nil -} - -func (t flatTable) int16Field(index int, defaultValue int16) (int16, error) { - position, ok, err := t.field(index, 2) - if err != nil || !ok { - return defaultValue, err - } - return int16(binary.LittleEndian.Uint16(t.data[position:])), nil -} - -func (t flatTable) int32Field(index int, defaultValue int32) (int32, error) { - position, ok, err := t.field(index, 4) - if err != nil || !ok { - return defaultValue, err - } - return int32(binary.LittleEndian.Uint32(t.data[position:])), nil -} - -func (t flatTable) int64Field(index int, defaultValue int64) (int64, error) { - position, ok, err := t.field(index, 8) - if err != nil || !ok { - return defaultValue, err - } - return int64(binary.LittleEndian.Uint64(t.data[position:])), nil -} - -func (t flatTable) indirect(position uint64) (uint64, error) { - if position+4 > uint64(len(t.data)) { - return 0, internalErrorf("flatbuffer offset is truncated") - } - offset := uint64(binary.LittleEndian.Uint32(t.data[position:])) - if offset == 0 || offset > uint64(len(t.data))-position { - return 0, internalErrorf("flatbuffer offset is invalid") - } - return position + offset, nil -} - -func (t flatTable) tableField(index int) (flatTable, bool, error) { - position, ok, err := t.field(index, 4) - if err != nil || !ok { - return flatTable{}, false, err - } - start, err := t.indirect(position) - if err != nil { - return flatTable{}, false, err - } - result := flatTable{data: t.data, start: start} - if _, _, err = result.vtable(); err != nil { - return flatTable{}, false, err - } - return result, true, nil -} - -func (t flatTable) vector(index int, elementWidth uint64) (uint64, int, bool, error) { - position, ok, err := t.field(index, 4) - if err != nil || !ok { - return 0, 0, ok, err - } - start, err := t.indirect(position) - if err != nil || start+4 > uint64(len(t.data)) { - return 0, 0, true, internalErrorf("flatbuffer vector is truncated") - } - length := uint64(binary.LittleEndian.Uint32(t.data[start:])) - dataStart := start + 4 - if elementWidth != 0 && length > (uint64(len(t.data))-dataStart)/elementWidth { - return 0, 0, true, internalErrorf("flatbuffer vector has invalid bounds") - } - if length > uint64(maxInt()) { - return 0, 0, true, internalErrorf("flatbuffer vector is too large") - } - return dataStart, int(length), true, nil -} - -func (t flatTable) stringField(index int) (string, bool, error) { - start, length, ok, err := t.vector(index, 1) - if err != nil || !ok { - return "", ok, err - } - if start+uint64(length) >= uint64(len(t.data)) || t.data[start+uint64(length)] != 0 { - return "", true, internalErrorf("flatbuffer string is not terminated") - } - return string(t.data[start : start+uint64(length)]), true, nil -} - -func (t flatTable) tableVector(index int) ([]flatTable, error) { - start, length, ok, err := t.vector(index, 4) - if err != nil || !ok { - return nil, err - } - result := make([]flatTable, length) - for i := range result { - element := start + uint64(i*4) - tableStart, indirectErr := t.indirect(element) - if indirectErr != nil { - return nil, indirectErr - } - result[i] = flatTable{data: t.data, start: tableStart} - if _, _, indirectErr = result[i].vtable(); indirectErr != nil { - return nil, indirectErr - } - } - return result, nil -} - -func (t flatTable) structVector(index int, width uint64) ([]byte, int, error) { - start, length, ok, err := t.vector(index, width) - if err != nil || !ok { - return nil, 0, err - } - return t.data[start : start+uint64(length)*width], length, nil -} diff --git a/pkg/sql/compile/sidecarflight/arrow_ipc_test.go b/pkg/sql/compile/sidecarflight/arrow_ipc_test.go deleted file mode 100644 index 7ad3aeea161c0..0000000000000 --- a/pkg/sql/compile/sidecarflight/arrow_ipc_test.go +++ /dev/null @@ -1,333 +0,0 @@ -// Copyright 2026 Matrix Origin -// -// 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 sidecarflight - -import ( - "encoding/binary" - "encoding/hex" - "math" - "testing" - - "github.com/matrixorigin/matrixone/pkg/common/mpool" - "github.com/matrixorigin/matrixone/pkg/container/types" - "github.com/matrixorigin/matrixone/pkg/container/vector" - planpb "github.com/matrixorigin/matrixone/pkg/pb/plan" - "github.com/stretchr/testify/require" -) - -// Generated by Arrow C++ 25.0, the same native Arrow major used by the -// sidecar. Keeping the bytes in the Go test makes the production package -// independent of both a Go Arrow package and a native Arrow installation. -const fixtureSchemaHex = "ffffffff980200001000000000000a000c000600050008000a000000000104000c000000080008000000040008000000040000000c0000003c020000f4010000c401000094010000640100003001000004010000dc000000a400000070000000400000000400000004feffff00000102100000002000000004000000000000000d0000007533325f7472616e73706f727400000044feffff00000001400000003cfeffff000001081000000018000000040000000000000004000000646174650000000026ffffff0000000068feffff0000010710000000180000000400000000000000040000006431323800000000d4ffffff260000000400000098feffff00000107100000001c0000000400000000000000030000006436340008000c0004000800080000001200000002000000ccfeffff00000105100000001400000004000000000000000100000073000000bcfefffff0feffff00000103100000001400000004000000000000000300000066363400d6ffffff0000020018ffffff00000103100000001c000000040000000000000003000000663332000000060008000600060000000000010048ffffff000001021000000014000000040000000000000003000000693634007cffffff000000014000000074ffffff00000102100000001400000004000000000000000300000069333200a8ffffff0000000120000000a0ffffff00000102100000001400000004000000000000000300000069313600d4ffffff0000000110000000ccffffff00000102100000001c0000000400000000000000020000006938000008000c0008000700080000000000000108000000100014000800060007000c00000010001000000000000106100000001800000004000000000000000100000062000000040004000400000000000000" - -const fixtureHeaderHex = "ffffffffa802000014000000000000000c0016000600050008000c000c0000000003040018000000180100000000000000000a0018000c00040008000a000000ac010000100000000200000000000000000000001900000000000000000000000100000000000000080000000000000001000000000000001000000000000000010000000000000018000000000000000200000000000000200000000000000001000000000000002800000000000000040000000000000030000000000000000100000000000000380000000000000008000000000000004000000000000000010000000000000048000000000000001000000000000000580000000000000001000000000000006000000000000000080000000000000068000000000000000100000000000000700000000000000010000000000000008000000000000000010000000000000088000000000000000c0000000000000098000000000000000400000000000000a0000000000000000100000000000000a8000000000000002000000000000000c8000000000000000100000000000000d0000000000000002000000000000000f0000000000000000100000000000000f80000000000000008000000000000000001000000000000010000000000000008010000000000001000000000000000000000000c000000020000000000000001000000000000000200000000000000010000000000000002000000000000000100000000000000020000000000000001000000000000000200000000000000010000000000000002000000000000000100000000000000020000000000000001000000000000000200000000000000010000000000000002000000000000000100000000000000020000000000000001000000000000000200000000000000010000000000000002000000000000000100000000000000" - -const fixtureBodyHex = "010000000000000001000000000000000100000000000000f8000000000000000100000000000000f0ff0000000000000100000000000000e0ffffff000000000100000000000000c0ffffffffffffff000000000000000001000000000000000000a03f0000000001000000000000000000000000000440000000000000000001000000000000000000000004000000040000000000000074706368000000000100000000000000c7cfffffffffffffffffffffffffffff00000000000000000000000000000000010000000000000007000000000000000100000000000000000000000000000000000000000000000100000000000000010000000000000001000000000000002a000000000000000000000000000000" - -func TestArrowIPCDecodesNegotiatedTPCHTypes(t *testing.T) { - schemaWire := mustHex(t, fixtureSchemaHex) - header := mustHex(t, fixtureHeaderHex) - body := mustHex(t, fixtureBodyHex) - typesOut := []planpb.Type{ - {Id: int32(types.T_bool)}, {Id: int32(types.T_int8)}, {Id: int32(types.T_int16)}, - {Id: int32(types.T_int32)}, {Id: int32(types.T_int64)}, {Id: int32(types.T_float32)}, - {Id: int32(types.T_float64)}, {Id: int32(types.T_varchar)}, - {Id: int32(types.T_decimal64), Width: 18, Scale: 2}, - {Id: int32(types.T_decimal128), Width: 38, Scale: 4}, - {Id: int32(types.T_date)}, {Id: int32(types.T_uint32)}, - } - headings := []string{"b", "i8", "i16", "i32", "i64", "f32", "f64", "s", "d64", "d128", "date", "u32_transport"} - schema, err := ParseSchema(schemaWire, typesOut, headings) - require.NoError(t, err) - require.NoError(t, schema.validateStreamSchema(schemaWire)) - - mp := mpool.MustNewZero() - bat, err := schema.decodeRecordBatch(header, body, 1<<20, mp) - require.NoError(t, err) - require.Equal(t, 2, bat.RowCount()) - for _, vec := range bat.Vecs { - require.False(t, vec.IsNull(0)) - require.True(t, vec.IsNull(1)) - } - require.True(t, vector.GetFixedAtNoTypeCheck[bool](bat.Vecs[0], 0)) - require.Equal(t, int8(-8), vector.GetFixedAtNoTypeCheck[int8](bat.Vecs[1], 0)) - require.Equal(t, int16(-16), vector.GetFixedAtNoTypeCheck[int16](bat.Vecs[2], 0)) - require.Equal(t, int32(-32), vector.GetFixedAtNoTypeCheck[int32](bat.Vecs[3], 0)) - require.Equal(t, int64(-64), vector.GetFixedAtNoTypeCheck[int64](bat.Vecs[4], 0)) - require.Equal(t, float32(1.25), vector.GetFixedAtNoTypeCheck[float32](bat.Vecs[5], 0)) - require.Equal(t, 2.5, vector.GetFixedAtNoTypeCheck[float64](bat.Vecs[6], 0)) - require.Equal(t, "tpch", bat.Vecs[7].GetStringAt(0)) - require.Equal(t, types.Decimal64(^uint64(12344)), vector.GetFixedAtNoTypeCheck[types.Decimal64](bat.Vecs[8], 0)) - require.Equal(t, types.Decimal128{B0_63: 7, B64_127: 1}, vector.GetFixedAtNoTypeCheck[types.Decimal128](bat.Vecs[9], 0)) - require.Equal(t, types.DaysFromUnixEpochToDate(1), vector.GetFixedAtNoTypeCheck[types.Date](bat.Vecs[10], 0)) - require.Equal(t, uint32(42), vector.GetFixedAtNoTypeCheck[uint32](bat.Vecs[11], 0)) - bat.Clean(mp) - - charTypes := append([]planpb.Type(nil), typesOut...) - charTypes[7] = planpb.Type{Id: int32(types.T_char), Width: 8} - charSchema, err := ParseSchema(schemaWire, charTypes, headings) - require.NoError(t, err) - charBatch, err := charSchema.decodeRecordBatch(header, body, 1<<20, mp) - require.NoError(t, err) - require.Equal(t, types.T_char, charBatch.Vecs[7].GetType().Oid) - require.Equal(t, "tpch", charBatch.Vecs[7].GetStringAt(0)) - charBatch.Clean(mp) - - requiredTypes := append([]planpb.Type(nil), typesOut...) - requiredTypes[0].NotNullable = true - requiredSchema, err := ParseSchema(schemaWire, requiredTypes, headings) - require.NoError(t, err) - _, err = requiredSchema.decodeRecordBatch(header, body, 1<<20, mp) - require.ErrorContains(t, err, "required field contains nulls") - - require.Equal(t, int64(0), mp.CurrNB()) -} - -func TestArrowIPCRejectsMalformedAndMismatchedData(t *testing.T) { - require.Equal(t, int64(math.MaxInt64/8+1), bitmapBytes(math.MaxInt64)) - schemaWire := mustHex(t, fixtureSchemaHex) - typesOut, headings := fixtureOutputShape() - for length := 0; length < len(schemaWire); length += 31 { - _, err := ParseSchema(schemaWire[:length], typesOut, headings) - require.Error(t, err) - } - _, err := ParseSchema(schemaWire, []planpb.Type{{Id: int32(types.T_int64)}}, []string{"wrong"}) - require.ErrorContains(t, err, "expects 1") - - schema, err := ParseSchema(schemaWire, typesOut, headings) - require.NoError(t, err) - mp := mpool.MustNewZero() - body := mustHex(t, fixtureBodyHex) - header := mustHex(t, fixtureHeaderHex) - for length := 0; length < len(header); length += 17 { - _, err = schema.decodeRecordBatch(header[:length], body, 1<<20, mp) - require.Error(t, err) - } - _, err = schema.decodeRecordBatch(mustHex(t, fixtureHeaderHex), body, 1, mp) - require.ErrorContains(t, err, "decoded-memory budget") - - required := *schema - required.fields = append([]arrowField(nil), schema.fields...) - required.fields[0].nullable = false - _, err = required.decodeRecordBatch(mustHex(t, fixtureHeaderHex), body, 1<<20, mp) - require.ErrorContains(t, err, "required field contains nulls") - - _, err = schema.decodeRecordBatch(mustHex(t, fixtureHeaderHex), body[:len(body)-1], 1<<20, mp) - require.ErrorContains(t, err, "body length mismatch") - require.Equal(t, int64(0), mp.CurrNB()) -} - -func TestArrowIPCLowLevelBoundsAndFraming(t *testing.T) { - _, err := ipcMetadata(nil) - require.ErrorContains(t, err, "truncated") - _, err = ipcMetadata([]byte{0xff, 0xff, 0xff, 0xff}) - require.ErrorContains(t, err, "continuation header") - continuation := make([]byte, 9) - binary.LittleEndian.PutUint32(continuation, math.MaxUint32) - binary.LittleEndian.PutUint32(continuation[4:], 2) - _, err = ipcMetadata(continuation) - require.ErrorContains(t, err, "length is invalid") - binary.LittleEndian.PutUint32(continuation[4:], 1) - metadata, err := ipcMetadata(continuation) - require.NoError(t, err) - require.Equal(t, []byte{0}, metadata) - legacy := []byte{1, 0, 0, 0, 7} - metadata, err = ipcMetadata(legacy) - require.NoError(t, err) - require.Equal(t, []byte{7}, metadata) - raw := []byte{0, 0, 0, 0, 9} - metadata, err = ipcMetadata(raw) - require.NoError(t, err) - require.Equal(t, raw, metadata) - - _, err = rootTable(nil) - require.ErrorContains(t, err, "root is truncated") - _, err = rootTable(make([]byte, 4)) - require.ErrorContains(t, err, "root offset") - - data := validFlatTableData() - table, err := rootTable(data) - require.NoError(t, err) - _, ok, err := table.field(10, 1) - require.NoError(t, err) - require.False(t, ok) - require.Equal(t, byte(9), mustByteField(t, table, 10, 9)) - - invalidBoolean := append([]byte(nil), data...) - invalidBoolean[20] = 2 - _, err = (flatTable{data: invalidBoolean, start: 16}).boolField(0, false) - require.ErrorContains(t, err, "boolean is invalid") - - invalidField := append([]byte(nil), data...) - binary.LittleEndian.PutUint16(invalidField[8:], 2) - _, _, err = (flatTable{data: invalidField, start: 16}).field(0, 8) - require.ErrorContains(t, err, "outside its table") - - _, err = (flatTable{data: data, start: 63}).indirect(63) - require.ErrorContains(t, err, "offset is truncated") - _, err = table.indirect(20) - require.ErrorContains(t, err, "offset is invalid") - - vectorData := validFlatTableData() - binary.LittleEndian.PutUint32(vectorData[20:], 8) - binary.LittleEndian.PutUint32(vectorData[28:], 2) - binary.LittleEndian.PutUint16(vectorData[32:], 1) - vectorData[34] = 1 - vectorTable := flatTable{data: vectorData, start: 16} - _, _, _, err = vectorTable.vector(0, 32) - require.ErrorContains(t, err, "invalid bounds") - _, _, err = vectorTable.structVector(0, 2) - require.NoError(t, err) - _, _, err = vectorTable.stringField(0) - require.ErrorContains(t, err, "not terminated") - _, err = vectorTable.tableVector(0) - require.Error(t, err) - - truncatedVector := validFlatTableData() - binary.LittleEndian.PutUint32(truncatedVector[20:], uint32(len(truncatedVector)-20)) - _, _, _, err = (flatTable{data: truncatedVector, start: 16}).vector(0, 1) - require.ErrorContains(t, err, "vector is truncated") - - invalidVTable := validFlatTableData() - binary.LittleEndian.PutUint32(invalidVTable[16:], 32) - _, _, err = (flatTable{data: invalidVTable, start: 16}).vtable() - require.ErrorContains(t, err, "vtable offset") - invalidBounds := validFlatTableData() - binary.LittleEndian.PutUint16(invalidBounds[4:], math.MaxUint16) - _, _, err = (flatTable{data: invalidBounds, start: 16}).vtable() - require.ErrorContains(t, err, "invalid bounds") -} - -func TestArrowSchemaRejectsNegotiationMismatches(t *testing.T) { - schemaWire := mustHex(t, fixtureSchemaHex) - typesOut, headings := fixtureOutputShape() - - _, err := ParseSchema(make([]byte, maxArrowMetadataBytes+1), typesOut, headings) - require.ErrorContains(t, err, "supported bound") - _, err = ParseSchema(schemaWire, nil, nil) - require.ErrorContains(t, err, "empty or inconsistent") - _, err = ParseSchema(schemaWire, typesOut, headings[:len(headings)-1]) - require.ErrorContains(t, err, "empty or inconsistent") - _, err = ParseSchema([]byte{0xff, 0xff, 0xff, 0xff}, typesOut, headings) - require.ErrorContains(t, err, "continuation header") - _, err = ParseSchema([]byte{4, 0, 0, 0}, typesOut, headings) - require.ErrorContains(t, err, "schema message") - - metadata, err := ipcMetadata(schemaWire) - require.NoError(t, err) - message, err := rootTable(metadata) - require.NoError(t, err) - schemaTable, ok, err := message.tableField(2) - require.NoError(t, err) - require.True(t, ok) - fieldTables, err := schemaTable.tableVector(1) - require.NoError(t, err) - require.Len(t, fieldTables, len(typesOut)) - extraTypes := append(append([]planpb.Type(nil), typesOut...), planpb.Type{Id: int32(types.T_int64)}) - extraHeadings := append(append([]string(nil), headings...), "extra") - _, err = ParseSchema(schemaWire, extraTypes, extraHeadings) - require.ErrorContains(t, err, "MatrixOne expects") - wrongHeadings := append([]string(nil), headings...) - wrongHeadings[0] = "wrong" - _, err = ParseSchema(schemaWire, typesOut, wrongHeadings) - require.ErrorContains(t, err, "is named") - - _, err = parseArrowField(flatTable{}, typesOut[0]) - require.ErrorContains(t, err, "missing name") - _, err = parseArrowField(fieldTables[0], planpb.Type{Id: int32(types.T_bool), NotNullable: true}) - require.NoError(t, err) - for _, tc := range []struct { - index int - expected planpb.Type - }{ - {index: 0, expected: planpb.Type{Id: int32(types.T_varchar)}}, - {index: 1, expected: planpb.Type{Id: int32(types.T_int16)}}, - {index: 5, expected: planpb.Type{Id: int32(types.T_float64)}}, - {index: 7, expected: planpb.Type{Id: int32(types.T_int64)}}, - {index: 8, expected: planpb.Type{Id: int32(types.T_decimal64), Width: 18, Scale: 3}}, - {index: 10, expected: planpb.Type{Id: int32(types.T_int32)}}, - } { - _, err = parseArrowField(fieldTables[tc.index], tc.expected) - require.ErrorContains(t, err, "does not match") - } - - mutated := append([]byte(nil), metadata...) - mutatedMessage, err := rootTable(mutated) - require.NoError(t, err) - mutatedSchema, ok, err := mutatedMessage.tableField(2) - require.NoError(t, err) - require.True(t, ok) - mutatedFields, err := mutatedSchema.tableVector(1) - require.NoError(t, err) - typePosition, ok, err := mutatedFields[0].field(2, 1) - require.NoError(t, err) - require.True(t, ok) - mutated[typePosition] = 99 - _, err = parseArrowField(mutatedFields[0], typesOut[0]) - require.ErrorContains(t, err, "unsupported Arrow type id") - - versionMutation := append([]byte(nil), metadata...) - versionMessage, err := rootTable(versionMutation) - require.NoError(t, err) - versionPosition, ok, err := versionMessage.field(0, 1) - require.NoError(t, err) - require.True(t, ok) - versionMutation[versionPosition] = 3 - _, err = ParseSchema(versionMutation, typesOut, headings) - require.ErrorContains(t, err, "metadata version") - - headerMutation := append([]byte(nil), metadata...) - headerMessage, err := rootTable(headerMutation) - require.NoError(t, err) - headerPosition, ok, err := headerMessage.field(1, 1) - require.NoError(t, err) - require.True(t, ok) - headerMutation[headerPosition] = 0 - _, err = ParseSchema(headerMutation, typesOut, headings) - require.ErrorContains(t, err, "header type") - - missingSchema := append([]byte(nil), metadata...) - missingMessage, err := rootTable(missingSchema) - require.NoError(t, err) - vtable, _, err := missingMessage.vtable() - require.NoError(t, err) - binary.LittleEndian.PutUint16(missingSchema[vtable+8:], 0) - _, err = ParseSchema(missingSchema, typesOut, headings) - require.ErrorContains(t, err, "missing its schema") -} - -func validFlatTableData() []byte { - data := make([]byte, 64) - binary.LittleEndian.PutUint32(data, 16) - binary.LittleEndian.PutUint16(data[4:], 6) - binary.LittleEndian.PutUint16(data[6:], 16) - binary.LittleEndian.PutUint16(data[8:], 4) - binary.LittleEndian.PutUint32(data[16:], 12) - return data -} - -func mustByteField(t *testing.T, table flatTable, index int, defaultValue byte) byte { - t.Helper() - value, err := table.byteField(index, defaultValue) - require.NoError(t, err) - return value -} - -func mustHex(t *testing.T, value string) []byte { - t.Helper() - decoded, err := hex.DecodeString(value) - require.NoError(t, err) - return decoded -} diff --git a/pkg/sql/compile/sidecarflight/client.go b/pkg/sql/compile/sidecarflight/client.go index 86caea65a08cf..ee0d9078df1b4 100644 --- a/pkg/sql/compile/sidecarflight/client.go +++ b/pkg/sql/compile/sidecarflight/client.go @@ -35,14 +35,18 @@ import ( ) const ( - flightService = "/arrow.flight.protocol.FlightService/" - getFlightInfoMethod = flightService + "GetFlightInfo" - doGetMethod = flightService + "DoGet" - doActionMethod = flightService + "DoAction" - commandDescriptor = int32(2) - ticketBytes = 32 - protocolVersion = uint32(3) - substraitVersion = "0.78.0" + flightService = "/arrow.flight.protocol.FlightService/" + getFlightInfoMethod = flightService + "GetFlightInfo" + doGetMethod = flightService + "DoGet" + doPutMethod = flightService + "DoPut" + doActionMethod = flightService + "DoAction" + commandDescriptor = int32(2) + ticketBytes = 32 + protocolVersion = uint32(5) + substraitVersion = "0.78.0" + maxNativeInputBatchBytes = uint64(4 << 20) + maxNativeInputs = 16 + maxPlanBytes = uint64(16 << 20) ) var serverStream = &grpc.StreamDesc{ServerStreams: true} @@ -88,7 +92,7 @@ type Execution struct { runtime *Runtime ticket []byte idempotencyKey []byte - schema *Schema + schema *nativeResultSchema deadline time.Time release func(context.Context) error @@ -104,6 +108,7 @@ type Execution struct { releaseDone chan struct{} releaseErr error reconciling bool + inputs []*NativeInput } // quiescenceUnknownError means the client could not prove that a possibly @@ -139,14 +144,18 @@ func NewRuntime(ctx context.Context, config Config, capabilityDocument string) ( if tlsConfig.MinVersion < tls.VersionTLS12 { tlsConfig.MinVersion = tls.VersionTLS12 } - if config.MaxBatchBytes > uint64(maxInt())-(1<<20) { - return nil, internalErrorf("sidecar flight: max batch bytes overflows platform int") + wirePayloadLimit := max(config.MaxBatchBytes, maxPlanBytes) + if wirePayloadLimit > uint64(maxInt())-(1<<20) { + return nil, internalErrorf("sidecar flight: transport payload limit overflows platform int") } - maximumMessage := config.MaxBatchBytes + 1<<20 + maximumMessage := wirePayloadLimit + 1<<20 conn, err := grpc.NewClient( config.Address, grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig)), - grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(int(maximumMessage))), + grpc.WithDefaultCallOptions( + grpc.MaxCallRecvMsgSize(int(maximumMessage)), + grpc.MaxCallSendMsgSize(int(maximumMessage)), + ), ) if err != nil { return nil, internalErrorf("sidecar flight: create client: %w", err) @@ -190,7 +199,7 @@ func (r *Runtime) Prepare( if release == nil { return nil, internalErrorf("sidecar flight: lease release owner is required") } - if len(queryID) != 16 || len(plan) == 0 || len(plan) > 16<<20 { + if len(queryID) != 16 || len(plan) == 0 || uint64(len(plan)) > maxPlanBytes { primary := internalErrorf("sidecar flight: query identity and Substrait plan are required") return nil, r.failBeforeVisibility(ctx, primary, release) } @@ -198,6 +207,17 @@ func (r *Runtime) Prepare( primary := internalErrorf("sidecar flight: lease-safe execution deadline is required") return nil, r.failBeforeVisibility(ctx, primary, release) } + r.mu.Lock() + stopped := r.stopped + r.mu.Unlock() + if stopped { + primary := internalErrorf("sidecar flight: runtime is stopping") + return nil, r.failBeforeVisibility(ctx, primary, release) + } + schema, schemaWire, err := newNativeResultSchema(outputTypes, headings) + if err != nil { + return nil, r.failBeforeVisibility(ctx, err, release) + } now := time.Now() deadline := now.Add(r.config.RequestTimeout) if callerDeadline, ok := ctx.Deadline(); ok && callerDeadline.Before(deadline) { @@ -237,7 +257,9 @@ func (r *Runtime) Prepare( CapabilityHash: r.capabilityHash[:], MaxBatchBytes: r.config.MaxBatchBytes, DeadlineUnixMS: uint64(deadline.UnixMilli()), Plan: plan, QueryID: append([]byte(nil), queryID...), IdempotencyKey: idempotencyKey[:], - AccountID: proto.Uint64(accountID), + AccountID: proto.Uint64(accountID), + MaxInputBatchBytes: min(maxNativeInputBatchBytes, r.config.MaxBatchBytes), + ResultSchema: schemaWire, } command, err := proto.Marshal(request) if err != nil { @@ -270,8 +292,7 @@ func (r *Runtime) Prepare( return nil, r.failPreparation( internalErrorf("sidecar flight: response capability hash mismatch"), ticket, idempotencyKey[:], release, deadline) } - schema, err := ParseSchema(info.Schema, outputTypes, headings) - if err != nil { + if err = schema.validateWire(info.Schema); err != nil { return nil, r.failPreparation( internalErrorf("sidecar flight: validate schema: %w", err), ticket, idempotencyKey[:], release, deadline) } diff --git a/pkg/sql/compile/sidecarflight/client_test.go b/pkg/sql/compile/sidecarflight/client_test.go index a87454a314fbf..d634b97c34996 100644 --- a/pkg/sql/compile/sidecarflight/client_test.go +++ b/pkg/sql/compile/sidecarflight/client_test.go @@ -15,6 +15,7 @@ package sidecarflight import ( + "bytes" "context" "crypto/ecdsa" "crypto/elliptic" @@ -23,9 +24,11 @@ import ( "crypto/tls" "crypto/x509" "crypto/x509/pkix" + "encoding/binary" "encoding/hex" "encoding/pem" "errors" + "io" "math/big" "net" "os" @@ -40,6 +43,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" planpb "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/perfcounter" "github.com/stretchr/testify/require" @@ -53,6 +57,22 @@ import ( type testFlightService interface{} +// These one-byte sentinels keep older test setup concise. GetFlightInfo now +// replaces schema with the request's negotiated native schema, and DoGet uses +// body only as the signal to emit the native result fixture. +const ( + fixtureSchemaHex = "00" + fixtureHeaderHex = "00" + fixtureBodyHex = "00" +) + +func mustHex(t *testing.T, value string) []byte { + t.Helper() + decoded, err := hex.DecodeString(value) + require.NoError(t, err) + return decoded +} + type testFlightServer struct { schema []byte header []byte @@ -75,6 +95,11 @@ type testFlightServer struct { doGetStarted chan struct{} blockDoGet chan struct{} doGetOnce sync.Once + doPutBatch chan struct{} + blockDoPutAck chan struct{} + doPutOnce sync.Once + doPutBatches atomic.Int32 + notNeededInput bool cancelFailures atomic.Int32 deadlineUnixMS atomic.Int64 } @@ -243,6 +268,9 @@ func TestFlightFallbackClassificationAndWireMessages(t *testing.T) { require.Contains(t, (&flightData{DataHeader: []byte{1}, DataBody: []byte{2}}).String(), "1,1 bytes") require.Contains(t, (&executeSubstraitRequest{Plan: []byte{1}}).String(), "1 bytes") require.Contains(t, (&cancelExecutionRequest{Ticket: []byte{1}}).String(), "1,0 bytes") + require.Contains(t, (&flightPutResult{AppMetadata: []byte{1}}).String(), "1 bytes") + require.Contains(t, (&uploadInputRequest{Ticket: []byte{1}, StreamRef: []byte{2}}).String(), "1,1 bytes") + require.Contains(t, (&uploadInputAck{AcknowledgedBatches: 1, Rows: 2, Bytes: 3}).String(), "1,false,false,false") } func TestExecutionStreamsOneOwnedBatchAndCancelsOnWriterFailure(t *testing.T) { @@ -286,6 +314,387 @@ func TestExecutionStreamsOneOwnedBatchAndCancelsOnWriterFailure(t *testing.T) { require.NoError(t, runtime.Close(cleanupCtx)) } +func TestNativeInputStreamsOneConsumedBatchAtATime(t *testing.T) { + server := &testFlightServer{ + schema: mustHex(t, fixtureSchemaHex), ticket: make([]byte, ticketBytes), hash: make([]byte, sha256.Size), + } + runtime := &Runtime{ + config: Config{MaxBatchBytes: 1 << 20, RequestTimeout: time.Second, CleanupTimeout: time.Second}, + conn: testFlightConnection(t, server), + executions: make(map[*Execution]struct{}), + } + copy(runtime.capabilityHash[:], server.hash) + typesOut, headings := fixtureOutputShape() + execution, err := runtime.Prepare(context.Background(), 1, make([]byte, 16), []byte("plan"), + typesOut, headings, testFlightDeadline(), testFlightRelease) + require.NoError(t, err) + input, err := execution.NewNativeInput(bytes.Repeat([]byte{7}, 32)) + require.NoError(t, err) + _, err = execution.NewNativeInput(bytes.Repeat([]byte{7}, 32)) + require.ErrorContains(t, err, "duplicate native input identity") + require.NoError(t, input.Start(context.Background())) + mp := mpool.MustNewZero() + bat := batch.NewWithSize(1) + bat.Attrs = []string{"value"} + bat.Vecs[0] = vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixed(bat.Vecs[0], int64(42), false, mp)) + bat.SetRowCount(1) + defer bat.Clean(mp) + require.NoError(t, input.Send(context.Background(), bat, mp)) + require.NoError(t, input.Finish(context.Background())) + require.NoError(t, input.Err()) +} + +func TestNativeInputRejectsInvalidAndTerminalStates(t *testing.T) { + ref := bytes.Repeat([]byte{7}, 32) + _, err := (*Execution)(nil).NewNativeInput(ref) + require.ErrorContains(t, err, "invalid native input identity") + _, err = (&Execution{}).NewNativeInput(ref) + require.ErrorContains(t, err, "invalid native input identity") + _, err = (&Execution{runtime: &Runtime{}, ticket: make([]byte, ticketBytes)}).NewNativeInput(ref[:31]) + require.ErrorContains(t, err, "invalid native input identity") + + for _, tc := range []struct { + name string + configure func(*Execution) + }{ + {name: "started", configure: func(e *Execution) { e.started = true }}, + {name: "cleanup", configure: func(e *Execution) { e.cleanupRunning = true }}, + {name: "terminal", configure: func(e *Execution) { e.terminal = true }}, + {name: "quiesced", configure: func(e *Execution) { e.quiesced = true }}, + } { + t.Run(tc.name, func(t *testing.T) { + execution := &Execution{runtime: &Runtime{}, ticket: make([]byte, ticketBytes)} + tc.configure(execution) + _, err := execution.NewNativeInput(ref) + require.ErrorContains(t, err, "no longer accepts native inputs") + }) + } + + full := &Execution{ + runtime: &Runtime{}, ticket: make([]byte, ticketBytes), + inputs: make([]*NativeInput, maxNativeInputs), + } + _, err = full.NewNativeInput(ref) + require.ErrorContains(t, err, "count exceeds") + + execution := &Execution{runtime: &Runtime{}, ticket: make([]byte, ticketBytes)} + input, err := execution.NewNativeInput(ref) + require.NoError(t, err) + ref[0] = 0 + require.Equal(t, byte(7), input.streamRef[0]) + + require.ErrorContains(t, (*NativeInput)(nil).Start(context.Background()), "nil native input") + require.NoError(t, (*NativeInput)(nil).Send(context.Background(), nil, nil)) + require.NoError(t, (*NativeInput)(nil).Finish(context.Background())) + require.NoError(t, (*NativeInput)(nil).Err()) + require.False(t, (*NativeInput)(nil).NotNeeded()) + (*NativeInput)(nil).Abort(nil) + + notNeeded := &NativeInput{notNeeded: true} + require.NoError(t, notNeeded.Start(context.Background())) + require.True(t, notNeeded.NotNeeded()) + + terminalCause := errors.New("terminal input") + terminal := &NativeInput{finished: true, terminalErr: terminalCause} + require.ErrorIs(t, terminal.Start(context.Background()), terminalCause) + require.ErrorIs(t, terminal.Finish(context.Background()), terminalCause) + require.ErrorIs(t, terminal.Err(), terminalCause) + + aborted := new(NativeInput) + aborted.Abort(nil) + require.ErrorIs(t, aborted.Err(), context.Canceled) + require.ErrorIs(t, aborted.Finish(context.Background()), context.Canceled) +} + +func TestCloneNativeWindowHandlesConstAndRejectsImpossibleRows(t *testing.T) { + mp := mpool.MustNewZero() + + withNil := batch.NewWithSize(1) + withNil.SetRowCount(1) + _, err := cloneNativeWindow(withNil, 0, 1, mp) + require.ErrorContains(t, err, "nil vector") + + constVec, err := vector.NewConstFixed(types.T_int64.ToType(), int64(42), 4, mp) + require.NoError(t, err) + constant := batch.NewWithSize(1) + constant.Vecs[0] = constVec + constant.SetRowCount(4) + window, err := cloneNativeWindow(constant, 1, 3, mp) + require.NoError(t, err) + require.Equal(t, 2, window.RowCount()) + require.True(t, window.Vecs[0].IsConst()) + require.Equal(t, 2, window.Vecs[0].Length()) + window.Clean(mp) + constant.Clean(mp) + + largeVec, err := vector.NewConstBytes(types.T_varchar.ToType(), bytes.Repeat([]byte{'x'}, 1024), 1, mp) + require.NoError(t, err) + large := batch.NewWithSize(1) + large.Vecs[0] = largeVec + large.SetRowCount(1) + require.ErrorContains(t, new(NativeInput).sendSplitLocked(large, 1, mp), "one native input row exceeds") + large.Clean(mp) + require.Equal(t, int64(0), mp.CurrNB()) +} + +func TestNativeInputAbortCancelsBlockedAcknowledgement(t *testing.T) { + server := &testFlightServer{ + schema: mustHex(t, fixtureSchemaHex), ticket: make([]byte, ticketBytes), hash: make([]byte, sha256.Size), + doPutBatch: make(chan struct{}), blockDoPutAck: make(chan struct{}), + } + runtime := &Runtime{ + config: Config{MaxBatchBytes: 128, RequestTimeout: time.Second, CleanupTimeout: time.Second}, + conn: testFlightConnection(t, server), executions: make(map[*Execution]struct{}), + } + copy(runtime.capabilityHash[:], server.hash) + typesOut, headings := fixtureOutputShape() + execution, err := runtime.Prepare(context.Background(), 1, make([]byte, 16), []byte("plan"), + typesOut, headings, testFlightDeadline(), testFlightRelease) + require.NoError(t, err) + input, err := execution.NewNativeInput(bytes.Repeat([]byte{7}, 32)) + require.NoError(t, err) + require.NoError(t, input.Start(context.Background())) + + mp := mpool.MustNewZero() + bat := batch.NewWithSize(1) + bat.Vecs[0] = vector.NewVec(types.T_int64.ToType()) + for row := int64(0); row < 20; row++ { + require.NoError(t, vector.AppendFixed(bat.Vecs[0], row, false, mp)) + } + bat.SetRowCount(20) + defer bat.Clean(mp) + sendDone := make(chan error, 1) + go func() { sendDone <- input.Send(context.Background(), bat, mp) }() + select { + case <-server.doPutBatch: + case <-time.After(time.Second): + t.Fatal("native input batch did not reach the server") + } + input.Abort(errors.New("injected cancellation")) + select { + case sendErr := <-sendDone: + require.Error(t, sendErr) + case <-time.After(time.Second): + t.Fatal("Abort did not release the blocked acknowledgement") + } +} + +func TestNativeInputRetireCancelsBlockedAcknowledgementWithoutError(t *testing.T) { + server := &testFlightServer{ + schema: mustHex(t, fixtureSchemaHex), ticket: make([]byte, ticketBytes), hash: make([]byte, sha256.Size), + doPutBatch: make(chan struct{}), blockDoPutAck: make(chan struct{}), + } + runtime := &Runtime{ + config: Config{MaxBatchBytes: 128, RequestTimeout: time.Second, CleanupTimeout: time.Second}, + conn: testFlightConnection(t, server), executions: make(map[*Execution]struct{}), + } + copy(runtime.capabilityHash[:], server.hash) + typesOut, headings := fixtureOutputShape() + execution, err := runtime.Prepare(context.Background(), 1, make([]byte, 16), []byte("plan"), + typesOut, headings, testFlightDeadline(), testFlightRelease) + require.NoError(t, err) + input, err := execution.NewNativeInput(bytes.Repeat([]byte{7}, 32)) + require.NoError(t, err) + require.NoError(t, input.Start(context.Background())) + + mp := mpool.MustNewZero() + bat := batch.NewWithSize(1) + bat.Vecs[0] = vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixed(bat.Vecs[0], int64(1), false, mp)) + bat.SetRowCount(1) + defer bat.Clean(mp) + sendDone := make(chan error, 1) + go func() { sendDone <- input.Send(context.Background(), bat, mp) }() + select { + case <-server.doPutBatch: + case <-time.After(time.Second): + t.Fatal("native input batch did not reach the server") + } + input.Retire() + select { + case sendErr := <-sendDone: + require.NoError(t, sendErr) + case <-time.After(time.Second): + t.Fatal("successful retirement did not release the blocked acknowledgement") + } + require.True(t, input.NotNeeded()) + require.NoError(t, input.Err()) +} + +func TestNativeInputSplitsOversizedBatchesWithinNegotiatedLimit(t *testing.T) { + server := &testFlightServer{ + schema: mustHex(t, fixtureSchemaHex), ticket: make([]byte, ticketBytes), hash: make([]byte, sha256.Size), + } + runtime := &Runtime{ + config: Config{MaxBatchBytes: 128, RequestTimeout: time.Second, CleanupTimeout: time.Second}, + conn: testFlightConnection(t, server), executions: make(map[*Execution]struct{}), + } + copy(runtime.capabilityHash[:], server.hash) + typesOut, headings := fixtureOutputShape() + execution, err := runtime.Prepare(context.Background(), 1, make([]byte, 16), []byte("plan"), + typesOut, headings, testFlightDeadline(), testFlightRelease) + require.NoError(t, err) + input, err := execution.NewNativeInput(bytes.Repeat([]byte{7}, 32)) + require.NoError(t, err) + require.NoError(t, input.Start(context.Background())) + + mp := mpool.MustNewZero() + bat := batch.NewWithSize(1) + bat.Vecs[0] = vector.NewVec(types.T_int64.ToType()) + for row := int64(0); row < 20; row++ { + require.NoError(t, vector.AppendFixed(bat.Vecs[0], row, false, mp)) + } + bat.SetRowCount(20) + defer bat.Clean(mp) + require.NoError(t, input.Send(context.Background(), bat, mp)) + require.NoError(t, input.Finish(context.Background())) + require.Greater(t, input.sequence, uint64(1)) + require.Equal(t, uint64(20), input.rows) + require.LessOrEqual(t, input.bytes, input.sequence*runtime.config.MaxBatchBytes) +} + +func TestNativeInputAcceptsEarlyNotNeeded(t *testing.T) { + server := &testFlightServer{ + schema: mustHex(t, fixtureSchemaHex), ticket: make([]byte, ticketBytes), hash: make([]byte, sha256.Size), + notNeededInput: true, + } + runtime := &Runtime{ + config: Config{MaxBatchBytes: 128, RequestTimeout: time.Second, CleanupTimeout: time.Second}, + conn: testFlightConnection(t, server), executions: make(map[*Execution]struct{}), + } + copy(runtime.capabilityHash[:], server.hash) + typesOut, headings := fixtureOutputShape() + execution, err := runtime.Prepare(context.Background(), 1, make([]byte, 16), []byte("plan"), + typesOut, headings, testFlightDeadline(), testFlightRelease) + require.NoError(t, err) + input, err := execution.NewNativeInput(bytes.Repeat([]byte{7}, 32)) + require.NoError(t, err) + require.NoError(t, input.Start(context.Background())) + + mp := mpool.MustNewZero() + bat := batch.NewWithSize(1) + bat.Vecs[0] = vector.NewVec(types.T_int64.ToType()) + for row := int64(0); row < 20; row++ { + require.NoError(t, vector.AppendFixed(bat.Vecs[0], row, false, mp)) + } + bat.SetRowCount(20) + defer bat.Clean(mp) + require.NoError(t, input.Send(context.Background(), bat, mp)) + require.True(t, input.notNeeded) + require.Equal(t, int32(1), server.doPutBatches.Load()) + require.NoError(t, input.Send(context.Background(), bat, mp)) + require.Equal(t, int32(1), server.doPutBatches.Load()) + require.NoError(t, input.Finish(context.Background())) +} + +func TestCleanupAfterResultEOFStillJoinsNativeInputs(t *testing.T) { + server := &testFlightServer{ + schema: mustHex(t, fixtureSchemaHex), ticket: make([]byte, ticketBytes), hash: make([]byte, sha256.Size), + } + runtime := &Runtime{ + config: Config{MaxBatchBytes: 1 << 20, RequestTimeout: time.Second, CleanupTimeout: time.Second}, + conn: testFlightConnection(t, server), executions: make(map[*Execution]struct{}), + } + copy(runtime.capabilityHash[:], server.hash) + typesOut, headings := fixtureOutputShape() + execution, err := runtime.Prepare(context.Background(), 1, make([]byte, 16), []byte("plan"), + typesOut, headings, testFlightDeadline(), testFlightRelease) + require.NoError(t, err) + input, err := execution.NewNativeInput(bytes.Repeat([]byte{7}, 32)) + require.NoError(t, err) + require.NoError(t, input.Start(context.Background())) + execution.mu.Lock() + execution.quiesced = true + execution.mu.Unlock() + require.NoError(t, execution.Cleanup(context.Background())) + require.ErrorIs(t, input.Err(), context.Canceled) + require.Equal(t, int32(1), server.cancels.Load()) +} + +func TestExecutionRejectsCompressedNativeResultBeforeFill(t *testing.T) { + mp := mpool.MustNewZero() + payload := func() []byte { + vec, err := vector.NewConstFixed(types.T_int64.ToType(), int64(7), 1, mp) + require.NoError(t, err) + vec.SetLength(1 << 30) + bat := batch.NewWithSize(1) + bat.Vecs[0] = vec + bat.SetRowCount(1 << 30) + defer bat.Clean(mp) + payload, err := bat.MarshalBinary() + require.NoError(t, err) + require.Less(t, len(payload), 1024) + return payload + }() + + server := &testFlightServer{ + schema: mustHex(t, fixtureSchemaHex), ticket: make([]byte, ticketBytes), hash: make([]byte, sha256.Size), + streamMessages: []*flightData{ + {DataHeader: []byte{1}}, + {DataHeader: marshalNativeBatchFrame(1, payload)}, + }, + } + runtime := &Runtime{ + config: Config{MaxBatchBytes: 1 << 20, RequestTimeout: time.Second, CleanupTimeout: time.Second}, + conn: testFlightConnection(t, server), executions: make(map[*Execution]struct{}), + } + copy(runtime.capabilityHash[:], server.hash) + execution, err := runtime.Prepare( + context.Background(), 1, make([]byte, 16), []byte("plan"), + []planpb.Type{{Id: int32(types.T_int64)}}, []string{"v"}, testFlightDeadline(), testFlightRelease, + ) + require.NoError(t, err) + fillCalls := 0 + err = execution.Run(context.Background(), mp, nil, func(*batch.Batch, *perfcounter.CounterSet) error { + fillCalls++ + return nil + }) + require.ErrorContains(t, err, "is not flat") + require.Zero(t, fillCalls) + require.Equal(t, int64(0), mp.CurrNB()) + require.NoError(t, execution.Cleanup(context.Background())) +} + +func TestResultEOFRetiresInputBeforeItsFirstBatch(t *testing.T) { + server := &testFlightServer{ + schema: mustHex(t, fixtureSchemaHex), ticket: make([]byte, ticketBytes), hash: make([]byte, sha256.Size), + } + runtime := &Runtime{ + config: Config{MaxBatchBytes: 1 << 20, RequestTimeout: time.Second, CleanupTimeout: time.Second}, + conn: testFlightConnection(t, server), executions: make(map[*Execution]struct{}), + } + copy(runtime.capabilityHash[:], server.hash) + typesOut, headings := fixtureOutputShape() + execution, err := runtime.Prepare(context.Background(), 1, make([]byte, 16), []byte("plan"), + typesOut, headings, testFlightDeadline(), testFlightRelease) + require.NoError(t, err) + input, err := execution.NewNativeInput(bytes.Repeat([]byte{7}, 32)) + require.NoError(t, err) + require.NoError(t, input.Start(context.Background())) + + mp := mpool.MustNewZero() + runDone := make(chan error, 1) + go func() { + runDone <- execution.Run(context.Background(), mp, nil, + func(*batch.Batch, *perfcounter.CounterSet) error { return nil }) + }() + select { + case runErr := <-runDone: + require.NoError(t, runErr) + case <-time.After(time.Second): + t.Fatal("result EOF waited for an input batch that the sidecar did not need") + } + require.True(t, input.NotNeeded()) + require.NoError(t, input.Finish(context.Background())) + require.NoError(t, input.Err()) + require.Zero(t, server.doPutBatches.Load()) + + require.NoError(t, execution.CleanupAfterRun(context.Background(), nil)) + require.Equal(t, int32(1), server.cancels.Load(), + "streamed success must join the server-side DoPut handler before cleanup") +} + func TestPrepareCancelsByIdempotencyWhenTicketIsUnknown(t *testing.T) { server := &testFlightServer{ schema: mustHex(t, fixtureSchemaHex), hash: make([]byte, sha256.Size), badInfo: true, @@ -514,9 +923,8 @@ func TestContextCancellationReachesSidecarWhileWriterIsBlocked(t *testing.T) { } func TestExecutionRejectsMalformedStreamsAndDuplicateClaims(t *testing.T) { - schemaWire := mustHex(t, fixtureSchemaHex) typesOut, headings := fixtureOutputShape() - schema, err := ParseSchema(schemaWire, typesOut, headings) + schema, schemaWire, err := newNativeResultSchema(typesOut, headings) require.NoError(t, err) for _, tc := range []struct { name string @@ -524,12 +932,12 @@ func TestExecutionRejectsMalformedStreamsAndDuplicateClaims(t *testing.T) { maximum uint64 want string }{ - {name: "missing schema", messages: []*flightData{}, maximum: 1 << 20, want: "before its schema"}, - {name: "empty header", messages: []*flightData{{}}, maximum: 1 << 20, want: "malformed or oversized"}, - {name: "schema body", messages: []*flightData{{DataHeader: schemaWire, DataBody: []byte{1}}}, maximum: 1 << 20, want: "schema message contains a body"}, - {name: "schema mismatch", messages: []*flightData{{DataHeader: append([]byte(nil), schemaWire[:len(schemaWire)-1]...)}}, maximum: 1 << 20, want: "stream schema"}, - {name: "oversized body", messages: []*flightData{{DataHeader: schemaWire}, {DataHeader: []byte{1}, DataBody: []byte{1, 2}}}, maximum: 1, want: "malformed or oversized"}, - {name: "invalid batch", messages: []*flightData{{DataHeader: schemaWire}, {DataHeader: []byte{1}}}, maximum: 1 << 20, want: "decode record batch"}, + {name: "missing schema", messages: []*flightData{}, maximum: 1 << 20, want: "before its Flight transport schema"}, + {name: "empty frame", messages: []*flightData{{}}, maximum: 1 << 20, want: "malformed Flight transport schema"}, + {name: "schema body", messages: []*flightData{{DataHeader: []byte{1}, DataBody: []byte{1}}}, maximum: 1 << 20, want: "malformed Flight transport schema"}, + {name: "schema metadata", messages: []*flightData{{AppMetadata: schemaWire}}, maximum: 1 << 20, want: "malformed Flight transport schema"}, + {name: "oversized frame", messages: []*flightData{{DataHeader: []byte{1}}, {DataHeader: make([]byte, nativeBatchFrameHeaderBytes+2)}}, maximum: 1, want: "malformed or oversized"}, + {name: "invalid batch", messages: []*flightData{{DataHeader: []byte{1}}, {DataHeader: marshalNativeBatchFrame(1, []byte{1})}}, maximum: 1 << 20, want: "decode MO native result batch"}, } { t.Run(tc.name, func(t *testing.T) { server := &testFlightServer{ @@ -577,6 +985,57 @@ func fixtureOutputShape() ([]planpb.Type, []string) { }, []string{"b", "i8", "i16", "i32", "i64", "f32", "f64", "s", "d64", "d128", "date", "u32_transport"} } +func fixtureNativeResultPayload() []byte { + mp := mpool.MustNewZero() + bat := batch.NewWithSize(12) + typesOut, _ := fixtureOutputShape() + for i := range typesOut { + typ := types.T(typesOut[i].Id).ToType() + typ.Width = typesOut[i].Width + typ.Scale = typesOut[i].Scale + typ.Charset = uint8(typesOut[i].Charset) + typ.SetNotNull(typesOut[i].NotNullable) + bat.Vecs[i] = vector.NewVec(typ) + } + must := func(err error) { + if err != nil { + panic(err) + } + } + must(vector.AppendFixed(bat.Vecs[0], true, false, mp)) + must(vector.AppendFixed(bat.Vecs[0], false, true, mp)) + must(vector.AppendFixed(bat.Vecs[1], int8(-8), false, mp)) + must(vector.AppendFixed(bat.Vecs[1], int8(0), true, mp)) + must(vector.AppendFixed(bat.Vecs[2], int16(-16), false, mp)) + must(vector.AppendFixed(bat.Vecs[2], int16(0), true, mp)) + must(vector.AppendFixed(bat.Vecs[3], int32(-32), false, mp)) + must(vector.AppendFixed(bat.Vecs[3], int32(0), true, mp)) + must(vector.AppendFixed(bat.Vecs[4], int64(-64), false, mp)) + must(vector.AppendFixed(bat.Vecs[4], int64(0), true, mp)) + must(vector.AppendFixed(bat.Vecs[5], float32(1.25), false, mp)) + must(vector.AppendFixed(bat.Vecs[5], float32(0), true, mp)) + must(vector.AppendFixed(bat.Vecs[6], 2.5, false, mp)) + must(vector.AppendFixed(bat.Vecs[6], float64(0), true, mp)) + must(vector.AppendBytes(bat.Vecs[7], []byte("tpch"), false, mp)) + must(vector.AppendBytes(bat.Vecs[7], nil, true, mp)) + must(vector.AppendFixed(bat.Vecs[8], types.Decimal64(^uint64(12344)), false, mp)) + must(vector.AppendFixed(bat.Vecs[8], types.Decimal64(0), true, mp)) + must(vector.AppendFixed(bat.Vecs[9], types.Decimal128{B0_63: 7, B64_127: 1}, false, mp)) + must(vector.AppendFixed(bat.Vecs[9], types.Decimal128{}, true, mp)) + must(vector.AppendFixed(bat.Vecs[10], types.DaysFromUnixEpochToDate(1), false, mp)) + must(vector.AppendFixed(bat.Vecs[10], types.Date(0), true, mp)) + must(vector.AppendFixed(bat.Vecs[11], uint32(42), false, mp)) + must(vector.AppendFixed(bat.Vecs[11], uint32(0), true, mp)) + bat.SetRowCount(2) + payload, err := bat.MarshalBinary() + bat.Clean(mp) + must(err) + if mp.CurrNB() != 0 { + panic("native result fixture leaked its memory pool") + } + return payload +} + func testFlightConnection(t *testing.T, implementation *testFlightServer) *grpc.ClientConn { t.Helper() listener := bufconn.Listen(1 << 20) @@ -587,6 +1046,7 @@ func testFlightConnection(t *testing.T, implementation *testFlightServer) *grpc. Methods: []grpc.MethodDesc{{MethodName: "GetFlightInfo", Handler: testGetFlightInfo}}, Streams: []grpc.StreamDesc{ {StreamName: "DoGet", Handler: testDoGet, ServerStreams: true}, + {StreamName: "DoPut", Handler: testDoPut, ServerStreams: true, ClientStreams: true}, {StreamName: "DoAction", Handler: testDoAction, ServerStreams: true}, }, }, implementation) @@ -610,6 +1070,7 @@ func testGetFlightInfo(service any, ctx context.Context, decode func(any) error, if request.Type != commandDescriptor || len(request.Path) != 0 || proto.Unmarshal(request.Cmd, command) != nil || command.ProtocolVersion != protocolVersion || command.SubstraitVersion != substraitVersion || len(command.CapabilityHash) != sha256.Size || command.MaxBatchBytes == 0 || command.DeadlineUnixMS == 0 || + command.MaxInputBatchBytes == 0 || len(command.ResultSchema) == 0 || len(command.Plan) == 0 || len(command.QueryID) != 16 || len(command.IdempotencyKey) != sha256.Size || command.AccountID == nil { return nil, status.Error(codes.InvalidArgument, "malformed ExecuteSubstrait command") } @@ -627,6 +1088,9 @@ func testGetFlightInfo(service any, ctx context.Context, decode func(any) error, if server.badInfo { return &flightInfo{Schema: server.schema, AppMetadata: server.hash}, nil } + if !bytes.Equal(server.schema, []byte("bad")) { + server.schema = bytes.Clone(command.ResultSchema) + } return &flightInfo{Schema: server.schema, Endpoint: []*flightEndpoint{{ Ticket: &flightTicket{Ticket: server.ticket}, Locations: server.locations, }}, AppMetadata: server.hash}, nil @@ -656,10 +1120,72 @@ func testDoGet(service any, stream grpc.ServerStream) error { } return nil } - if err := stream.SendMsg(&flightData{DataHeader: server.schema}); err != nil { + if err := stream.SendMsg(&flightData{DataHeader: []byte{1}}); err != nil { + return err + } + if len(server.body) == 0 { + return nil + } + return stream.SendMsg(&flightData{DataHeader: marshalNativeBatchFrame(1, fixtureNativeResultPayload())}) +} + +func testDoPut(service any, stream grpc.ServerStream) error { + server := service.(*testFlightServer) + first := new(flightData) + if err := stream.RecvMsg(first); err != nil { + return err + } + request := new(uploadInputRequest) + if first.Descriptor == nil || first.Descriptor.Type != commandDescriptor || + proto.Unmarshal(first.Descriptor.Cmd, request) != nil || len(request.Ticket) != ticketBytes || + len(request.StreamRef) != 32 { + return status.Error(codes.InvalidArgument, "malformed native input descriptor") + } + attached, _ := proto.Marshal(&uploadInputAck{Ready: true}) + if err := stream.SendMsg(&flightPutResult{AppMetadata: attached}); err != nil { return err } - return stream.SendMsg(&flightData{DataHeader: server.header, DataBody: server.body}) + var batches, rows, bytesSeen uint64 + for { + message := new(flightData) + err := stream.RecvMsg(message) + if err == io.EOF { + ack, _ := proto.Marshal(&uploadInputAck{AcknowledgedBatches: batches, Rows: rows, Bytes: bytesSeen, Complete: true}) + return stream.SendMsg(&flightPutResult{AppMetadata: ack}) + } + if err != nil { + return err + } + if len(message.AppMetadata) < nativeBatchFrameHeaderBytes || string(message.AppMetadata[:4]) != "MOB1" { + return status.Error(codes.InvalidArgument, "malformed native input frame") + } + batches++ + server.doPutBatches.Add(1) + if binary.LittleEndian.Uint64(message.AppMetadata[8:16]) != batches { + return status.Error(codes.InvalidArgument, "non-contiguous native input sequence") + } + payload := message.AppMetadata[nativeBatchFrameHeaderBytes:] + if server.notNeededInput { + ack, _ := proto.Marshal(&uploadInputAck{Complete: true, NotNeeded: true}) + return stream.SendMsg(&flightPutResult{AppMetadata: ack}) + } + rows += binary.LittleEndian.Uint64(payload[:8]) + bytesSeen += uint64(len(payload)) + if server.doPutBatch != nil { + server.doPutOnce.Do(func() { close(server.doPutBatch) }) + } + if server.blockDoPutAck != nil { + select { + case <-server.blockDoPutAck: + case <-stream.Context().Done(): + return context.Cause(stream.Context()) + } + } + ack, _ := proto.Marshal(&uploadInputAck{AcknowledgedBatches: batches, Rows: rows, Bytes: bytesSeen}) + if err = stream.SendMsg(&flightPutResult{AppMetadata: ack}); err != nil { + return err + } + } } func testDoAction(service any, stream grpc.ServerStream) error { @@ -745,6 +1271,7 @@ func testTLSFlightServer(t *testing.T, implementation *testFlightServer) (string Methods: []grpc.MethodDesc{{MethodName: "GetFlightInfo", Handler: testGetFlightInfo}}, Streams: []grpc.StreamDesc{ {StreamName: "DoGet", Handler: testDoGet, ServerStreams: true}, + {StreamName: "DoPut", Handler: testDoPut, ServerStreams: true, ClientStreams: true}, {StreamName: "DoAction", Handler: testDoAction, ServerStreams: true}, }, }, implementation) diff --git a/pkg/sql/compile/sidecarflight/native_input.go b/pkg/sql/compile/sidecarflight/native_input.go new file mode 100644 index 0000000000000..f005617a94d39 --- /dev/null +++ b/pkg/sql/compile/sidecarflight/native_input.go @@ -0,0 +1,467 @@ +// Copyright 2026 Matrix Origin +// +// 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 sidecarflight + +import ( + "bytes" + "context" + "encoding/binary" + "errors" + "io" + "sync" + "sync/atomic" + + "github.com/gogo/protobuf/proto" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "google.golang.org/grpc" +) + +var putStream = &grpc.StreamDesc{ServerStreams: true, ClientStreams: true} + +// NativeInput is one single-use, acknowledged MO-batch stream for a StreamRead. +// Send returns only after the sidecar scan has fully consumed that batch. +type NativeInput struct { + execution *Execution + streamRef []byte + + mu sync.Mutex + stream grpc.ClientStream + sequence uint64 + rows uint64 + bytes uint64 + finished bool + notNeeded bool + terminalErr error + // retired is the success-valued terminal signal published by result EOF. + // It is independent from mu so EOF can interrupt a DoPut acknowledgement + // wait before the producer has another frame to send. + retired atomic.Bool + + // cancelMu is intentionally independent from mu. Send and Finish hold mu + // while waiting for a sidecar acknowledgement; Abort must be able to cancel + // that RPC before it waits to publish terminal state under mu. + cancelMu sync.Mutex + cancel context.CancelFunc +} + +func (e *Execution) NewNativeInput(streamRef []byte) (*NativeInput, error) { + if e == nil || e.runtime == nil || len(e.ticket) != ticketBytes || len(streamRef) != 32 { + return nil, internalErrorf("sidecar flight: invalid native input identity") + } + input := &NativeInput{execution: e, streamRef: append([]byte(nil), streamRef...)} + e.mu.Lock() + if e.started || e.cleanupRunning || e.terminal || e.quiesced { + e.mu.Unlock() + return nil, internalErrorf("sidecar flight: execution no longer accepts native inputs") + } + if len(e.inputs) >= maxNativeInputs { + e.mu.Unlock() + return nil, internalErrorf("sidecar flight: native input count exceeds the protocol limit") + } + for _, existing := range e.inputs { + if bytes.Equal(existing.streamRef, streamRef) { + e.mu.Unlock() + return nil, internalErrorf("sidecar flight: duplicate native input identity") + } + } + e.inputs = append(e.inputs, input) + e.mu.Unlock() + return input, nil +} + +func (n *NativeInput) open(ctx context.Context) error { + if n.stream != nil { + return nil + } + streamCtx := ctx + if streamCtx == nil { + streamCtx = context.Background() + } + streamCtx, cancel := context.WithCancel(streamCtx) + n.setCancel(cancel) + stream, err := n.execution.runtime.conn.NewStream(streamCtx, putStream, doPutMethod) + if err != nil { + n.cancelStream() + return internalErrorf("sidecar flight: open native input stream: %w", err) + } + request, err := proto.Marshal(&uploadInputRequest{Ticket: n.execution.ticket, StreamRef: n.streamRef}) + if err != nil { + n.cancelStream() + return internalErrorf("sidecar flight: encode native input request: %w", err) + } + if err = stream.SendMsg(&flightData{Descriptor: &flightDescriptor{Type: commandDescriptor, Cmd: request}}); err != nil { + n.cancelStream() + return internalErrorf("sidecar flight: send native input descriptor: %w", err) + } + n.stream = stream + ack, err := n.recvAck() + if err != nil { + n.cancelStream() + return err + } + if ack.AcknowledgedBatches != 0 || ack.Rows != 0 || ack.Bytes != 0 || ack.Complete || ack.NotNeeded || !ack.Ready { + n.cancelStream() + return internalErrorf("sidecar flight: invalid native input attachment acknowledgement") + } + return nil +} + +// Start attaches the input to its prepared execution. All inputs are attached +// before DoGet starts so a sidecar plan that prunes a read cannot retire its +// ticket before the matching DoPut handler exists. +func (n *NativeInput) Start(ctx context.Context) error { + if n == nil { + return internalErrorf("sidecar flight: nil native input") + } + n.mu.Lock() + defer n.mu.Unlock() + if n.retired.Load() || n.notNeeded { + return nil + } + if n.finished || n.terminalErr != nil { + return errors.Join(internalErrorf("sidecar flight: native input is terminal"), n.terminalErr) + } + if err := n.open(ctx); err != nil { + if n.retired.Load() { + return nil + } + n.terminalErr = err + return err + } + return nil +} + +func (n *NativeInput) Send(ctx context.Context, bat *batch.Batch, mp *mpool.MPool) error { + if n == nil || bat == nil || bat.IsEmpty() { + return nil + } + n.mu.Lock() + defer n.mu.Unlock() + if n.retired.Load() || n.notNeeded { + return nil + } + if n.finished || n.terminalErr != nil { + return errors.Join(internalErrorf("sidecar flight: native input is terminal"), n.terminalErr) + } + if err := n.open(ctx); err != nil { + if n.retired.Load() { + return nil + } + n.terminalErr = err + return err + } + if err := bat.CheckLength(); err != nil { + n.terminalErr = err + return err + } + if (len(bat.Attrs) != 0 && len(bat.Attrs) != len(bat.Vecs)) || len(bat.ExtraBuf) != 0 || + bat.Recursive != 0 || bat.ShuffleIDX != 0 { + n.terminalErr = internalErrorf("sidecar flight: native input contains unsupported batch metadata") + return n.terminalErr + } + size, err := bat.MarshalBinarySize() + if err != nil { + n.terminalErr = err + return err + } + limit := min(maxNativeInputBatchBytes, n.execution.runtime.config.MaxBatchBytes) + if uint64(size) > limit { + err = n.sendSplitLocked(bat, limit, mp) + if err != nil { + n.terminalErr = err + } + return err + } + payload, err := bat.MarshalBinary() + if err != nil { + n.terminalErr = err + return err + } + return n.sendPayloadLocked(payload) +} + +func (n *NativeInput) sendPayloadLocked(payload []byte) error { + n.sequence++ + frame := marshalNativeBatchFrame(n.sequence, payload) + if err := n.stream.SendMsg(&flightData{AppMetadata: frame}); err != nil { + if n.retired.Load() { + return nil + } + n.terminalErr = internalErrorf("sidecar flight: send native input batch: %w", err) + return n.terminalErr + } + ack, err := n.recvAck() + if err != nil { + if n.retired.Load() { + return nil + } + n.terminalErr = err + return err + } + if ack.NotNeeded { + acknowledgedCurrent := ack.AcknowledgedBatches == n.sequence + acknowledgedPrevious := ack.AcknowledgedBatches == n.sequence-1 + expectedRows, expectedBytes := n.rows, n.bytes + if acknowledgedCurrent { + expectedRows += binary.LittleEndian.Uint64(payload[:8]) + expectedBytes += uint64(len(payload)) + } + if !ack.Complete || ack.Ready || (!acknowledgedCurrent && !acknowledgedPrevious) || + ack.Rows != expectedRows || ack.Bytes != expectedBytes { + n.terminalErr = internalErrorf("sidecar flight: invalid native input not-needed acknowledgement") + return n.terminalErr + } + var trailing flightPutResult + if err = n.stream.RecvMsg(&trailing); err != io.EOF { + if n.retired.Load() { + return nil + } + n.terminalErr = internalErrorf("sidecar flight: native input not-needed stream has trailing results: %w", err) + return n.terminalErr + } + n.sequence = ack.AcknowledgedBatches + n.rows, n.bytes = ack.Rows, ack.Bytes + n.notNeeded = true + n.finished = true + n.cancelStream() + return nil + } + expectedRows := n.rows + binary.LittleEndian.Uint64(payload[:8]) + expectedBytes := n.bytes + uint64(len(payload)) + if ack.AcknowledgedBatches != n.sequence || ack.Rows != expectedRows || ack.Bytes != expectedBytes || + ack.Complete || ack.NotNeeded || ack.Ready { + n.terminalErr = internalErrorf("sidecar flight: invalid native input acknowledgement") + return n.terminalErr + } + n.rows = ack.Rows + n.bytes = ack.Bytes + return nil +} + +func (n *NativeInput) sendSplitLocked(source *batch.Batch, limit uint64, mp *mpool.MPool) error { + for start := 0; start < source.RowCount(); { + low, high, best := start+1, source.RowCount(), -1 + for low <= high { + middle := low + (high-low)/2 + window, err := cloneNativeWindow(source, start, middle, mp) + if err != nil { + return err + } + size, sizeErr := window.MarshalBinarySize() + window.Clean(mp) + if sizeErr != nil { + return sizeErr + } + if uint64(size) <= limit { + best = middle + low = middle + 1 + } else { + high = middle - 1 + } + } + if best < 0 { + return internalErrorf("sidecar flight: one native input row exceeds the negotiated limit") + } + window, err := cloneNativeWindow(source, start, best, mp) + if err != nil { + return err + } + payload, marshalErr := window.MarshalBinary() + window.Clean(mp) + if marshalErr != nil { + return marshalErr + } + if err = n.sendPayloadLocked(payload); err != nil { + return err + } + if n.retired.Load() || n.notNeeded { + return nil + } + start = best + } + return nil +} + +func cloneNativeWindow(source *batch.Batch, start, end int, mp *mpool.MPool) (*batch.Batch, error) { + result := batch.NewWithSize(len(source.Vecs)) + result.Attrs = append([]string(nil), source.Attrs...) + for i, sourceVec := range source.Vecs { + if sourceVec == nil { + result.Clean(mp) + return nil, internalErrorf("sidecar flight: native input contains a nil vector") + } + cloneStart, cloneEnd := start, end + if sourceVec.IsConst() { + cloneStart, cloneEnd = 0, min(1, sourceVec.Length()) + } + cloned, err := sourceVec.CloneWindow(cloneStart, cloneEnd, mp) + if err != nil { + result.Clean(mp) + return nil, err + } + if sourceVec.IsConst() { + cloned.SetLength(end - start) + } + result.Vecs[i] = cloned + } + result.SetRowCount(end - start) + return result, nil +} + +func (n *NativeInput) Finish(ctx context.Context) error { + if n == nil { + return nil + } + n.mu.Lock() + defer n.mu.Unlock() + if n.retired.Load() { + return nil + } + if n.finished { + return n.terminalErr + } + if n.terminalErr != nil { + n.finished = true + return n.terminalErr + } + if err := n.open(ctx); err != nil { + if n.retired.Load() { + return nil + } + n.terminalErr = err + n.finished = true + return err + } + if err := n.stream.CloseSend(); err != nil { + if n.retired.Load() { + return nil + } + n.terminalErr = internalErrorf("sidecar flight: close native input: %w", err) + n.finished = true + return n.terminalErr + } + ack, err := n.recvAck() + if err != nil { + if n.retired.Load() { + return nil + } + n.terminalErr = err + n.finished = true + return err + } + if !ack.Complete || ack.Ready || ack.AcknowledgedBatches != n.sequence || ack.Rows != n.rows || + ack.Bytes != n.bytes { + n.terminalErr = internalErrorf("sidecar flight: missing final native input acknowledgement") + n.finished = true + return n.terminalErr + } + var trailing flightPutResult + if err = n.stream.RecvMsg(&trailing); err != io.EOF { + if n.retired.Load() { + return nil + } + n.terminalErr = internalErrorf("sidecar flight: native input stream has trailing results: %w", err) + n.finished = true + return n.terminalErr + } + n.rows, n.bytes = ack.Rows, ack.Bytes + n.finished = true + n.cancelStream() + return nil +} + +// Retire publishes successful result-side EOF to the producer and interrupts +// any blocked DoPut operation. Unlike Abort it does not manufacture a query +// error: the sidecar has already produced the complete result and no longer +// consumes this input. +func (n *NativeInput) Retire() { + if n == nil { + return + } + n.retired.Store(true) + n.cancelStream() +} + +func (n *NativeInput) recvAck() (*uploadInputAck, error) { + result := new(flightPutResult) + if err := n.stream.RecvMsg(result); err != nil { + return nil, internalErrorf("sidecar flight: receive native input acknowledgement: %w", err) + } + ack := new(uploadInputAck) + if len(result.AppMetadata) == 0 || proto.Unmarshal(result.AppMetadata, ack) != nil { + return nil, internalErrorf("sidecar flight: malformed native input acknowledgement") + } + return ack, nil +} + +func (n *NativeInput) Abort(cause error) { + if n == nil { + return + } + n.cancelStream() + if n.retired.Load() { + return + } + if cause == nil { + cause = context.Canceled + } + n.mu.Lock() + defer n.mu.Unlock() + if n.retired.Load() { + return + } + if n.terminalErr == nil { + n.terminalErr = cause + } + n.finished = true +} + +func (n *NativeInput) setCancel(cancel context.CancelFunc) { + n.cancelMu.Lock() + n.cancel = cancel + n.cancelMu.Unlock() +} + +func (n *NativeInput) cancelStream() { + n.cancelMu.Lock() + cancel := n.cancel + n.cancel = nil + n.cancelMu.Unlock() + if cancel != nil { + cancel() + } +} + +func (n *NativeInput) Err() error { + if n == nil { + return nil + } + n.mu.Lock() + defer n.mu.Unlock() + return n.terminalErr +} + +func (n *NativeInput) NotNeeded() bool { + if n == nil { + return false + } + if n.retired.Load() { + return true + } + n.mu.Lock() + defer n.mu.Unlock() + return n.notNeeded +} diff --git a/pkg/sql/compile/sidecarflight/native_result.go b/pkg/sql/compile/sidecarflight/native_result.go new file mode 100644 index 0000000000000..bf7fcb8b4af8d --- /dev/null +++ b/pkg/sql/compile/sidecarflight/native_result.go @@ -0,0 +1,206 @@ +// Copyright 2026 Matrix Origin +// +// 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 sidecarflight + +import ( + "bytes" + "encoding/binary" + "math" + + "github.com/gogo/protobuf/proto" + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + planpb "github.com/matrixorigin/matrixone/pkg/pb/plan" +) + +const ( + nativeBatchFrameHeaderBytes = 24 + nativeResultSchemaVersion = uint32(1) + maxNativeResultSchemaBytes = 1 << 20 + maxNativeResultColumns = 4096 +) + +type nativeResultSchema struct { + Version uint32 `protobuf:"varint,1,opt,name=version,proto3"` + Columns []*nativeResultColumn `protobuf:"bytes,2,rep,name=columns,proto3"` +} + +func (s *nativeResultSchema) Reset() { *s = nativeResultSchema{} } +func (s *nativeResultSchema) String() string { return proto.CompactTextString(s) } +func (s *nativeResultSchema) ProtoMessage() {} + +type nativeResultColumn struct { + Name string `protobuf:"bytes,1,opt,name=name,proto3"` + Oid uint32 `protobuf:"varint,2,opt,name=oid,proto3"` + Width int32 `protobuf:"varint,3,opt,name=width,proto3"` + Scale int32 `protobuf:"varint,4,opt,name=scale,proto3"` + Charset uint32 `protobuf:"varint,5,opt,name=charset,proto3"` + NotNullable bool `protobuf:"varint,6,opt,name=not_nullable,json=notNullable,proto3"` +} + +func (m *nativeResultColumn) Reset() { *m = nativeResultColumn{} } +func (m *nativeResultColumn) String() string { return proto.CompactTextString(m) } +func (m *nativeResultColumn) ProtoMessage() {} + +func newNativeResultSchema(expected []planpb.Type, headings []string) (*nativeResultSchema, []byte, error) { + if len(expected) == 0 || len(expected) > maxNativeResultColumns || len(headings) != len(expected) { + return nil, nil, internalErrorf("sidecar flight: MatrixOne result schema is empty or inconsistent") + } + result := &nativeResultSchema{ + Version: nativeResultSchemaVersion, + Columns: make([]*nativeResultColumn, len(expected)), + } + for i := range expected { + column, err := nativeResultColumnFromPlan(headings[i], expected[i]) + if err != nil { + return nil, nil, internalErrorf("sidecar flight: result column %d: %w", i, err) + } + result.Columns[i] = column + } + wire, err := proto.Marshal(result) + if err != nil { + return nil, nil, internalErrorf("sidecar flight: encode native result schema: %w", err) + } + if len(wire) == 0 || len(wire) > maxNativeResultSchemaBytes { + return nil, nil, internalErrorf("sidecar flight: native result schema exceeds the supported bound") + } + return result, wire, nil +} + +func nativeResultColumnFromPlan(name string, expected planpb.Type) (*nativeResultColumn, error) { + oid := types.T(expected.Id) + switch oid { + case types.T_bool, + types.T_int8, types.T_int16, types.T_int32, types.T_int64, types.T_uint32, + types.T_float32, types.T_float64, + types.T_date: + case types.T_char, types.T_varchar: + if expected.Width < 0 || expected.Charset > math.MaxUint8 { + return nil, internalErrorf("unsupported MatrixOne result type %s", oid.String()) + } + case types.T_decimal64: + if expected.Width <= 0 || expected.Width > 18 || expected.Scale < 0 || expected.Scale > expected.Width { + return nil, internalErrorf("invalid MatrixOne decimal64(%d,%d)", expected.Width, expected.Scale) + } + case types.T_decimal128: + if expected.Width <= 18 || expected.Width > 38 || expected.Scale < 0 || expected.Scale > expected.Width { + return nil, internalErrorf("invalid MatrixOne decimal128(%d,%d)", expected.Width, expected.Scale) + } + default: + return nil, internalErrorf("unsupported MatrixOne result type %s", oid.String()) + } + return &nativeResultColumn{ + Name: name, Oid: uint32(oid), Width: expected.Width, Scale: expected.Scale, + Charset: expected.Charset, NotNullable: expected.NotNullable, + }, nil +} + +func (s *nativeResultSchema) validateWire(wire []byte) error { + if s == nil { + return internalErrorf("sidecar flight: missing native result schema") + } + expected, err := proto.Marshal(s) + if err != nil { + return internalErrorf("sidecar flight: encode expected native result schema: %w", err) + } + if !bytes.Equal(wire, expected) { + return internalErrorf("sidecar flight: native result schema mismatch") + } + return nil +} + +func marshalNativeBatchFrame(sequence uint64, payload []byte) []byte { + frame := make([]byte, nativeBatchFrameHeaderBytes+len(payload)) + copy(frame[:4], "MOB1") + binary.LittleEndian.PutUint16(frame[4:6], 1) + binary.LittleEndian.PutUint64(frame[8:16], sequence) + binary.LittleEndian.PutUint64(frame[16:24], uint64(len(payload))) + copy(frame[nativeBatchFrameHeaderBytes:], payload) + return frame +} + +func unmarshalNativeBatchFrame(frame []byte, maximum uint64) (uint64, []byte, error) { + if len(frame) < nativeBatchFrameHeaderBytes || string(frame[:4]) != "MOB1" || + binary.LittleEndian.Uint16(frame[4:6]) != 1 || frame[6] != 0 || frame[7] != 0 { + return 0, nil, internalErrorf("sidecar flight: invalid MO native batch frame") + } + sequence := binary.LittleEndian.Uint64(frame[8:16]) + payloadBytes := binary.LittleEndian.Uint64(frame[16:24]) + if sequence == 0 || payloadBytes == 0 || payloadBytes > maximum || + payloadBytes != uint64(len(frame)-nativeBatchFrameHeaderBytes) { + return 0, nil, internalErrorf("sidecar flight: invalid MO native batch frame bounds") + } + return sequence, frame[nativeBatchFrameHeaderBytes:], nil +} + +func (s *nativeResultSchema) decodeBatch(payload []byte, mp *mpool.MPool) (result *batch.Batch, err error) { + if s == nil || mp == nil { + return nil, internalErrorf("sidecar flight: missing native result schema or memory pool") + } + result = batch.NewOffHeapEmpty() + defer func() { + if err != nil && result != nil { + result.Clean(mp) + result = nil + } + }() + if err = result.UnmarshalBinaryWithAnyMp(payload, mp); err != nil { + return nil, internalErrorf("sidecar flight: decode MO native result batch: %w", err) + } + if result.RowCount() <= 0 || len(result.Vecs) != len(s.Columns) || + (len(result.Attrs) != 0 && len(result.Attrs) != len(s.Columns)) || + len(result.ExtraBuf) != 0 || result.Recursive != 0 || result.ShuffleIDX != 0 { + return nil, internalErrorf( + "sidecar flight: MO native result batch metadata mismatch: rows=%d vectors=%d attrs=%d extra=%d recursive=%d shuffle=%d", + result.RowCount(), len(result.Vecs), len(result.Attrs), len(result.ExtraBuf), result.Recursive, result.ShuffleIDX, + ) + } + for _, attr := range result.Attrs { + if attr != "" { + return nil, internalErrorf("sidecar flight: MO native result batch contains unexpected attributes") + } + } + decodedAttrs := result.Attrs + result.Attrs = nil + canonical, marshalErr := result.MarshalBinary() + result.Attrs = decodedAttrs + if marshalErr != nil || !bytes.Equal(canonical, payload) { + return nil, internalErrorf("sidecar flight: MO native result batch is non-canonical or has trailing data") + } + if err = result.CheckLength(); err != nil { + return nil, internalErrorf("sidecar flight: MO native result batch length: %w", err) + } + for i, column := range s.Columns { + vec := result.Vecs[i] + if vec == nil { + return nil, internalErrorf("sidecar flight: MO native result column %d is nil", i) + } + if !vec.IsFlat() { + return nil, internalErrorf("sidecar flight: MO native result column %d is not flat", i) + } + actual := vec.GetType() + expectedSize := types.T(column.Oid).ToType().Size + if uint32(actual.Oid) != column.Oid || actual.Size != expectedSize || + actual.Width != column.Width || actual.Scale != column.Scale || + uint32(actual.Charset) != column.Charset || actual.GetNotNull() != column.NotNullable { + return nil, internalErrorf("sidecar flight: MO native result column %d type mismatch", i) + } + if column.NotNullable && vec.HasNull() { + return nil, internalErrorf("sidecar flight: required MO native result column %d contains nulls", i) + } + } + return result, nil +} diff --git a/pkg/sql/compile/sidecarflight/native_result_test.go b/pkg/sql/compile/sidecarflight/native_result_test.go new file mode 100644 index 0000000000000..9cd441543d4ed --- /dev/null +++ b/pkg/sql/compile/sidecarflight/native_result_test.go @@ -0,0 +1,130 @@ +// Copyright 2026 Matrix Origin +// +// 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 sidecarflight + +import ( + "encoding/binary" + "testing" + + "github.com/matrixorigin/matrixone/pkg/common/mpool" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/stretchr/testify/require" +) + +func TestNativeResultCodecDecodesNegotiatedTypes(t *testing.T) { + typesOut, headings := fixtureOutputShape() + schema, wire, err := newNativeResultSchema(typesOut, headings) + require.NoError(t, err) + require.NoError(t, schema.validateWire(wire)) + + mp := mpool.MustNewZero() + bat, err := schema.decodeBatch(fixtureNativeResultPayload(), mp) + require.NoError(t, err) + require.Equal(t, 2, bat.RowCount()) + for _, vec := range bat.Vecs { + require.False(t, vec.IsNull(0)) + require.True(t, vec.IsNull(1)) + } + require.True(t, vector.GetFixedAtNoTypeCheck[bool](bat.Vecs[0], 0)) + require.Equal(t, "tpch", bat.Vecs[7].GetStringAt(0)) + require.Equal(t, types.DaysFromUnixEpochToDate(1), vector.GetFixedAtNoTypeCheck[types.Date](bat.Vecs[10], 0)) + require.Equal(t, uint32(42), vector.GetFixedAtNoTypeCheck[uint32](bat.Vecs[11], 0)) + bat.Clean(mp) + require.Equal(t, int64(0), mp.CurrNB()) +} + +func TestNativeResultCodecRejectsFramesAndSchemaMismatch(t *testing.T) { + sequence, payload, err := unmarshalNativeBatchFrame(marshalNativeBatchFrame(7, []byte("payload")), 1<<20) + require.NoError(t, err) + require.Equal(t, uint64(7), sequence) + require.Equal(t, []byte("payload"), payload) + + _, _, err = unmarshalNativeBatchFrame([]byte("MOB1"), 1<<20) + require.Error(t, err) + typesOut, headings := fixtureOutputShape() + schema, wire, err := newNativeResultSchema(typesOut, headings) + require.NoError(t, err) + wire[len(wire)-1] ^= 1 + require.ErrorContains(t, schema.validateWire(wire), "schema mismatch") + + mp := mpool.MustNewZero() + trailing := append(fixtureNativeResultPayload(), 0) + _, err = schema.decodeBatch(trailing, mp) + require.ErrorContains(t, err, "non-canonical or has trailing data") + + wrongSize := fixtureNativeResultPayload() + // batch header (12), vector length (4), class (1), then the MO Type's + // four-byte Size field begins four bytes into the 16-byte Type. + binary.LittleEndian.PutUint32(wrongSize[21:25], 8) + _, err = schema.decodeBatch(wrongSize, mp) + require.ErrorContains(t, err, "invalid vector type size") + require.Equal(t, int64(0), mp.CurrNB()) +} + +func TestNativeResultCodecRejectsNonFlatVectors(t *testing.T) { + schema := &nativeResultSchema{ + Version: nativeResultSchemaVersion, + Columns: []*nativeResultColumn{{Name: "v", Oid: uint32(types.T_int64)}}, + } + + for _, tc := range []struct { + name string + rows int + make func(*testing.T, *mpool.MPool) *vector.Vector + }{ + { + name: "constant billion-row broadcast", + rows: 1 << 30, + make: func(t *testing.T, mp *mpool.MPool) *vector.Vector { + vec, err := vector.NewConstFixed(types.T_int64.ToType(), int64(7), 1, mp) + require.NoError(t, err) + vec.SetLength(1 << 30) + return vec + }, + }, + { + name: "dictionary class", + rows: 1, + make: func(t *testing.T, mp *mpool.MPool) *vector.Vector { + vec := vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixed(vec, int64(7), false, mp)) + vec.SetClass(vector.DIST) + return vec + }, + }, + } { + t.Run(tc.name, func(t *testing.T) { + mp := mpool.MustNewZero() + payload := func() []byte { + bat := batch.NewWithSize(1) + bat.Vecs[0] = tc.make(t, mp) + bat.SetRowCount(tc.rows) + defer bat.Clean(mp) + payload, err := bat.MarshalBinary() + require.NoError(t, err) + if tc.rows > 1 { + require.Less(t, len(payload), 1024, "constant logical work must remain compressed on the wire") + } + return payload + }() + + _, err := schema.decodeBatch(payload, mp) + require.ErrorContains(t, err, "is not flat") + require.Equal(t, int64(0), mp.CurrNB()) + }) + } +} diff --git a/pkg/sql/compile/sidecarflight/stream.go b/pkg/sql/compile/sidecarflight/stream.go index 073f8e5658e56..1b3383cc05019 100644 --- a/pkg/sql/compile/sidecarflight/stream.go +++ b/pkg/sql/compile/sidecarflight/stream.go @@ -84,35 +84,50 @@ func (e *Execution) Run( if err = stream.CloseSend(); err != nil { return internalErrorf("sidecar flight: close ticket request: %w", err) } - seenSchema := false + var sequence uint64 + seenTransportSchema := false for { data := new(flightData) err = stream.RecvMsg(data) if err == io.EOF { - if !seenSchema { - return internalErrorf("sidecar flight: stream ended before its schema") + if !seenTransportSchema { + return internalErrorf("sidecar flight: stream ended before its Flight transport schema") } return e.finishSuccess() } if err != nil { return internalErrorf("sidecar flight: receive result: %w", err) } - if len(data.DataHeader) == 0 || uint64(len(data.DataBody)) > e.runtime.config.MaxBatchBytes { - return internalErrorf("sidecar flight: malformed or oversized FlightData") - } - if !seenSchema { - if len(data.DataBody) != 0 { - return internalErrorf("sidecar flight: schema message contains a body") - } - if err = e.schema.validateStreamSchema(data.DataHeader); err != nil { - return internalErrorf("sidecar flight: stream schema: %w", err) + if !seenTransportSchema { + if data.Descriptor != nil || len(data.DataHeader) == 0 || len(data.DataHeader) > maxNativeResultSchemaBytes || + len(data.DataBody) != 0 || len(data.AppMetadata) != 0 { + return internalErrorf("sidecar flight: malformed Flight transport schema") } - seenSchema = true + seenTransportSchema = true continue } - bat, decodeErr := e.schema.decodeRecordBatch(data.DataHeader, data.DataBody, e.runtime.config.MaxBatchBytes, mp) + if data.Descriptor != nil || len(data.DataHeader) == 0 || len(data.DataBody) != 0 || len(data.AppMetadata) != 0 || + uint64(len(data.DataHeader)) > e.runtime.config.MaxBatchBytes+nativeBatchFrameHeaderBytes { + return internalErrorf( + "sidecar flight: malformed or oversized MO native result frame: descriptor=%t header=%d body=%d metadata=%d", + data.Descriptor != nil, len(data.DataHeader), len(data.DataBody), len(data.AppMetadata), + ) + } + frameSequence, payload, frameErr := unmarshalNativeBatchFrame(data.DataHeader, e.runtime.config.MaxBatchBytes) + if frameErr != nil { + return frameErr + } + if frameSequence != sequence+1 { + return internalErrorf("sidecar flight: non-contiguous MO native result sequence") + } + sequence = frameSequence + bat, decodeErr := e.schema.decodeBatch(payload, mp) if decodeErr != nil { - return internalErrorf("sidecar flight: decode record batch: %w", decodeErr) + return decodeErr + } + bat.Attrs = make([]string, len(e.schema.Columns)) + for i := range e.schema.Columns { + bat.Attrs[i] = e.schema.Columns[i].Name } fillErr := func() error { defer bat.Clean(mp) @@ -146,8 +161,12 @@ func (e *Execution) finishSuccess() error { e.mu.Unlock() return internalErrorf("sidecar flight: execution was cancelled while finishing") } + inputs := slices.Clone(e.inputs) e.quiesced = true e.mu.Unlock() + for _, input := range inputs { + input.Retire() + } return nil } @@ -162,7 +181,10 @@ func (e *Execution) CancelAndJoin(ctx context.Context) error { ctx = context.Background() } e.mu.Lock() - if e.terminal || e.quiesced { + // Result EOF proves the execution worker is quiescent, but streamed reads + // may still own live DoPut handlers. Only direct executions can return from + // that state without explicitly aborting and joining their input streams. + if e.terminal || (e.quiesced && len(e.inputs) == 0) { e.mu.Unlock() return nil } @@ -183,12 +205,16 @@ func (e *Execution) CancelAndJoin(ctx context.Context) error { e.cleanupDone = make(chan struct{}) done := e.cleanupDone cancel := e.streamCancel + inputs := append([]*NativeInput(nil), e.inputs...) ticket := slices.Clone(e.ticket) var idempotencyKey []byte if len(ticket) == 0 { idempotencyKey = slices.Clone(e.idempotencyKey) } e.mu.Unlock() + for _, input := range inputs { + input.Abort(context.Canceled) + } if cancel != nil { cancel() } @@ -205,17 +231,16 @@ func (e *Execution) CancelAndJoin(ctx context.Context) error { } // CleanupAfterRun preserves the cancellation-before-lease-release ordering. -// Successful EOF is already quiescent; every other terminal path requires the -// explicit cancel-and-join acknowledgement. +// Direct successful EOF is already quiescent. Streamed successful EOF also +// retires local producers, but still requires CancelExecution to join the +// server-side DoPut handlers before leases are released. func (e *Execution) CleanupAfterRun(ctx context.Context, runErr error) error { if e == nil { return runErr } - if runErr != nil { - if err := e.CancelAndJoin(ctx); err != nil { - e.runtime.scheduleReconciliation(e) - return errors.Join(runErr, err) - } + if err := e.CancelAndJoin(ctx); err != nil { + e.runtime.scheduleReconciliation(e) + return errors.Join(runErr, err) } releaseErr := e.releaseLeases(ctx) if releaseErr != nil { diff --git a/pkg/sql/compile/sidecarflight/wire.go b/pkg/sql/compile/sidecarflight/wire.go index 127cb2c2203ee..3f2aaf9e41012 100644 --- a/pkg/sql/compile/sidecarflight/wire.go +++ b/pkg/sql/compile/sidecarflight/wire.go @@ -104,16 +104,28 @@ func (m *flightData) String() string { } func (*flightData) ProtoMessage() {} +type flightPutResult struct { + AppMetadata []byte `protobuf:"bytes,1,opt,name=app_metadata,json=appMetadata,proto3"` +} + +func (m *flightPutResult) Reset() { *m = flightPutResult{} } +func (m *flightPutResult) String() string { + return fmt.Sprintf("FlightPutResult{%d bytes}", len(m.AppMetadata)) +} +func (*flightPutResult) ProtoMessage() {} + type executeSubstraitRequest struct { - ProtocolVersion uint32 `protobuf:"varint,1,opt,name=protocol_version,json=protocolVersion,proto3"` - SubstraitVersion string `protobuf:"bytes,2,opt,name=substrait_version,json=substraitVersion,proto3"` - CapabilityHash []byte `protobuf:"bytes,3,opt,name=capability_hash,json=capabilityHash,proto3"` - MaxBatchBytes uint64 `protobuf:"varint,4,opt,name=max_batch_bytes,json=maxBatchBytes,proto3"` - DeadlineUnixMS uint64 `protobuf:"varint,5,opt,name=deadline_unix_ms,json=deadlineUnixMs,proto3"` - Plan []byte `protobuf:"bytes,6,opt,name=plan,proto3"` - QueryID []byte `protobuf:"bytes,7,opt,name=query_id,json=queryId,proto3"` - IdempotencyKey []byte `protobuf:"bytes,8,opt,name=idempotency_key,json=idempotencyKey,proto3"` - AccountID *uint64 `protobuf:"varint,9,opt,name=account_id,json=accountId,proto3,oneof"` + ProtocolVersion uint32 `protobuf:"varint,1,opt,name=protocol_version,json=protocolVersion,proto3"` + SubstraitVersion string `protobuf:"bytes,2,opt,name=substrait_version,json=substraitVersion,proto3"` + CapabilityHash []byte `protobuf:"bytes,3,opt,name=capability_hash,json=capabilityHash,proto3"` + MaxBatchBytes uint64 `protobuf:"varint,4,opt,name=max_batch_bytes,json=maxBatchBytes,proto3"` + DeadlineUnixMS uint64 `protobuf:"varint,5,opt,name=deadline_unix_ms,json=deadlineUnixMs,proto3"` + Plan []byte `protobuf:"bytes,6,opt,name=plan,proto3"` + QueryID []byte `protobuf:"bytes,7,opt,name=query_id,json=queryId,proto3"` + IdempotencyKey []byte `protobuf:"bytes,8,opt,name=idempotency_key,json=idempotencyKey,proto3"` + AccountID *uint64 `protobuf:"varint,9,opt,name=account_id,json=accountId,proto3,oneof"` + MaxInputBatchBytes uint64 `protobuf:"varint,10,opt,name=max_input_batch_bytes,json=maxInputBatchBytes,proto3"` + ResultSchema []byte `protobuf:"bytes,11,opt,name=result_schema,json=resultSchema,proto3"` } func (m *executeSubstraitRequest) Reset() { *m = executeSubstraitRequest{} } @@ -132,3 +144,29 @@ func (m *cancelExecutionRequest) String() string { return fmt.Sprintf("CancelExecutionRequest{%d,%d bytes}", len(m.Ticket), len(m.IdempotencyKey)) } func (*cancelExecutionRequest) ProtoMessage() {} + +type uploadInputRequest struct { + Ticket []byte `protobuf:"bytes,1,opt,name=ticket,proto3"` + StreamRef []byte `protobuf:"bytes,2,opt,name=stream_ref,json=streamRef,proto3"` +} + +func (m *uploadInputRequest) Reset() { *m = uploadInputRequest{} } +func (m *uploadInputRequest) String() string { + return fmt.Sprintf("UploadInputRequest{%d,%d bytes}", len(m.Ticket), len(m.StreamRef)) +} +func (*uploadInputRequest) ProtoMessage() {} + +type uploadInputAck struct { + AcknowledgedBatches uint64 `protobuf:"varint,1,opt,name=acknowledged_batches,json=acknowledgedBatches,proto3"` + Rows uint64 `protobuf:"varint,2,opt,name=rows,proto3"` + Bytes uint64 `protobuf:"varint,3,opt,name=bytes,proto3"` + Complete bool `protobuf:"varint,4,opt,name=complete,proto3"` + NotNeeded bool `protobuf:"varint,5,opt,name=not_needed,json=notNeeded,proto3"` + Ready bool `protobuf:"varint,6,opt,name=ready,proto3"` +} + +func (m *uploadInputAck) Reset() { *m = uploadInputAck{} } +func (m *uploadInputAck) String() string { + return fmt.Sprintf("UploadInputAck{%d,%t,%t,%t}", m.AcknowledgedBatches, m.Complete, m.NotNeeded, m.Ready) +} +func (*uploadInputAck) ProtoMessage() {} diff --git a/pkg/sql/compile/sirius_read.go b/pkg/sql/compile/sirius_read.go index 33ee02b839403..0172438aaa865 100644 --- a/pkg/sql/compile/sirius_read.go +++ b/pkg/sql/compile/sirius_read.go @@ -35,6 +35,12 @@ type SiriusReadPlan struct { OutputTypes []planpb.Type Headings []string LeaseExpiresAt time.Time + StreamInputs []SiriusStreamInput +} + +type SiriusStreamInput struct { + NodeID int32 + StreamRef []byte } func (p *SiriusReadPlan) Release(ctx context.Context, leases *substrait.LeaseManager) error { diff --git a/pkg/sql/compile/sirius_read_test.go b/pkg/sql/compile/sirius_read_test.go index 58658ee49aee9..49ae4af63849d 100644 --- a/pkg/sql/compile/sirius_read_test.go +++ b/pkg/sql/compile/sirius_read_test.go @@ -74,6 +74,8 @@ func TestCompileSiriusReadRejectsMissingPlan(t *testing.T) { func TestSiriusOffloadContextIsExplicit(t *testing.T) { require.False(t, siriusOffloadRequested(context.Background())) require.True(t, siriusOffloadRequested(WithSiriusOffload(context.Background()))) + require.Equal(t, siriusOffloadDirect, siriusOffloadModeFrom(WithSiriusOffload(context.Background()))) + require.Equal(t, siriusOffloadStream, siriusOffloadModeFrom(WithSiriusStreamOffload(context.Background()))) require.True(t, siriusStatementEligible(&tree.Select{})) require.False(t, siriusStatementEligible(&tree.Select{IsPerform: true})) require.False(t, siriusStatementEligible(&tree.Select{Ep: &tree.ExportParam{}})) diff --git a/pkg/sql/compile/sirius_runtime.go b/pkg/sql/compile/sirius_runtime.go index 31641fe0d9665..0ac64af596740 100644 --- a/pkg/sql/compile/sirius_runtime.go +++ b/pkg/sql/compile/sirius_runtime.go @@ -20,6 +20,7 @@ import ( "time" "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/common/mpool" moruntime "github.com/matrixorigin/matrixone/pkg/common/runtime" "github.com/matrixorigin/matrixone/pkg/defines" planpb "github.com/matrixorigin/matrixone/pkg/pb/plan" @@ -32,17 +33,34 @@ import ( // resolver ownership through the process default runtime. const SiriusRuntimeKey = "sql-compile-sirius-runtime" +var errSiriusResultComplete = moerr.NewInternalErrorNoCtx("substrait: Sirius result stream completed") + type siriusOffloadContextKey struct{} +type siriusOffloadMode uint8 + +const ( + siriusOffloadDirect siriusOffloadMode = iota + 1 + siriusOffloadStream +) + // WithSiriusOffload marks an explicitly hinted statement. Absence of this // marker leaves every native compile and execution path unchanged. func WithSiriusOffload(ctx context.Context) context.Context { - return context.WithValue(ctx, siriusOffloadContextKey{}, true) + return context.WithValue(ctx, siriusOffloadContextKey{}, siriusOffloadDirect) +} + +func WithSiriusStreamOffload(ctx context.Context) context.Context { + return context.WithValue(ctx, siriusOffloadContextKey{}, siriusOffloadStream) } func siriusOffloadRequested(ctx context.Context) bool { - requested, _ := ctx.Value(siriusOffloadContextKey{}).(bool) - return requested + return siriusOffloadModeFrom(ctx) != 0 +} + +func siriusOffloadModeFrom(ctx context.Context) siriusOffloadMode { + mode, _ := ctx.Value(siriusOffloadContextKey{}).(siriusOffloadMode) + return mode } func siriusStatementEligible(stmt tree.Statement) bool { @@ -135,12 +153,17 @@ func lookupSiriusRuntime(service string) (*SiriusRuntime, bool) { type siriusReadOwner struct { execution *sidecarflight.Execution runtime *SiriusRuntime + inputs []*sidecarflight.NativeInput } func newSiriusReadOwner(execution *sidecarflight.Execution, runtime *SiriusRuntime) *siriusReadOwner { return &siriusReadOwner{execution: execution, runtime: runtime} } +func newSiriusStreamOwner(execution *sidecarflight.Execution, runtime *SiriusRuntime, inputs []*sidecarflight.NativeInput) *siriusReadOwner { + return &siriusReadOwner{execution: execution, runtime: runtime, inputs: inputs} +} + func (o *siriusReadOwner) finish(ctx context.Context, succeeded bool) error { if o == nil { return nil @@ -154,13 +177,25 @@ func (o *siriusReadOwner) finish(ctx context.Context, succeeded bool) error { } func (c *Compile) tryCompileSiriusRead(ctx context.Context, queryPlan *planpb.Plan) (bool, error) { - if c == nil || !siriusOffloadRequested(ctx) || c.isPrepare || c.isInternal || !siriusStatementEligible(c.stmt) { + if c == nil || !siriusOffloadRequested(ctx) || c.isPrepare || c.isInternal { + return false, nil + } + if !siriusStatementEligible(c.stmt) { + if siriusOffloadModeFrom(ctx) == siriusOffloadStream { + return false, moerr.NewInternalError(ctx, "substrait: streamed Sirius offload requires a SELECT statement") + } return false, nil } runtime, ok := lookupSiriusRuntime(c.proc.GetService()) if !ok { + if siriusOffloadModeFrom(ctx) == siriusOffloadStream { + return false, moerr.NewInternalError(ctx, "substrait: streamed Sirius runtime is not configured") + } return false, nil } + if siriusOffloadModeFrom(ctx) == siriusOffloadStream { + return c.tryCompileSiriusStreamRead(ctx, queryPlan, runtime) + } accountID, err := defines.GetAccountId(ctx) if err != nil { return false, nil @@ -237,11 +272,17 @@ func releaseReadRefs(ctx context.Context, leases *substrait.LeaseManager, readRe return result } -func (c *Compile) runSiriusRead(ctx context.Context) (err error) { +func (c *Compile) runSiriusRead( + ctx context.Context, + allocationExporter func(mpool.AllocationAccountTerminalSnapshot), +) (err error) { owner := c.siriusRead if owner == nil { return moerr.NewInternalError(ctx, "substrait: missing Sirius execution owner") } + if len(owner.inputs) != 0 { + return c.runSiriusStreamRead(ctx, owner, allocationExporter) + } defer func() { if recovered := recover(); recovered != nil { _ = owner.finish(ctx, false) @@ -251,3 +292,131 @@ func (c *Compile) runSiriusRead(ctx context.Context) (err error) { runErr := owner.execution.Run(ctx, c.proc.Mp(), c.counterSet, c.fill) return errors.Join(runErr, owner.finish(ctx, runErr == nil)) } + +func (c *Compile) runSiriusStreamRead( + ctx context.Context, + owner *siriusReadOwner, + allocationExporter func(mpool.AllocationAccountTerminalSnapshot), +) (err error) { + runCtx, cancel := context.WithCancelCause(ctx) + defer cancel(context.Canceled) + if c.MessageBoard == nil { + setupErr := moerr.NewInternalError(ctx, "substrait: streamed Sirius execution has no message board") + return errors.Join(setupErr, owner.finish(ctx, false)) + } + c.remoteFragmentCounts = collectRemoteFragmentCounts(c.scopes, c.addr) + if len(c.remoteFragmentCounts) != 0 { + setupErr := moerr.NewInternalError(ctx, "substrait: streamed Sirius execution is not local-CN") + return errors.Join(setupErr, owner.finish(ctx, false)) + } + if setupErr := c.ensureAllocationAccountLifecycle(allocationExporter); setupErr != nil { + return errors.Join(setupErr, owner.finish(ctx, false)) + } + allocationAttempt, setupErr := c.beginAllocationAccountAttempt() + if setupErr != nil { + return errors.Join(setupErr, owner.finish(ctx, false)) + } + defer func() { + if allocationAttempt == nil { + return + } + _, finishErr := allocationAttempt.finish() + if c.allocationAttempt == allocationAttempt { + c.allocationAttempt = nil + } + err = errors.Join(err, finishErr) + }() + for _, input := range owner.inputs { + if startErr := input.Start(runCtx); startErr != nil { + cancel(startErr) + for _, pending := range owner.inputs { + pending.Abort(startErr) + } + return errors.Join(startErr, owner.finish(ctx, false)) + } + } + producerDone := make(chan error, 1) + producerJoined := false + defer func() { + if recovered := recover(); recovered != nil { + panicErr := moerr.ConvertPanicError(c.proc.Ctx, recovered) + cancel(panicErr) + if c.proc != nil && c.proc.Cancel != nil { + c.proc.Cancel(panicErr) + } + _ = owner.finish(ctx, false) + if !producerJoined { + <-producerDone + } + panic(recovered) + } + }() + go func() { + var producerErr error + defer func() { + if recovered := recover(); recovered != nil { + producerErr = moerr.ConvertPanicError(c.proc.Ctx, recovered) + } + if producerErr != nil { + cancel(producerErr) + if c.proc != nil && c.proc.Cancel != nil { + c.proc.Cancel(producerErr) + } + } + producerDone <- producerErr + }() + producerErr = c.prePipelineInitializer() + if producerErr == nil { + c.MessageBoard.BeforeRunonce() + producerErr = c.runOnce() + } + }() + + resultErr := owner.execution.Run(runCtx, c.proc.Mp(), c.counterSet, c.fill) + if resultErr == nil { + // Result EOF is authoritative success. Retire the local producer even if + // the sidecar pruned an input before its first batch; NativeInput.Retire + // has already interrupted any blocked DoPut acknowledgement. + cancel(errSiriusResultComplete) + if c.proc != nil && c.proc.Cancel != nil { + c.proc.Cancel(errSiriusResultComplete) + } + } else { + cancel(resultErr) + if c.proc != nil && c.proc.Cancel != nil { + c.proc.Cancel(resultErr) + } + } + producerErr := <-producerDone + producerJoined = true + for _, input := range owner.inputs { + producerErr = errors.Join(producerErr, input.Err()) + } + if resultErr == nil && context.Cause(ctx) == nil && isOnlySiriusResultCompletion(producerErr) { + producerErr = nil + } + runErr := errors.Join(resultErr, producerErr) + return errors.Join(runErr, owner.finish(ctx, runErr == nil)) +} + +func isOnlySiriusResultCompletion(err error) bool { + if err == nil { + return true + } + if joined, ok := err.(interface{ Unwrap() []error }); ok { + children := joined.Unwrap() + if len(children) == 0 { + return false + } + for _, child := range children { + if !isOnlySiriusResultCompletion(child) { + return false + } + } + return true + } + if wrapped := errors.Unwrap(err); wrapped != nil { + return isOnlySiriusResultCompletion(wrapped) + } + return errors.Is(err, errSiriusResultComplete) || errors.Is(err, context.Canceled) +} diff --git a/pkg/sql/compile/sirius_runtime_test.go b/pkg/sql/compile/sirius_runtime_test.go index a5dd55a893a1c..6e0ebac463f21 100644 --- a/pkg/sql/compile/sirius_runtime_test.go +++ b/pkg/sql/compile/sirius_runtime_test.go @@ -174,12 +174,39 @@ func TestSiriusCompileFastRejections(t *testing.T) { offloaded, err := (&Compile{}).tryCompileSiriusRead(context.Background(), nil) require.NoError(t, err) require.False(t, offloaded) + offloaded, err = (&Compile{stmt: &tree.ExplainAnalyze{}}).tryCompileSiriusRead(WithSiriusStreamOffload(context.Background()), nil) + require.False(t, offloaded) + require.ErrorContains(t, err, "requires a SELECT statement") require.NoError(t, (*siriusReadOwner)(nil).finish(context.Background(), false)) - err = (&Compile{}).runSiriusRead(context.Background()) + err = (&Compile{}).runSiriusRead(context.Background(), nil) require.ErrorContains(t, err, "missing Sirius execution owner") } +func TestSiriusStreamCompileInitializesOneCNPhysicalState(t *testing.T) { + query := &planpb.Query{ + StmtType: planpb.Query_SELECT, Steps: []int32{0}, Headings: []string{"a"}, + Nodes: []*planpb.Node{{ + NodeId: 0, NodeType: planpb.Node_TABLE_SCAN, Stats: &planpb.Stats{}, + ObjRef: &planpb.ObjectRef{Obj: 42, ObjName: "t"}, + TableDef: &planpb.TableDef{DbId: 7, TblId: 42, Name: "t", TableType: "r", Cols: []*planpb.ColDef{{ + Name: "a", Typ: planpb.Type{Id: int32(types.T_int64)}, + }}}, + }}, + } + queryPlan := &planpb.Plan{Plan: &planpb.Plan_Query{Query: query}} + proc := testutil.NewProcess(t) + defer proc.Free() + c := &Compile{proc: proc, ncpu: 8} + c.initSiriusStreamCompile(queryPlan) + require.NotNil(t, c.anal) + require.Same(t, query, c.anal.qry) + require.NotNil(t, query.Nodes[0].AnalyzeInfo) + require.GreaterOrEqual(t, query.Nodes[0].Stats.Dop, int32(1)) + require.NotEqual(t, plan2.ExecTypeAP_MULTICN, c.execType) + c.anal.release() +} + func TestSQLSelectLimitIsMaterializedBeforeSiriusExport(t *testing.T) { proc := testutil.NewProcess(t) proc.Base.SessionInfo.ApplySQLSelectLimit = true diff --git a/pkg/sql/compile/sirius_stream.go b/pkg/sql/compile/sirius_stream.go new file mode 100644 index 0000000000000..3e389617f4f5c --- /dev/null +++ b/pkg/sql/compile/sirius_stream.go @@ -0,0 +1,243 @@ +// Copyright 2026 Matrix Origin +// +// 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 compile + +import ( + "context" + "crypto/rand" + "crypto/sha256" + "errors" + "time" + + "github.com/matrixorigin/matrixone/pkg/common/moerr" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/defines" + planpb "github.com/matrixorigin/matrixone/pkg/pb/plan" + "github.com/matrixorigin/matrixone/pkg/perfcounter" + "github.com/matrixorigin/matrixone/pkg/sql/colexec/output" + "github.com/matrixorigin/matrixone/pkg/sql/compile/sidecarflight" + plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan" + "github.com/matrixorigin/matrixone/pkg/sql/plan/substrait" +) + +func (c *Compile) CompileSiriusStreamRead( + ctx context.Context, + queryPlan *planpb.Plan, + accountID uint64, + queryID []byte, + ttl time.Duration, +) (*SiriusReadPlan, error) { + if c == nil || c.proc == nil || queryPlan == nil || queryPlan.GetQuery() == nil { + return nil, moerr.NewInternalError(ctx, "substrait: stream compile has no SELECT plan") + } + candidate, err := substrait.Export(queryPlan.GetQuery()) + if err != nil { + return nil, err + } + reads, err := candidate.StreamReads() + if err != nil { + return nil, err + } + if len(reads) == 0 || len(reads) > 16 || len(queryID) != 16 || ttl <= 0 { + return nil, substrait.NotEligible(substrait.EligibilityPlanShape, "streamed read count or identity is unsupported") + } + txnOp := c.proc.GetTxnOperator() + if txnOp == nil || txnOp.GetWorkspace() == nil || !txnOp.GetWorkspace().Readonly() || + txnOp.GetWorkspace().WriteOffset() != 0 || txnOp.GetWorkspace().GetSnapshotWriteOffset() != 0 { + return nil, substrait.NotEligible(substrait.EligibilityTransaction, "stream mode requires a read-only snapshot without prior writes") + } + snapshot := types.TimestampToTS(txnOp.SnapshotTS()) + snapshotBytes, err := snapshot.Marshal() + if err != nil { + return nil, err + } + expires := time.Now().Add(ttl) + bindings := make(map[int32]substrait.ReadBinding, len(reads)) + inputs := make([]SiriusStreamInput, 0, len(reads)) + for _, read := range reads { + if read.Occurrences != 1 { + return nil, substrait.NotEligible(substrait.EligibilityPlanShape, "stream mode does not replay shared scan nodes") + } + ref := make([]byte, 32) + if _, err = rand.Read(ref); err != nil { + return nil, err + } + digest := sha256.Sum256(read.StreamSchema) + wire, marshalErr := substrait.MarshalStreamRead(&substrait.StreamRead{ + ProtocolVersion: substrait.StreamReadProtocolVersion, + StreamRef: ref, QueryID: append([]byte(nil), queryID...), AccountID: accountID, + SnapshotTS: snapshotBytes, SchemaDigest: digest[:], CapabilityHash: substrait.CapabilityHash[:], + ExpiresAtUnixMS: uint64(expires.UnixMilli()), + }) + if marshalErr != nil { + return nil, marshalErr + } + bindings[read.NodeID] = substrait.ReadBinding{ + TypeURL: substrait.StreamReadTypeURL, Value: wire, Schema: read.StreamSchema, + } + inputs = append(inputs, SiriusStreamInput{ + NodeID: read.NodeID, StreamRef: append([]byte(nil), ref...), + }) + } + planBytes, err := candidate.BuildWithBindings(bindings) + if err != nil { + return nil, err + } + return &SiriusReadPlan{ + Plan: planBytes, OutputTypes: candidate.OutputTypes(), + Headings: append([]string(nil), queryPlan.GetQuery().Headings...), + LeaseExpiresAt: expires, StreamInputs: inputs, + }, nil +} + +func (c *Compile) tryCompileSiriusStreamRead( + ctx context.Context, + queryPlan *planpb.Plan, + runtime *SiriusRuntime, +) (bool, error) { + account, err := defines.GetAccountId(ctx) + if err != nil { + return false, err + } + accountID := uint64(account) + statementID := c.proc.GetStmtProfile().GetStmtId() + queryID := append([]byte(nil), statementID[:]...) + readPlan, err := c.CompileSiriusStreamRead(ctx, queryPlan, accountID, queryID, runtime.LeaseTTL) + if err != nil { + return false, err + } + c.initSiriusStreamCompile(queryPlan) + execution, err := runtime.Flight.Prepare( + ctx, accountID, queryID, readPlan.Plan, readPlan.OutputTypes, readPlan.Headings, + readPlan.LeaseExpiresAt.Add(-runtime.CleanupTimeout), func(context.Context) error { return nil }, + ) + if err != nil { + return false, err + } + inputs, scopes, err := c.compileSiriusStreamScopes(queryPlan.GetQuery(), readPlan.StreamInputs, execution) + if err != nil { + cleanupCtx, cancel := context.WithTimeoutCause( + context.WithoutCancel(ctx), runtime.CleanupTimeout, + moerr.NewInternalErrorNoCtx("substrait: timed out cleaning up failed streamed compile"), + ) + defer cancel() + return false, errors.Join(err, execution.Cleanup(cleanupCtx)) + } + c.scopes = scopes + c.siriusRead = newSiriusStreamOwner(execution, runtime, inputs) + return true, nil +} + +func (c *Compile) initSiriusStreamCompile(queryPlan *planpb.Plan) { + execType := plan2.GetExecType(queryPlan.GetQuery(), c.getHaveDDL(), c.isPrepare) + if execType == plan2.ExecTypeAP_MULTICN { + execType = plan2.ExecTypeAP_ONECN + } + c.execType = execType + ncpu := int32(c.ncpu) + if ncpu < 1 { + ncpu = 1 + } + plan2.CalcQueryDOP(queryPlan, ncpu, 1, execType) + c.initAnalyzeModule(queryPlan.GetQuery()) +} + +func (c *Compile) compileSiriusLocalTableScan(node *planpb.Node) ([]*Scope, error) { + if _, _, _, err := c.handleDbRelContext(node, false); err != nil { + return nil, err + } + local := getEngineNode(c) + local.Addr = c.addr + if node.Stats != nil && node.Stats.Dop > 0 { + local.Mcpu = min(local.Mcpu, int(node.Stats.Dop)) + } + local.Mcpu = normalizeMcpu(local.Mcpu) + local.CNCNT = 1 + local.CNIDX = 0 + scope, err := c.compileTableScanWithNode(node, local, c.anal.isFirst) + if err != nil { + return nil, err + } + c.anal.isFirst = false + return []*Scope{scope}, nil +} + +func (c *Compile) compileSiriusStreamScopes( + query *planpb.Query, + streamInputs []SiriusStreamInput, + execution *sidecarflight.Execution, +) ([]*sidecarflight.NativeInput, []*Scope, error) { + inputs := make([]*sidecarflight.NativeInput, 0, len(streamInputs)) + roots := make([]*Scope, 0, len(streamInputs)) + succeeded := false + defer func() { + if !succeeded { + ReleaseScopes(roots) + } + }() + for _, spec := range streamInputs { + if spec.NodeID < 0 || int(spec.NodeID) >= len(query.Nodes) || query.Nodes[spec.NodeID] == nil { + return nil, nil, moerr.NewInternalErrorNoCtx("substrait: streamed scan node is missing") + } + node := plan2.DeepCopyNode(query.Nodes[spec.NodeID]) + c.appendMetaTables(node.ObjRef) + node.RuntimeFilterProbeList = nil + node.RuntimeFilterBuildList = nil + node.RecvMsgList = nil + // Scan-level aggregation is a native storage optimization. The streamed + // relation must contain ordinary post-filter/project rows because the + // semantic aggregate remains in the exported Substrait plan. + node.AggList = nil + if len(node.ProjectList) == 0 { + for position, column := range node.TableDef.Cols { + if column == nil || column.Hidden { + continue + } + node.ProjectList = append(node.ProjectList, &planpb.Expr{ + Typ: column.Typ, + Expr: &planpb.Expr_Col{Col: &planpb.ColRef{RelPos: 0, ColPos: int32(position)}}, + }) + } + } + c.setAnalyzeCurrent(nil, int(spec.NodeID)) + scans, err := c.compileSiriusLocalTableScan(node) + if err != nil { + return nil, nil, err + } + scans = c.compileTableScanFiltersAndProjection(node, scans) + if node.Offset != nil { + scans = c.compileOffset(node, scans) + } + if node.Limit != nil { + scans = c.compileLimit(node, scans) + } + root := c.newMergeScope(scans) + roots = append(roots, root) + nativeInput, err := execution.NewNativeInput(spec.StreamRef) + if err != nil { + return nil, nil, err + } + root.setRootOperator(output.NewArgument().WithFunc(func(bat *batch.Batch, _ *perfcounter.CounterSet) error { + if bat == nil { + return nativeInput.Finish(c.proc.Ctx) + } + return nativeInput.Send(c.proc.Ctx, bat, c.proc.Mp()) + }).WithShouldStop(nativeInput.NotNeeded)) + inputs = append(inputs, nativeInput) + } + succeeded = true + return inputs, roots, nil +} diff --git a/pkg/sql/compile/sirius_stream_test.go b/pkg/sql/compile/sirius_stream_test.go new file mode 100644 index 0000000000000..2cbcef8577df0 --- /dev/null +++ b/pkg/sql/compile/sirius_stream_test.go @@ -0,0 +1,299 @@ +// Copyright 2026 Matrix Origin +// +// 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 compile + +import ( + "context" + "errors" + "fmt" + "testing" + "time" + + "github.com/golang/mock/gomock" + "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/types" + "github.com/matrixorigin/matrixone/pkg/container/vector" + "github.com/matrixorigin/matrixone/pkg/defines" + mock_frontend "github.com/matrixorigin/matrixone/pkg/frontend/test" + planpb "github.com/matrixorigin/matrixone/pkg/pb/plan" + "github.com/matrixorigin/matrixone/pkg/pb/timestamp" + "github.com/matrixorigin/matrixone/pkg/perfcounter" + "github.com/matrixorigin/matrixone/pkg/sql/compile/sidecarflight" + "github.com/matrixorigin/matrixone/pkg/sql/parsers/tree" + plan2 "github.com/matrixorigin/matrixone/pkg/sql/plan" + "github.com/matrixorigin/matrixone/pkg/sql/plan/substrait" + "github.com/matrixorigin/matrixone/pkg/testutil" + "github.com/matrixorigin/matrixone/pkg/vm/engine" + "github.com/matrixorigin/matrixone/pkg/vm/message" + "github.com/stretchr/testify/require" + spb "github.com/substrait-io/substrait-protobuf/go/substraitpb" + "google.golang.org/protobuf/proto" +) + +func siriusStreamTestPlan() *planpb.Plan { + return &planpb.Plan{Plan: &planpb.Plan_Query{Query: &planpb.Query{ + StmtType: planpb.Query_SELECT, + Steps: []int32{0}, + Headings: []string{"a"}, + Nodes: []*planpb.Node{{ + NodeId: 0, + NodeType: planpb.Node_TABLE_SCAN, + Stats: &planpb.Stats{}, + ObjRef: &planpb.ObjectRef{Obj: 42, ObjName: "t"}, + TableDef: &planpb.TableDef{ + DbId: 7, TblId: 42, Version: 3, Name: "t", TableType: "r", + Cols: []*planpb.ColDef{{ + Name: "a", ColId: 11, Seqnum: 5, + Typ: planpb.Type{Id: int32(types.T_int64)}, + }}, + }, + }}, + }}} +} + +func TestCompileSiriusStreamReadBindsSnapshotAndNativeInput(t *testing.T) { + ctrl := gomock.NewController(t) + workspace := mock_frontend.NewMockWorkspace(ctrl) + workspace.EXPECT().Readonly().Return(true) + workspace.EXPECT().WriteOffset().Return(uint64(0)) + workspace.EXPECT().GetSnapshotWriteOffset().Return(0) + txnOp := mock_frontend.NewMockTxnOperator(ctrl) + txnOp.EXPECT().GetWorkspace().Return(workspace).AnyTimes() + txnOp.EXPECT().SnapshotTS().Return(timestamp.Timestamp{PhysicalTime: 42, LogicalTime: 3}) + + proc := testutil.NewProcess(t) + t.Cleanup(proc.Free) + proc.Base.TxnOperator = txnOp + queryPlan := siriusStreamTestPlan() + queryID := []byte("0123456789abcdef") + started := time.Now() + + readPlan, err := (&Compile{proc: proc}).CompileSiriusStreamRead( + context.Background(), queryPlan, 7, queryID, time.Minute, + ) + require.NoError(t, err) + require.Len(t, readPlan.StreamInputs, 1) + require.Len(t, readPlan.StreamInputs[0].StreamRef, 32) + require.Equal(t, int32(0), readPlan.StreamInputs[0].NodeID) + require.Equal(t, queryPlan.GetQuery().Headings, readPlan.Headings) + require.Len(t, readPlan.OutputTypes, 1) + require.WithinDuration(t, started.Add(time.Minute), readPlan.LeaseExpiresAt, time.Second) + + var exported spb.Plan + require.NoError(t, proto.Unmarshal(readPlan.Plan, &exported)) + require.Equal(t, []string{substrait.StreamReadTypeURL}, exported.ExpectedTypeUrls) + read := exported.Relations[0].GetRoot().Input.GetRead() + require.NotNil(t, read) + require.Equal(t, substrait.StreamReadTypeURL, read.GetExtensionTable().Detail.TypeUrl) + require.NotEmpty(t, read.GetExtensionTable().Detail.Value) +} + +func TestCompileSiriusStreamReadRejectsActualUint32NativeBatchBeforePrepare(t *testing.T) { + proc := testutil.NewProcess(t) + t.Cleanup(proc.Free) + + native := batch.NewWithSize(1) + native.Vecs[0] = vector.NewVec(types.T_uint32.ToType()) + require.NoError(t, vector.AppendFixed(native.Vecs[0], uint32(42), false, proc.Mp())) + native.SetRowCount(1) + defer native.Clean(proc.Mp()) + payload, err := native.MarshalBinary() + require.NoError(t, err) + + decoded := batch.NewOffHeapEmpty() + require.NoError(t, decoded.UnmarshalBinaryWithAnyMp(payload, proc.Mp())) + defer decoded.Clean(proc.Mp()) + require.Equal(t, types.T_uint32, decoded.Vecs[0].GetType().Oid, + "the MO-native codec must not silently widen uint32 to int64") + + queryPlan := siriusStreamTestPlan() + queryPlan.GetQuery().Nodes[0].TableDef.Cols[0].Typ = planpb.Type{ + Id: int32(decoded.Vecs[0].GetType().Oid), + } + _, err = (&Compile{proc: proc}).CompileSiriusStreamRead( + context.Background(), queryPlan, 7, []byte("0123456789abcdef"), time.Minute, + ) + require.True(t, substrait.IsNotEligible(err)) + require.ErrorContains(t, err, "unsupported native input type INT UNSIGNED") +} + +func TestCompileSiriusStreamReadRejectsBeforeOpeningTransport(t *testing.T) { + valid := siriusStreamTestPlan() + proc := testutil.NewProcess(t) + t.Cleanup(proc.Free) + + for _, tc := range []struct { + name string + compile *Compile + plan *planpb.Plan + queryID []byte + ttl time.Duration + want string + }{ + {name: "nil compile", plan: valid, queryID: make([]byte, 16), ttl: time.Minute, want: "no SELECT plan"}, + {name: "missing process", compile: &Compile{}, plan: valid, queryID: make([]byte, 16), ttl: time.Minute, want: "no SELECT plan"}, + {name: "unsupported plan", compile: &Compile{proc: proc}, plan: &planpb.Plan{Plan: &planpb.Plan_Query{Query: &planpb.Query{StmtType: planpb.Query_SELECT}}}, queryID: make([]byte, 16), ttl: time.Minute, want: "SELECT query root"}, + {name: "bad query identity", compile: &Compile{proc: proc}, plan: valid, queryID: make([]byte, 15), ttl: time.Minute, want: "identity is unsupported"}, + {name: "bad lease ttl", compile: &Compile{proc: proc}, plan: valid, queryID: make([]byte, 16), want: "identity is unsupported"}, + {name: "missing transaction", compile: &Compile{proc: proc}, plan: valid, queryID: make([]byte, 16), ttl: time.Minute, want: "read-only snapshot"}, + } { + t.Run(tc.name, func(t *testing.T) { + _, err := tc.compile.CompileSiriusStreamRead(context.Background(), tc.plan, 0, tc.queryID, tc.ttl) + require.ErrorContains(t, err, tc.want) + }) + } +} + +func TestTryCompileSiriusStreamReadCleansUpPreVisibilityPrepareFailure(t *testing.T) { + ctrl := gomock.NewController(t) + workspace := mock_frontend.NewMockWorkspace(ctrl) + workspace.EXPECT().Readonly().Return(true) + workspace.EXPECT().WriteOffset().Return(uint64(0)) + workspace.EXPECT().GetSnapshotWriteOffset().Return(0) + workspace.EXPECT().GetHaveDDL().Return(false) + txnOp := mock_frontend.NewMockTxnOperator(ctrl) + txnOp.EXPECT().GetWorkspace().Return(workspace).AnyTimes() + txnOp.EXPECT().SnapshotTS().Return(timestamp.Timestamp{PhysicalTime: 42}) + + proc := testutil.NewProcess(t) + t.Cleanup(proc.Free) + proc.Base.TxnOperator = txnOp + c := &Compile{proc: proc} + runtime := &SiriusRuntime{ + Flight: new(sidecarflight.Runtime), LeaseTTL: time.Minute, CleanupTimeout: time.Second, + } + offloaded, err := c.tryCompileSiriusStreamRead( + defines.AttachAccountId(context.Background(), 7), siriusStreamTestPlan(), runtime, + ) + require.False(t, offloaded) + require.ErrorContains(t, err, "lease-safe execution deadline has expired") + require.NotNil(t, c.anal) + c.anal.release() +} + +func TestSiriusStreamRuntimeRejectsIncompleteLocalExecution(t *testing.T) { + execution := new(sidecarflight.Execution) + runtime := &SiriusRuntime{CleanupTimeout: time.Second} + directOwner := newSiriusReadOwner(execution, runtime) + require.NoError(t, directOwner.finish(context.Background(), true)) + require.NoError(t, directOwner.finish(context.Background(), false)) + + proc := testutil.NewProcess(t) + t.Cleanup(proc.Free) + c := &Compile{ + proc: proc, + fill: func(*batch.Batch, *perfcounter.CounterSet) error { + return nil + }, + siriusRead: directOwner, + } + require.ErrorContains(t, c.runSiriusRead(context.Background(), nil), "invalid execution") + + c.siriusRead = newSiriusStreamOwner(execution, runtime, []*sidecarflight.NativeInput{nil}) + require.ErrorContains(t, c.runSiriusRead(context.Background(), nil), "has no message board") + + board := message.NewMessageBoard() + t.Cleanup(func() { board.CloseAndDrain() }) + c.MessageBoard = board + c.addr = "local" + c.scopes = []*Scope{{Magic: Remote, NodeInfo: engine.Node{Addr: "remote"}}} + require.ErrorContains(t, c.runSiriusRead(context.Background(), nil), "is not local-CN") + + startErr := context.DeadlineExceeded + input := new(sidecarflight.NativeInput) + input.Abort(startErr) + c.scopes = nil + streamOwner := newSiriusStreamOwner(execution, runtime, []*sidecarflight.NativeInput{input}) + require.ErrorIs(t, c.runSiriusStreamRead(context.Background(), streamOwner, nil), startErr) + + emptyOwner := newSiriusStreamOwner(execution, runtime, nil) + require.ErrorContains(t, c.runSiriusStreamRead(context.Background(), emptyOwner, nil), "invalid execution") +} + +func TestSiriusStreamCompileRejectsMissingScanAndRuntime(t *testing.T) { + c := new(Compile) + _, _, err := c.compileSiriusStreamScopes( + &planpb.Query{}, []SiriusStreamInput{{NodeID: 0, StreamRef: make([]byte, 32)}}, nil, + ) + require.ErrorContains(t, err, "streamed scan node is missing") + + proc := testutil.NewProcess(t) + t.Cleanup(proc.Free) + offloaded, err := (&Compile{proc: proc, stmt: &tree.Select{}}).tryCompileSiriusRead( + WithSiriusStreamOffload(context.Background()), siriusStreamTestPlan(), + ) + require.False(t, offloaded) + require.ErrorContains(t, err, "runtime is not configured") + + streamRuntime := &SiriusRuntime{LeaseTTL: time.Minute} + offloaded, err = (&Compile{proc: proc}).tryCompileSiriusStreamRead( + context.Background(), siriusStreamTestPlan(), streamRuntime, + ) + require.False(t, offloaded) + require.Error(t, err) + offloaded, err = (&Compile{proc: proc}).tryCompileSiriusStreamRead( + defines.AttachAccountId(context.Background(), 7), siriusStreamTestPlan(), streamRuntime, + ) + require.False(t, offloaded) + require.True(t, substrait.IsNotEligible(err)) +} + +func TestSiriusResultCompletionErrorClassification(t *testing.T) { + require.True(t, isOnlySiriusResultCompletion(nil)) + require.True(t, isOnlySiriusResultCompletion(errSiriusResultComplete)) + require.True(t, isOnlySiriusResultCompletion(errors.Join( + fmt.Errorf("wrapped: %w", errSiriusResultComplete), context.Canceled, + ))) + require.False(t, isOnlySiriusResultCompletion(errors.New("producer failed"))) + require.False(t, isOnlySiriusResultCompletion(errors.Join( + errSiriusResultComplete, errors.New("producer failed"), + ))) +} + +func TestCompileSiriusStreamScopesReleasesTreeWhenInputRegistrationFails(t *testing.T) { + ctrl := gomock.NewController(t) + storage := mock_frontend.NewMockEngine(ctrl) + database := mock_frontend.NewMockDatabase(ctrl) + relation := mock_frontend.NewMockRelation(ctrl) + storage.EXPECT().Database(gomock.Any(), "db", gomock.Any()).Return(database, nil) + database.EXPECT().Relation(gomock.Any(), "t", gomock.Any()).Return(relation, nil) + + proc := testutil.NewProcess(t) + t.Cleanup(proc.Free) + queryPlan := siriusStreamTestPlan() + node := queryPlan.GetQuery().Nodes[0] + node.ObjRef.SchemaName = "db" + node.Stats = &planpb.Stats{Dop: 2} + node.TableDef.Cols = append([]*planpb.ColDef{ + {Name: "hidden", Hidden: true, Typ: planpb.Type{Id: int32(types.T_int64)}}, + }, node.TableDef.Cols...) + node.Offset = plan2.MakePlan2Uint64ConstExprWithType(1) + node.Limit = plan2.MakePlan2Uint64ConstExprWithType(2) + node.AggList = []*planpb.Expr{{}} + + c := &Compile{proc: proc, e: storage, addr: "local", ncpu: 4} + c.initSiriusStreamCompile(queryPlan) + t.Cleanup(c.anal.release) + before := proc.Mp().CurrNB() + inputs, scopes, err := c.compileSiriusStreamScopes( + queryPlan.GetQuery(), + []SiriusStreamInput{{NodeID: 0, StreamRef: make([]byte, 32)}}, + new(sidecarflight.Execution), + ) + require.ErrorContains(t, err, "invalid native input identity") + require.Nil(t, inputs) + require.Nil(t, scopes) + require.Equal(t, before, proc.Mp().CurrNB()) +} diff --git a/pkg/sql/plan/substrait/export.go b/pkg/sql/plan/substrait/export.go index fa70fc6791375..8ffaafe4e671c 100644 --- a/pkg/sql/plan/substrait/export.go +++ b/pkg/sql/plan/substrait/export.go @@ -18,9 +18,11 @@ package substrait import ( + "bytes" "context" "encoding/binary" "math" + "strconv" "strings" "sync" @@ -37,9 +39,10 @@ import ( ) const ( - Version = "0.78.0" - TaeReadTypeURL = "type.googleapis.com/matrixone.sirius.v1.TaeRead" - MaxPlanBytes = 16 << 20 + Version = "0.78.0" + TaeReadTypeURL = "type.googleapis.com/matrixone.sirius.v1.TaeRead" + StreamReadTypeURL = "type.googleapis.com/matrixone.sirius.v1.StreamRead" + MaxPlanBytes = 16 << 20 // Bound the materialization of optimizer-folded IN lists before allocating // one Substrait expression per member. maxLiteralVectorValues = 1 << 16 @@ -57,6 +60,9 @@ type Candidate struct { // Read identifies one physical table scan which needs a TaeRead lease. type Read struct { NodeID int32 + // Occurrences counts references to this scan node in the exported relation + // graph. One-pass StreamRead inputs reject values other than one. + Occurrences uint32 // AccountID is bound by Admit immediately before snapshot preparation; // logical export deliberately leaves it unset. AccountID uint64 @@ -65,6 +71,12 @@ type Read struct { SchemaVersion uint32 Columns []ColumnMapping Schema []byte // deterministic Substrait NamedStruct bytes + StreamSchema []byte + StreamNames []string + // streamErr records a physical MO input type which the native StreamRead + // wire cannot represent. It must not make the semantic TaeRead candidate + // ineligible: direct TAE decoding has its own, wider type contract. + streamErr error } // ColumnMapping binds one exported ordinal to the physical TAE column used at @@ -79,11 +91,28 @@ func (c *Candidate) Reads() []Read { result := append([]Read(nil), c.reads...) for i := range result { result[i].Schema = append([]byte(nil), result[i].Schema...) + result[i].StreamSchema = append([]byte(nil), result[i].StreamSchema...) + result[i].StreamNames = append([]string(nil), result[i].StreamNames...) result[i].Columns = append([]ColumnMapping(nil), result[i].Columns...) } return result } +// StreamReads returns the validated one-pass native input contract. A plan can +// remain eligible for direct TaeRead while being ineligible for StreamRead +// when its semantic Substrait type is wider than its unchanged MO vector type. +func (c *Candidate) StreamReads() ([]Read, error) { + if c == nil { + return nil, moerr.NewInternalErrorNoCtx("substrait: nil candidate") + } + for _, read := range c.reads { + if read.streamErr != nil { + return nil, read.streamErr + } + } + return c.Reads(), nil +} + // OutputTypes returns the MatrixOne result contract corresponding to the // Substrait root names. Transport decoding uses it to restore MO physical // representations (notably DATE and unsigned EXTRACT results). @@ -150,10 +179,41 @@ func Export(q *planpb.Query) (*Candidate, error) { // Build binds admitted TaeRead messages to every scan and serializes the plan. func (c *Candidate) Build(readValues map[int32][]byte) ([]byte, error) { + bindings := make(map[int32]ReadBinding, len(readValues)) + for nodeID, value := range readValues { + bindings[nodeID] = ReadBinding{TypeURL: TaeReadTypeURL, Value: value} + } + return c.BuildWithBindings(bindings) +} + +type ReadBinding struct { + TypeURL string + Value []byte + Schema []byte +} + +func (c *Candidate) BuildWithBindings(bindings map[int32]ReadBinding) ([]byte, error) { if c == nil { return nil, moerr.NewInternalErrorNoCtxf("substrait: nil candidate") } - e := exporter{query: c.query, readValues: readValues} + if len(bindings) != len(c.reads) { + return nil, moerr.NewInternalErrorNoCtx("substrait: plan has no admitted TaeRead or StreamRead for every scan") + } + for _, read := range c.reads { + binding, ok := bindings[read.NodeID] + if !ok { + return nil, moerr.NewInternalErrorNoCtxf("substrait: node %d has no admitted TaeRead or StreamRead", read.NodeID) + } + if binding.TypeURL == StreamReadTypeURL { + if read.streamErr != nil { + return nil, read.streamErr + } + if !bytes.Equal(binding.Schema, read.StreamSchema) { + return nil, moerr.NewInternalErrorNoCtxf("substrait: node %d stream schema differs from the validated scan", read.NodeID) + } + } + } + e := exporter{query: c.query, readValues: make(map[int32][]byte), readBindings: bindings, streamReads: make(map[int32]bool), expectedTypeURLs: make(map[string]bool)} relations := make([]*spb.PlanRel, 0, len(c.query.Steps)) for step, rootID := range c.query.Steps { e.stepOrdinal = int32(step) @@ -167,10 +227,16 @@ func (c *Candidate) Build(readValues map[int32][]byte) ([]byte, error) { relations = append(relations, &spb.PlanRel{RelType: &spb.PlanRel_Root{Root: &spb.RelRoot{Input: relation, Names: append([]string(nil), c.headings...)}}}) } } + expected := make([]string, 0, 2) + for _, typeURL := range []string{TaeReadTypeURL, StreamReadTypeURL} { + if e.expectedTypeURLs[typeURL] { + expected = append(expected, typeURL) + } + } p := &spb.Plan{ Version: &spb.Version{MajorNumber: 0, MinorNumber: 78, PatchNumber: 0, Producer: "matrixone"}, Relations: relations, - ExpectedTypeUrls: []string{TaeReadTypeURL}, + ExpectedTypeUrls: expected, Extensions: e.extensions(), } b, err := proto.MarshalOptions{Deterministic: true}.Marshal(p) @@ -184,14 +250,17 @@ func (c *Candidate) Build(readValues map[int32][]byte) ([]byte, error) { } type exporter struct { - query *planpb.Query - readValues map[int32][]byte - reads []Read - functions map[string]uint32 - validateOnly bool - visiting map[int32]bool - readSeen map[int32]bool - stepOrdinal int32 + query *planpb.Query + readValues map[int32][]byte + readBindings map[int32]ReadBinding + streamReads map[int32]bool + expectedTypeURLs map[string]bool + reads []Read + functions map[string]uint32 + validateOnly bool + visiting map[int32]bool + readIndexes map[int32]int + stepOrdinal int32 } func (e *exporter) node(id int32) (*spb.Rel, error) { @@ -257,6 +326,9 @@ func (e *exporter) node(id int32) (*spb.Rel, error) { if err != nil { return nil, err } + if n.NodeType == planpb.Node_TABLE_SCAN && e.streamReads[n.NodeId] { + return rel, nil + } return e.fetch(rel, n) } @@ -596,21 +668,64 @@ func (e *exporter) read(n *planpb.Node) (*spb.Rel, error) { return nil, err } if e.validateOnly { - if e.readSeen == nil { - e.readSeen = make(map[int32]bool) + if e.readIndexes == nil { + e.readIndexes = make(map[int32]int) } - if !e.readSeen[n.NodeId] { - e.readSeen[n.NodeId] = true + if index, seen := e.readIndexes[n.NodeId]; seen { + e.reads[index].Occurrences++ + } else { + streamTypes, typeErr := e.outputTypes(n.NodeId) + if typeErr != nil { + return nil, typeErr + } + streamSchema, streamNames, streamErr := namedStructFromTypes(streamTypes) + var streamSchemaBytes []byte + if streamErr == nil { + streamSchemaBytes, streamErr = proto.MarshalOptions{Deterministic: true}.Marshal(streamSchema) + } + if streamErr != nil && !IsNotEligible(streamErr) { + return nil, streamErr + } + e.readIndexes[n.NodeId] = len(e.reads) e.reads = append(e.reads, Read{ NodeID: n.NodeId, + Occurrences: 1, DatabaseID: n.TableDef.DbId, TableID: n.TableDef.TblId, SchemaVersion: n.TableDef.Version, Columns: columns, Schema: schemaBytes, + StreamSchema: streamSchemaBytes, + StreamNames: streamNames, + streamErr: streamErr, }) } } + if !e.validateOnly && e.readBindings != nil { + binding, ok := e.readBindings[n.NodeId] + if !ok || len(binding.Value) == 0 { + return nil, moerr.NewInternalErrorNoCtxf("substrait: node %d has no admitted TaeRead or StreamRead", n.NodeId) + } + if binding.TypeURL == StreamReadTypeURL { + var streamSchema spb.NamedStruct + if len(binding.Schema) == 0 || proto.Unmarshal(binding.Schema, &streamSchema) != nil { + return nil, moerr.NewInternalErrorNoCtxf("substrait: node %d has an invalid stream schema", n.NodeId) + } + e.streamReads[n.NodeId] = true + e.expectedTypeURLs[StreamReadTypeURL] = true + return &spb.Rel{RelType: &spb.Rel_Read{Read: &spb.ReadRel{ + BaseSchema: &streamSchema, + ReadType: &spb.ReadRel_ExtensionTable_{ExtensionTable: &spb.ReadRel_ExtensionTable{ + Detail: &anypb.Any{TypeUrl: StreamReadTypeURL, Value: binding.Value}, + }}, + }}}, nil + } + if binding.TypeURL != TaeReadTypeURL { + return nil, moerr.NewInternalErrorNoCtxf("substrait: node %d has an unsupported read binding", n.NodeId) + } + e.expectedTypeURLs[TaeReadTypeURL] = true + e.readValues[n.NodeId] = binding.Value + } value := e.readValues[n.NodeId] if !e.validateOnly && len(value) == 0 { return nil, moerr.NewInternalErrorNoCtxf("substrait: node %d has no admitted TaeRead", n.NodeId) @@ -1198,6 +1313,64 @@ func namedStruct(t *planpb.TableDef) (*spb.NamedStruct, error) { return &spb.NamedStruct{Names: names, Struct: &spb.Type_Struct{Types: fields, Nullability: spb.Type_NULLABILITY_REQUIRED}}, nil } +func namedStructFromTypes(typesIn []planpb.Type) (*spb.NamedStruct, []string, error) { + if len(typesIn) == 0 { + return nil, nil, moerr.NewInternalErrorNoCtx("substrait: streamed read has no columns") + } + names := make([]string, len(typesIn)) + fields := make([]*spb.Type, len(typesIn)) + for i := range typesIn { + names[i] = "col_" + strconv.Itoa(i) + field, err := nativeInputType(&typesIn[i]) + if err != nil { + return nil, nil, err + } + fields[i] = field + } + return &spb.NamedStruct{Names: append([]string(nil), names...), Struct: &spb.Type_Struct{ + Types: fields, Nullability: spb.Type_NULLABILITY_REQUIRED, + }}, names, nil +} + +// nativeInputType is the StreamRead physical MO input contract, not the semantic +// Substrait result contract. Keep the allow-list explicit: substraitType may +// map a logical result to a wider signed type (currently uint32 to i64), while +// StreamRead sends the original MO vector without conversion. Direct TaeRead +// keeps using substraitType because Sirius performs its own physical decoding. +func nativeInputType(t *planpb.Type) (*spb.Type, error) { + if t == nil { + return nil, moerr.NewInternalErrorNoCtx("substrait: missing native input type") + } + switch types.T(t.Id) { + case types.T_bool, + types.T_int8, types.T_int16, types.T_int32, types.T_int64, + types.T_float32, types.T_float64, + types.T_char, types.T_varchar, + types.T_date: + return substraitType(t) + case types.T_decimal64, types.T_decimal128: + field, err := substraitType(t) + if err != nil { + return nil, err + } + if (types.T(t.Id) == types.T_decimal64 && t.Width > 18) || + (types.T(t.Id) == types.T_decimal128 && t.Width <= 18) { + return nil, notEligiblef( + EligibilityType, + "native input type %s does not match decimal(%d,%d)", + types.T(t.Id).String(), t.Width, t.Scale, + ) + } + return field, nil + default: + return nil, notEligiblef( + EligibilityType, + "unsupported native input type %s", + types.T(t.Id).String(), + ) + } +} + func columnMapping(t *planpb.TableDef) ([]ColumnMapping, error) { if t == nil { return nil, moerr.NewInternalErrorNoCtx("substrait: missing table for column mapping") @@ -1906,7 +2079,7 @@ func validateScalarSignature(name string, out *planpb.Type, args []*planpb.Expr) return notEligiblef(EligibilityExpression, "unsupported %s signature", name) } case "substring": - if types.T(out.Id) != types.T_varchar || len(args) != 3 || !isTPCHStringType(types.T(args[0].Typ.Id)) || types.T(args[1].Typ.Id) != types.T_int64 || types.T(args[2].Typ.Id) != types.T_int64 { + if !isTPCHStringType(types.T(out.Id)) || len(args) != 3 || !isTPCHStringType(types.T(args[0].Typ.Id)) || types.T(args[1].Typ.Id) != types.T_int64 || types.T(args[2].Typ.Id) != types.T_int64 { return notEligiblef(EligibilityExpression, "unsupported substring signature") } } diff --git a/pkg/sql/plan/substrait/protocol.go b/pkg/sql/plan/substrait/protocol.go index fa8693a7b9044..25150af924fbb 100644 --- a/pkg/sql/plan/substrait/protocol.go +++ b/pkg/sql/plan/substrait/protocol.go @@ -24,19 +24,63 @@ import ( ) const ( - ProtocolVersion = 1 - TaeReadProtocolVersion = 2 - maxTaeReadSize = 16 << 10 - maxResolveRequestSize = maxTaeReadSize + maxCanonicalSchemaSize + 32 + ProtocolVersion = 1 + TaeReadProtocolVersion = 2 + StreamReadProtocolVersion = 1 + maxTaeReadSize = 16 << 10 + maxResolveRequestSize = maxTaeReadSize + maxCanonicalSchemaSize + 32 // GC protection stores the expiry as Unix nanoseconds. Keep the wire value // within that signed range before any consumer converts it through // time.UnixMilli(...).UnixNano(). maxTaeReadExpiryUnixMS = uint64(math.MaxInt64 / 1_000_000) - CapabilityDocument = `{"protocol_version":3,"substrait_version":"0.78.0","tae_read_protocol_version":2,"tae_read_feature_bits":0,"operators":["read","filter","project","aggregate","sort","fetch","join","reference"],"join_types":["inner","left","right","left_semi","left_anti","right_semi","right_anti"],"expressions":["literal","selection","scalar_function","cast","if_then","singular_or_list"],"types":["bool","i8","i16","i32","i64","fp32","fp64","varchar","decimal","date"],"semantic_registry":"exact-mo-bound-overload-and-tpch-family-v1","scalar_functions":["and","or","not","equal","not_equal","lt","lte","gt","gte","is_null","is_not_null","is_not_distinct_from","add","subtract","multiply","divide","modulus","between","like","starts_with","substring","extract"],"aggregate_functions":["count","sum","min","max","avg"],"transport":"arrow-flight","sirius_execution_contract":1,"max_plan_bytes":16777216}` + CapabilityDocument = `{"protocol_version":5,"substrait_version":"0.78.0","tae_read_protocol_version":2,"tae_read_feature_bits":0,"stream_read_protocol_version":1,"stream_read_feature_bits":0,"operators":["read","filter","project","aggregate","sort","fetch","join","reference"],"join_types":["inner","left","right","left_semi","left_anti","right_semi","right_anti"],"expressions":["literal","selection","scalar_function","cast","if_then","singular_or_list"],"types":["bool","i8","i16","i32","i64","fp32","fp64","varchar","decimal","date"],"semantic_registry":"exact-mo-bound-overload-and-tpch-family-v1","scalar_functions":["and","or","not","equal","not_equal","lt","lte","gt","gte","is_null","is_not_null","is_not_distinct_from","add","subtract","multiply","divide","modulus","between","like","starts_with","substring","extract"],"aggregate_functions":["count","sum","min","max","avg"],"transport":"arrow-flight","stream_input_transport":"flight-do-put-mo-native","result_transport":"flight-doget-mo-native","mo_native_result_schema_version":1,"gpu_mo_batch_ingress":true,"adaptive_gpu_result_pack":true,"mo_native_batch_codec_version":1,"mo_type_size":16,"mo_varlena_size":24,"mo_native_endian":"little","stream_input_ack":"ready-consumed-final-v1","stream_input_slots_per_read":1,"max_stream_inputs":16,"max_stream_input_batch_bytes":4194304,"sirius_execution_contract":1,"max_plan_bytes":16777216}` ) var CapabilityHash = sha256.Sum256([]byte(CapabilityDocument)) +// StreamRead identifies one post-scan relation whose MO-native batches are +// supplied over the execution's authenticated Flight input stream. +type StreamRead struct { + ProtocolVersion uint32 + FeatureBits uint64 + StreamRef, QueryID []byte + AccountID uint64 + SnapshotTS, SchemaDigest, CapabilityHash []byte + ExpiresAtUnixMS uint64 +} + +func (r *StreamRead) Validate(nowUnixMS uint64) error { + if r == nil || r.ProtocolVersion != StreamReadProtocolVersion || r.FeatureBits != 0 { + return moerr.NewInternalErrorNoCtx("invalid StreamRead protocol") + } + if len(r.StreamRef) != 32 || len(r.QueryID) != 16 || len(r.SnapshotTS) != 12 || + len(r.SchemaDigest) != sha256.Size || len(r.CapabilityHash) != sha256.Size { + return moerr.NewInternalErrorNoCtx("invalid StreamRead identity or digest length") + } + if r.ExpiresAtUnixMS <= nowUnixMS || r.ExpiresAtUnixMS > maxTaeReadExpiryUnixMS || + !equalBytes(r.CapabilityHash, CapabilityHash[:]) { + return moerr.NewInternalErrorNoCtx("invalid or expired StreamRead") + } + return nil +} + +func MarshalStreamRead(r *StreamRead) ([]byte, error) { + if err := r.Validate(0); err != nil { + return nil, err + } + var b []byte + b = appendUint(b, 1, uint64(r.ProtocolVersion)) + b = appendUint(b, 2, r.FeatureBits) + b = appendBytes(b, 3, r.StreamRef) + b = appendBytes(b, 4, r.QueryID) + b = appendRequiredUint(b, 5, r.AccountID) + b = appendBytes(b, 6, r.SnapshotTS) + b = appendBytes(b, 7, r.SchemaDigest) + b = appendBytes(b, 8, r.CapabilityHash) + b = appendUint(b, 9, r.ExpiresAtUnixMS) + return b, nil +} + // TaeRead is the exact matrixone.sirius.v1.TaeRead v1 wire contract. type TaeRead struct { ProtocolVersion uint32 diff --git a/pkg/sql/plan/substrait/substrait_test.go b/pkg/sql/plan/substrait/substrait_test.go index c781049a197d7..3506c735f1893 100644 --- a/pkg/sql/plan/substrait/substrait_test.go +++ b/pkg/sql/plan/substrait/substrait_test.go @@ -62,6 +62,139 @@ func TestExportBuildSupportedSubset(t *testing.T) { require.Equal(t, TaeReadTypeURL, p.Relations[0].GetRoot().Input.GetFilter().Input.GetRead().GetExtensionTable().Detail.TypeUrl) } +func TestBuildStreamReadReplacesTheCompleteScanNode(t *testing.T) { + q := scanQuery() + q.Nodes[0].FilterList = []*planpb.Expr{fn(">", boolType(), col(0), i64(7))} + q.Nodes[0].ProjectList = []*planpb.Expr{col(0)} + q.Nodes[0].Limit = i64(10) + candidate, err := Export(q) + require.NoError(t, err) + reads := candidate.Reads() + require.Len(t, reads, 1) + require.Equal(t, []string{"col_0"}, reads[0].StreamNames) + digest := sha256.Sum256(reads[0].StreamSchema) + wire, err := MarshalStreamRead(&StreamRead{ + ProtocolVersion: StreamReadProtocolVersion, + StreamRef: bytes.Repeat([]byte{1}, 32), QueryID: bytes.Repeat([]byte{2}, 16), + SnapshotTS: bytes.Repeat([]byte{3}, 12), SchemaDigest: digest[:], + CapabilityHash: CapabilityHash[:], ExpiresAtUnixMS: 2000, + }) + require.NoError(t, err) + _, err = candidate.BuildWithBindings(map[int32]ReadBinding{0: { + TypeURL: StreamReadTypeURL, Value: wire, Schema: append(append([]byte(nil), reads[0].StreamSchema...), 0), + }}) + require.ErrorContains(t, err, "stream schema differs") + built, err := candidate.BuildWithBindings(map[int32]ReadBinding{0: { + TypeURL: StreamReadTypeURL, Value: wire, Schema: reads[0].StreamSchema, + }}) + require.NoError(t, err) + var plan spb.Plan + require.NoError(t, proto.Unmarshal(built, &plan)) + require.Equal(t, []string{StreamReadTypeURL}, plan.ExpectedTypeUrls) + read := plan.Relations[0].GetRoot().Input.GetRead() + require.NotNil(t, read) + require.Equal(t, StreamReadTypeURL, read.GetExtensionTable().Detail.TypeUrl) + require.Equal(t, []string{"col_0"}, read.BaseSchema.Names) +} + +func TestStreamReadValidationRejectsEveryInvalidIdentityClass(t *testing.T) { + valid := func() *StreamRead { + return &StreamRead{ + ProtocolVersion: StreamReadProtocolVersion, + StreamRef: bytes.Repeat([]byte{1}, 32), + QueryID: bytes.Repeat([]byte{2}, 16), + SnapshotTS: bytes.Repeat([]byte{3}, 12), + SchemaDigest: bytes.Repeat([]byte{4}, sha256.Size), + CapabilityHash: append([]byte(nil), CapabilityHash[:]...), + ExpiresAtUnixMS: 2000, + } + } + require.ErrorContains(t, (*StreamRead)(nil).Validate(1000), "protocol") + _, err := MarshalStreamRead(nil) + require.ErrorContains(t, err, "protocol") + + for _, tc := range []struct { + name string + configure func(*StreamRead) + want string + }{ + {name: "version", configure: func(r *StreamRead) { r.ProtocolVersion++ }, want: "protocol"}, + {name: "features", configure: func(r *StreamRead) { r.FeatureBits = 1 }, want: "protocol"}, + {name: "stream identity", configure: func(r *StreamRead) { r.StreamRef = r.StreamRef[:31] }, want: "identity"}, + {name: "query identity", configure: func(r *StreamRead) { r.QueryID = r.QueryID[:15] }, want: "identity"}, + {name: "snapshot", configure: func(r *StreamRead) { r.SnapshotTS = r.SnapshotTS[:11] }, want: "identity"}, + {name: "schema digest", configure: func(r *StreamRead) { r.SchemaDigest = r.SchemaDigest[:31] }, want: "identity"}, + {name: "capability length", configure: func(r *StreamRead) { r.CapabilityHash = r.CapabilityHash[:31] }, want: "identity"}, + {name: "expired", configure: func(r *StreamRead) { r.ExpiresAtUnixMS = 1000 }, want: "expired"}, + {name: "expiry overflow", configure: func(r *StreamRead) { r.ExpiresAtUnixMS = maxTaeReadExpiryUnixMS + 1 }, want: "expired"}, + {name: "capability mismatch", configure: func(r *StreamRead) { r.CapabilityHash[0] ^= 1 }, want: "expired"}, + } { + t.Run(tc.name, func(t *testing.T) { + streamRead := valid() + tc.configure(streamRead) + require.ErrorContains(t, streamRead.Validate(1000), tc.want) + }) + } +} + +func TestNativeInputTypeRejectsSemanticOnlyUint32Mapping(t *testing.T) { + uint32Type := planpb.Type{Id: int32(types.T_uint32)} + semantic, err := substraitType(&uint32Type) + require.NoError(t, err) + require.NotNil(t, semantic.GetI64(), "uint32 results still use signed i64 transport") + + _, err = nativeInputType(&uint32Type) + require.True(t, IsNotEligible(err)) + require.ErrorContains(t, err, "unsupported native input type INT UNSIGNED") + _, err = nativeInputType(nil) + require.ErrorContains(t, err, "missing native input type") + + for _, typ := range []planpb.Type{ + {Id: int32(types.T_bool)}, + {Id: int32(types.T_int8)}, {Id: int32(types.T_int16)}, + {Id: int32(types.T_int32)}, {Id: int32(types.T_int64)}, + {Id: int32(types.T_float32)}, {Id: int32(types.T_float64)}, + {Id: int32(types.T_char), Width: 8}, {Id: int32(types.T_varchar), Width: 32}, + {Id: int32(types.T_decimal64), Width: 18, Scale: 2}, + {Id: int32(types.T_decimal128), Width: 38, Scale: 4}, + {Id: int32(types.T_date)}, + } { + _, err = nativeInputType(&typ) + require.NoError(t, err, types.T(typ.Id).String()) + } + for _, typ := range []planpb.Type{ + {Id: int32(types.T_decimal64), Width: 19, Scale: 2}, + {Id: int32(types.T_decimal128), Width: 18, Scale: 2}, + } { + _, err = nativeInputType(&typ) + require.True(t, IsNotEligible(err)) + require.ErrorContains(t, err, "does not match decimal") + } +} + +func TestUint32CandidateKeepsDirectTaeReadAndRejectsOnlyStreamRead(t *testing.T) { + query := scanQuery() + query.Nodes[0].TableDef.Cols[0].Typ = planpb.Type{Id: int32(types.T_uint32)} + + candidate, err := Export(query) + require.NoError(t, err) + direct, err := candidate.Build(map[int32][]byte{0: {1}}) + require.NoError(t, err) + var exported spb.Plan + require.NoError(t, proto.Unmarshal(direct, &exported)) + require.NotNil(t, exported.Relations[0].GetRoot().Input.GetRead().BaseSchema.Struct.Types[0].GetI64(), + "the direct TaeRead contract keeps the intentional uint32-to-i64 widening") + + _, err = candidate.StreamReads() + require.True(t, IsNotEligible(err)) + require.ErrorContains(t, err, "unsupported native input type INT UNSIGNED") + _, err = candidate.BuildWithBindings(map[int32]ReadBinding{0: { + TypeURL: StreamReadTypeURL, Value: []byte{1}, + }}) + require.True(t, IsNotEligible(err)) + require.ErrorContains(t, err, "unsupported native input type INT UNSIGNED") +} + func TestExportRejectsFullOuterBeforeSnapshotAccess(t *testing.T) { q := scanQuery() q.Nodes = append(q.Nodes, &planpb.Node{NodeId: 1, NodeType: planpb.Node_JOIN, JoinType: planpb.Node_OUTER, Children: []int32{0, 0}}) @@ -156,6 +289,7 @@ func TestSharedScanNodeProducesOneAdmissionRead(t *testing.T) { candidate, err := Export(q) require.NoError(t, err) require.Len(t, candidate.Reads(), 1) + require.Equal(t, uint32(2), candidate.Reads()[0].Occurrences) } func TestProjectEmitsOnlyMOProjectColumns(t *testing.T) { @@ -433,7 +567,7 @@ func TestBoundInt64SumIsAdvertisedWithDecimalResult(t *testing.T) { } func TestCapabilityHashMatchesSidecarContract(t *testing.T) { - require.Equal(t, "6f788b3d6665ecdd1ac734043fb757968893f14fd7d197fabcfa287764ee6bad", hex.EncodeToString(CapabilityHash[:])) + require.Equal(t, "e72e3c64e9519fb2824c7773ea40564a2c76e7ec36e46560d7c6de7d1444fc11", hex.EncodeToString(CapabilityHash[:])) } func TestBoundDecimalUnaryMinusLowersToSubtractFromTypedZero(t *testing.T) { @@ -470,6 +604,15 @@ func TestBoundDecimalUnaryMinusLowersToSubtractFromTypedZero(t *testing.T) { require.NotNil(t, subtract.Arguments[0].GetValue().GetLiteral().GetDecimal()) } +func TestCharSubstringResultUsesTPCHStringFamily(t *testing.T) { + charType := planpb.Type{Id: int32(types.T_char), Width: 15} + integerType := planpb.Type{Id: int32(types.T_int64), NotNullable: true} + args := []*planpb.Expr{{Typ: charType}, {Typ: integerType}, {Typ: integerType}} + // CHAR is intentionally transported as Substrait string/VARCHAR and + // restored to the negotiated MatrixOne physical result type. + require.NoError(t, validateScalarSignature("substring", &charType, args)) +} + func TestFetchAcceptsBoundUint64AndPreservesAbsentCount(t *testing.T) { q := scanQuery() q.Nodes[0].Offset = u64(2) diff --git a/proto/sidecar/v1/sidecar.proto b/proto/sidecar/v1/sidecar.proto index 0cb5312299564..64d72c431ed61 100644 --- a/proto/sidecar/v1/sidecar.proto +++ b/proto/sidecar/v1/sidecar.proto @@ -23,6 +23,44 @@ message ExecuteSubstraitRequest { uint64 max_batch_bytes = 4; uint64 deadline_unix_ms = 5; bytes plan = 6; + bytes query_id = 7; + bytes idempotency_key = 8; + optional uint64 account_id = 9; + uint64 max_input_batch_bytes = 10; + bytes result_schema = 11; +} + +message NativeResultSchema { + uint32 version = 1; + repeated NativeResultColumn columns = 2; +} + +message NativeResultColumn { + string name = 1; + uint32 oid = 2; + int32 width = 3; + int32 scale = 4; + uint32 charset = 5; + bool not_nullable = 6; +} + +message UploadInputRequest { + bytes ticket = 1; + bytes stream_ref = 2; +} + +message UploadInputAck { + uint64 acknowledged_batches = 1; + uint64 rows = 2; + uint64 bytes = 3; + bool complete = 4; + bool not_needed = 5; + bool ready = 6; +} + +message CancelExecutionRequest { + bytes ticket = 1; + bytes idempotency_key = 2; } message ResolveTaeReadRequest {