From 5d4699db1cb3a01dec3b3ec77cef67887051d65a Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 10 Aug 2026 10:45:37 +0800 Subject: [PATCH 01/92] docs: refine coding agent maps (#8790) * docs: refine coding agent maps Signed-off-by: Dennis Zhuang * docs: trim license header guidance Signed-off-by: Dennis Zhuang * docs: update README links and project status Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- .agents/README.md | 13 +++-- .agents/architecture-invariants.md | 47 +++++++++--------- .agents/generated-files.md | 14 ++++-- .../skills/greptimedb-release-note/SKILL.md | 7 +-- AGENTS.md | 49 +++++++++++++++++-- README.md | 4 +- src/common/meta/AGENTS.md | 38 ++++++++++++++ src/flow/AGENTS.md | 21 +++++--- src/frontend/AGENTS.md | 9 ++-- src/metric-engine/AGENTS.md | 8 +-- src/operator/AGENTS.md | 40 +++++++++++++++ src/query/AGENTS.md | 39 +++++++++++++++ src/servers/AGENTS.md | 37 ++++++++++++++ tests/perf/AGENTS.md | 5 +- 14 files changed, 273 insertions(+), 58 deletions(-) create mode 100644 src/common/meta/AGENTS.md create mode 100644 src/operator/AGENTS.md create mode 100644 src/query/AGENTS.md create mode 100644 src/servers/AGENTS.md diff --git a/.agents/README.md b/.agents/README.md index d55aac50f06..8d21de0f1a2 100644 --- a/.agents/README.md +++ b/.agents/README.md @@ -11,16 +11,23 @@ Shared agent skills live in `.agents/skills`. - Add or update shared skills under `.agents/skills//SKILL.md`. - If a new skill does not appear, restart the agent or start a new thread. -## Per-crate guides +## Per-directory guides -Hot crates carry an `AGENTS.md` next to their code as a navigation aid (module -map, read/write paths, change-coupling points, test commands, gotchas): +Read every `AGENTS.md` from the repository root down to the path being changed. +The root guide applies everywhere; nested guides add rules and navigation for +their subtree. +- [`src/common/meta/AGENTS.md`](../src/common/meta/AGENTS.md) — metadata keys, KV backends, DDL procedures, and caches +- [`src/query/AGENTS.md`](../src/query/AGENTS.md) — query planning, optimization, and distributed execution +- [`src/servers/AGENTS.md`](../src/servers/AGENTS.md) — wire protocols and network servers +- [`src/operator/AGENTS.md`](../src/operator/AGENTS.md) — statement, DDL/DML, and write orchestration - [`src/mito2/AGENTS.md`](../src/mito2/AGENTS.md) — primary time-series storage engine - [`src/metric-engine/AGENTS.md`](../src/metric-engine/AGENTS.md) — metrics engine (logical/physical regions) - [`src/flow/AGENTS.md`](../src/flow/AGENTS.md) — stream processing / continuous aggregation - [`src/frontend/AGENTS.md`](../src/frontend/AGENTS.md) — request entry point and orchestration - [`src/meta-srv/AGENTS.md`](../src/meta-srv/AGENTS.md) — metadata and cluster coordination +- [`tests/compatibility/AGENTS.md`](../tests/compatibility/AGENTS.md) — persisted/wire compatibility cases +- [`tests/perf/AGENTS.md`](../tests/perf/AGENTS.md) — query regression harness and case DSL ## Architecture invariants diff --git a/.agents/architecture-invariants.md b/.agents/architecture-invariants.md index b7ce954c972..78c78f3a976 100644 --- a/.agents/architecture-invariants.md +++ b/.agents/architecture-invariants.md @@ -20,15 +20,18 @@ metadata KV values (`common-meta` keys, metric-engine metadata), and gRPC messages. A node running an old version may read data written by a new one and vice versa. -- Add fields, don't repurpose or reorder them. For serde types use - `#[serde(default)]` / `#[serde(alias)]`; never change the meaning of an existing - field or the discriminant of an existing enum variant. -- Monotonic version counters (e.g. the mito2 manifest version) only ever move - forward — never reset or skip. +- Preserve the contract of the actual encoding: Protobuf field numbers/types, + serde field and variant names, positional field order, and explicitly encoded + enum discriminants. Do not treat Rust declaration order as a wire contract + for name-based formats such as JSON. +- Use `#[serde(default)]` for newly optional data and `#[serde(alias)]` for a + compatible rename. Do not reset or reuse versions within one persisted + history. - When a change touches a persisted or wire format, add a case to the - compatibility test suite. See `docs/rfcs/2025-07-04-compatibility-test-framework.md` - — compatibility has been broken on releases before (v0.14.1, v0.15.1) precisely - because this step was skipped. + compatibility test suite. Follow [`tests/compatibility/README.md`](../tests/compatibility/README.md) + and [`tests/compatibility/AGENTS.md`](../tests/compatibility/AGENTS.md). The + original compatibility RFC describes a superseded case format and is not the + current runbook. - Wire types are generated from the external `greptime-proto` crate; change the format there first, then bump the dependency (see invariant 6's pattern). @@ -54,9 +57,9 @@ The workspace is layered; dependencies point downward only. Runtimes are partitioned by workload so one workload can't starve another. They live in `common-runtime` (`src/common/runtime/`). -- Use the categorized spawns — `spawn_global`, `spawn_query`, `spawn_ingest`, - `spawn_compact`, `spawn_hb` — instead of constructing your own tokio runtime, - and pick the category that matches the work. +- Runtime construction belongs in process bootstrap or test/benchmark harnesses. + Product components use `spawn_global`, `spawn_query`, `spawn_ingest`, + `spawn_compact`, or `spawn_hb` instead of a component-local Tokio runtime. - Run CPU-bound or synchronous-blocking work via `spawn_blocking_*`; never do heavy CPU or blocking syscalls directly inside an async task. - Never call `block_on*` from inside an async context or an engine worker — it @@ -85,23 +88,19 @@ unfinished work merge without freezing it into the stable config surface. When you stabilize such a feature, drop the prefix and document the migration. -## 6. DataFusion is a pinned fork — two sections, two forms +## 6. DataFusion is a pinned fork — workspace dependencies plus patches GreptimeDB uses a fork at `GreptimeTeam/datafusion`, wired up in the root -`Cargo.toml` through **two sections that hold different things**: +`Cargo.toml` through two sections: -- `[workspace.dependencies]` pins each DataFusion sub-crate to an **exact - crates.io version** (e.g. `datafusion = "=53.1.0"`). -- `[patch.crates-io]` redirects those same crates to the **fork at a git rev** - (e.g. `datafusion = { git = ".../GreptimeTeam/datafusion.git", rev = "..." }`). +- `[workspace.dependencies]` pins directly referenced sub-crates to an exact + crates.io version. +- `[patch.crates-io]` redirects resolution to the fork. It may also contain + transitive-only sub-crates with no workspace dependency. -So: - -- Adding a new DataFusion sub-crate dependency means adding it to **both** - sections — the `=` entry under `[workspace.dependencies]` and the - matching git-rev patch under `[patch.crates-io]`. -- Upgrading DataFusion means bumping the version in `[workspace.dependencies]` - **and** the rev in `[patch.crates-io]` together, for all of them. +Add a direct dependency to `[workspace.dependencies]` and patch it if needed; a +transitive-only override needs only the patch. Upgrade all exact pins and fork +revisions together. ## 7. Enterprise-gated code is licensed differently from the rest diff --git a/.agents/generated-files.md b/.agents/generated-files.md index 1a16fcb28cf..23f7e81daf3 100644 --- a/.agents/generated-files.md +++ b/.agents/generated-files.md @@ -35,7 +35,14 @@ Edit the example TOMLs and/or the template, then regenerate. Do not edit ## Grafana dashboards -Generated artifacts; regenerate rather than editing by hand. +For the metrics dashboards, edit +`grafana/dashboards/metrics/cluster/dashboard.json`. The generator derives the +standalone `dashboard.json` by removing instance filters, then produces +`dashboard.yaml` and `dashboard.md` for both cluster and standalone variants. +Do not hand-edit those five derived files. + +`grafana/dashboards/events/dashboard.json` and +`grafana/dashboards/logs/dashboard.json` are not outputs of this generator. ```bash make dashboards @@ -56,5 +63,6 @@ dependency. ## License headers -Managed by the license-header tooling and the pre-commit hooks; let the tooling -add/fix them rather than editing headers by hand. +License headers are checked by `korandoru/hawkeye@v5` in +`.github/workflows/checks.yml`. For enterprise-gated files, also update both +license lists and run `make check-enterprise-license`. diff --git a/.agents/skills/greptimedb-release-note/SKILL.md b/.agents/skills/greptimedb-release-note/SKILL.md index a5cb47e9b02..893b99445ce 100644 --- a/.agents/skills/greptimedb-release-note/SKILL.md +++ b/.agents/skills/greptimedb-release-note/SKILL.md @@ -167,9 +167,10 @@ The release note is also published as a blog post in **`GreptimeTeam/docs`**. git -C fetch origin main git -C worktree add -b chore/X.Y.Z-release-note /tmp/docs-release-note origin/main ``` -- **The PR body must follow the docs repo's template** (`.github/pull_request_template.md` — - "What's Changed in this PR" + a Checklist). Fill in the description; leave checklist boxes - for the reviewer. +- **Read and follow the current docs repo PR template** + (`.github/pull_request_template.md`) from the checkout used for the release. + Fill in every required section and leave reviewer-owned checklist boxes unchecked; + do not rely on hard-coded section names from an older template. - Commit with sign-off, push, open a **draft** PR, then remove the worktree: ``` git -C /tmp/docs-release-note add blog/release-X-Y-Z.md diff --git a/AGENTS.md b/AGENTS.md index a00889b855f..c99856c2941 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -18,7 +18,8 @@ provides sub-second querying at PB scale with high cost efficiency. | Build (debug) | `make build` | | Build (release) | `make build RELEASE=true` | | Run standalone | `cargo run -- standalone start` | -| Test | `cargo nextest run` (preferred over `cargo test`) | +| Targeted Rust tests | `cargo nextest run -p ` (preferred over `cargo test`) | +| Full local Rust suite | `make test` | | SQL tests | `cargo sqlness bare` (single case: `cargo sqlness bare -t `) | | Format | `make fmt` (and `make fmt-toml` for TOML) | | Lint | `make clippy` (= `cargo clippy --workspace --all-targets --all-features -- -D warnings`) | @@ -43,14 +44,30 @@ GreptimeDB is a Cargo workspace rooted at the repository; most crates live under - **Shared**: `src/common/`, `src/datatypes/`, `src/store-api/` (engine contract), `src/catalog/`, `src/table/` -Hot crates carry their own `AGENTS.md` with a module map, read/write paths, -change-coupling points, and gotchas: +Before editing a path, read every `AGENTS.md` from the repository root down to +that path. The root guide always applies; a nested guide adds or overrides rules +for its subtree. Changes spanning multiple subtrees must follow every applicable +guide. +`AGENTS.md` files are maps, not manuals. Keep them short and stable: record +where code lives, important boundaries or change-coupling, high-cost gotchas, +and the validation entry point. Implementation details belong in code. Do not +copy exhaustive test lists or workflow logic; link to the source of truth. + +High-change areas and specialized test suites carry their own `AGENTS.md` with +a module map, read/write paths, change-coupling points, and gotchas: + +- [`src/common/meta/AGENTS.md`](src/common/meta/AGENTS.md) +- [`src/query/AGENTS.md`](src/query/AGENTS.md) +- [`src/servers/AGENTS.md`](src/servers/AGENTS.md) +- [`src/operator/AGENTS.md`](src/operator/AGENTS.md) - [`src/mito2/AGENTS.md`](src/mito2/AGENTS.md) - [`src/metric-engine/AGENTS.md`](src/metric-engine/AGENTS.md) - [`src/flow/AGENTS.md`](src/flow/AGENTS.md) - [`src/frontend/AGENTS.md`](src/frontend/AGENTS.md) - [`src/meta-srv/AGENTS.md`](src/meta-srv/AGENTS.md) +- [`tests/compatibility/AGENTS.md`](tests/compatibility/AGENTS.md) +- [`tests/perf/AGENTS.md`](tests/perf/AGENTS.md) ## Read before changing code @@ -71,11 +88,35 @@ change-coupling points, and gotchas: - Error handling: `src/common/error/` (`ErrorExt`, `StatusCode`) - Protocol implementations: `src/servers/src/` +## Worktree safety + +- Check `git status --short` before editing. Preserve unrelated tracked changes + and untracked files; re-read files changed by another process. +- Do not rewrite history, force-push, or remove files in bulk unless the user + explicitly requests it. +- Update generated artifacts only through the generators documented in + [`.agents/generated-files.md`](.agents/generated-files.md). + +## Validation by change type + +Use the narrowest command that covers the change, then expand only when its +blast radius requires it. + +| Change | Minimum validation | +| --- | --- | +| One Rust crate | `cargo nextest run -p ` | +| Cross-workspace Rust behavior | `make test`; inspect `.github/workflows/rust.yml` when CI parity matters | +| SQL parsing, planning, execution, or output | `cargo sqlness bare -t `; inspect regenerated `.result` files | +| Persisted metadata or wire format | Add/run a case under `tests/compatibility/`; follow its `README.md` and `AGENTS.md` | +| Public configuration | Update example TOMLs, loading/serialization snapshots, and docs; run `make config-docs` | +| Query regression harness or DSL | Follow `tests/perf/AGENTS.md` | +| Enterprise-gated code | Build/test with `--features enterprise` where applicable and run `make check-enterprise-license` | + ## Before opening a PR 1. `make fmt` 2. `make clippy` -3. `make test` (or `cargo nextest run`) +3. `make test` 4. `make check-udeps` (run `make fix-udeps` if it reports unused dependencies). 5. If you added or changed a public configuration option, update the applicable example TOMLs, configuration-loading and serialized-config snapshot tests, diff --git a/README.md b/README.md index 019a7c848ba..21980b1a22b 100644 --- a/README.md +++ b/README.md @@ -203,13 +203,13 @@ See the [Contribution Guidelines](CONTRIBUTING.md) for the full developer workfl - **Kubernetes**: [GreptimeDB Operator](https://github.com/GreptimeTeam/greptimedb-operator) - **Helm Charts**: [Greptime Helm Charts](https://github.com/GreptimeTeam/helm-charts) - **Dashboard**: [Web UI](https://github.com/GreptimeTeam/dashboard) -- **gRPC Ingester**: [Go](https://github.com/GreptimeTeam/greptimedb-ingester-go), [Java](https://github.com/GreptimeTeam/greptimedb-ingester-java), [C++](https://github.com/GreptimeTeam/greptimedb-ingester-cpp), [Erlang](https://github.com/GreptimeTeam/greptimedb-ingester-erl), [Rust](https://github.com/GreptimeTeam/greptimedb-ingester-rust), [.NET](https://github.com/GreptimeTeam/greptimedb-ingester-dotnet) +- **gRPC Ingester**: [Go](https://github.com/GreptimeTeam/greptimedb-ingester-go), [Java](https://github.com/GreptimeTeam/greptimedb-ingester-java), [C++](https://github.com/GreptimeTeam/greptimedb-ingester-cpp), [Erlang](https://github.com/GreptimeTeam/greptimedb-ingester-erl), [Rust](https://github.com/GreptimeTeam/greptimedb-ingester-rust), [.NET](https://github.com/GreptimeTeam/greptimedb-ingester-dotnet), [TypeScript](https://github.com/GreptimeTeam/greptimedb-ingester-ts) - **Grafana Data Source**: [GreptimeDB Grafana data source plugin](https://github.com/GreptimeTeam/greptimedb-grafana-datasource) - **Grafana Dashboard**: [Official Dashboard for monitoring](https://github.com/GreptimeTeam/greptimedb/blob/main/grafana/README.md) ## Project Status -GreptimeDB is at [v1.0 GA](https://github.com/GreptimeTeam/greptimedb/releases/tag/v1.0.0) with stable APIs and regular releases. It runs in production at scale — [OceanBase Cloud](https://greptime.com/blogs/2025-07-22-user-case-obcloud-log-management-greptimedb) operates 80+ GreptimeDB clusters managing 300 TB of logs, cutting log storage cost by 60% after migrating from Grafana Loki. See more in [case studies](https://greptime.com/blogs/?category=Use%20Case). +GreptimeDB is generally available, with stable APIs and regular releases. It runs in production at scale — [OceanBase Cloud](https://greptime.com/blogs/2025-07-22-user-case-obcloud-log-management-greptimedb) operates 80+ GreptimeDB clusters managing 300 TB of logs, cutting log storage cost by 60% after migrating from Grafana Loki. See more in [case studies](https://greptime.com/blogs/?category=Use%20Case). Read the [v1.0 highlights](https://greptime.com/blogs/2025-11-05-greptimedb-v1-highlights) and [2026 roadmap](https://greptime.com/blogs/2026-02-11-greptimedb-roadmap-2026), or browse the [version reference](https://docs.greptime.com/nightly/reference/about-greptimedb-version). diff --git a/src/common/meta/AGENTS.md b/src/common/meta/AGENTS.md new file mode 100644 index 00000000000..fd5da08a933 --- /dev/null +++ b/src/common/meta/AGENTS.md @@ -0,0 +1,38 @@ +# common-meta — Agent & Contributor Guide + +Navigation map for `src/common/meta`. Repo-wide invariants: +[`.agents/architecture-invariants.md`](../../../.agents/architecture-invariants.md). + +`common-meta` contains shared metadata keys and values, KV backends, caches, +and durable DDL procedures. Metasrv process and service wiring live in +`src/meta-srv`. + +## Module map + +| Area | Path | Entry point | +| --- | --- | --- | +| Metadata model | `src/common/meta/src/key.rs`, `src/common/meta/src/key/` | Typed keys, values, and `TableMetadataManager` | +| KV storage | `src/common/meta/src/kv_backend.rs`, `src/common/meta/src/kv_backend/` | `KvBackend`, transactions, memory/etcd/RDS backends | +| DDL procedures | `src/common/meta/src/ddl.rs`, `src/common/meta/src/ddl/`, `src/common/meta/src/ddl_manager.rs` | Durable DDL state machines and task dispatch | +| Caches | `src/common/meta/src/cache.rs`, `src/common/meta/src/cache/`, `src/common/meta/src/cache_invalidator.rs` | Metadata caches and invalidation | +| RPC types | `src/common/meta/src/rpc.rs`, `src/common/meta/src/rpc/` | Shared metasrv request/response types | +| Recovery | `src/common/meta/src/reconciliation.rs`, `src/common/meta/src/reconciliation/` | Catalog/table/region reconciliation | + +## Change coupling + +- Key/value encoding changes require backward-compatible decoding and a case in + `tests/compatibility/`. +- Multi-key metadata changes must preserve transaction boundaries and cache + invalidation. +- Procedure changes must preserve persisted state and `TYPE_NAME`; new + procedures need loader registration in `src/common/meta/src/ddl_manager.rs`. +- `KvBackend` behavior changes should extend the shared tests in + `src/common/meta/src/kv_backend/test.rs`. + +## Testing + +```bash +cargo nextest run -p common-meta +``` + +Backend-specific coverage may also need `pg_kvbackend` or `mysql_kvbackend`. diff --git a/src/flow/AGENTS.md b/src/flow/AGENTS.md index aba019e02a0..0ac0aa0e243 100644 --- a/src/flow/AGENTS.md +++ b/src/flow/AGENTS.md @@ -8,15 +8,19 @@ Repo-wide rules that apply here: [`.agents/architecture-invariants.md`](../../.a ## What this crate does Flownode is the stream-processing engine behind continuous aggregation / -materialized views. It runs in two modes: +materialized views. It has two execution paths: -- **Batching mode** (the default and the actively developed path): splits data +- **Batching mode** (the actively developed path): splits data into time windows and periodically runs aggregation SQL through the frontend, writing results back to a sink table. -- **Streaming mode** (the older dataflow engine): an incremental DFIR/dataflow +- **Streaming mode** (the legacy dataflow path): an incremental DFIR/dataflow compute graph that processes row-level diffs. -A flow without an explicit `flow_type` is created as **batching**. +Users cannot select a mode directly: `flow_type` is a reserved internal option. +`StatementExecutor::determine_flow_type` in `src/operator/src/statement/ddl.rs` +owns current mode selection. `FlowDualEngine` defaults missing internal +`flow_type` metadata to batching for compatibility. Read those paths before +changing routing rules. ## Module map @@ -64,7 +68,9 @@ Flow metadata lives in `common-meta`, not here: `compute/render.rs` (streaming) and ensure batching SQL handles it. - **Persisted flow metadata**: keep `FlowInfoValue` backward compatible (`serde(default)` / `serde(alias)`). -- A change to one engine often needs the mirror change in the other. +- **Shared routing, metadata, or sink contracts**: check both engines. An + engine-specific implementation change does not automatically need a mirror + change in the other path. ## Testing @@ -76,9 +82,8 @@ Helpers in `src/flow/src/test_utils.rs` (test context, test query engine). ## Gotchas -- Batching vs streaming differ a lot in latency, debuggability, and code path — - confirm which mode a flow uses before reasoning about it. Batching is the - default and the primary target. +- Batching vs streaming differ in latency, state, and execution. Confirm the + selected mode before reasoning about a flow. - Streaming workers are `!Send`; cross-thread interaction goes through `WorkerHandle`, not the worker directly. - Internal flow timestamps (`repr::Timestamp`, ms) are not necessarily the diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md index 2321b657327..3afd75634d0 100644 --- a/src/frontend/AGENTS.md +++ b/src/frontend/AGENTS.md @@ -36,12 +36,9 @@ distributed mode it talks to remote datanodes via `operator`/`client`. ## Request lifecycles -- **SQL query** (`instance.rs`): `do_query` → `pre_parsing` interceptor → parse → - `post_parsing` interceptor → then per statement: `check_permission` → - `statement_executor.plan` (logical plan) → `query_engine.execute` → for - distributed reads, `region_query.rs` fetches from datanodes → cancellable - `RecordBatch` stream. Interceptors run around parsing, before the per-statement - permission check — preserve that ordering. +- **SQL query** (`instance.rs`): `do_query_inner` handles parsing, interceptors, + permission checks, timeout/cancellation, and delegates planning/execution to + `StatementExecutor`. Distributed scans enter through `region_query.rs`. - **Insert** (`instance/grpc.rs`): `handle_inserts` / `handle_row_inserts` → `check_permission` → `operator`'s `Inserter` (schema validation, optional auto-create, partition routing) → local `RegionServer` (standalone) or RPC to diff --git a/src/metric-engine/AGENTS.md b/src/metric-engine/AGENTS.md index 4471de3faee..d1049478978 100644 --- a/src/metric-engine/AGENTS.md +++ b/src/metric-engine/AGENTS.md @@ -30,7 +30,8 @@ The architecture is documented at the top of `src/metric-engine/src/lib.rs`. | `state` | `src/metric-engine/src/engine/state.rs` | In-memory cache of physical columns and logical column metadata | | `repeated_task` | `src/metric-engine/src/repeated_task.rs` | Periodic metadata-region flush task | | `utils` | `src/metric-engine/src/utils.rs` | `RegionId` conversions (data vs metadata group), manifest encoding | -| `config` | `src/metric-engine/src/config.rs` | `EngineConfig` (metadata flush interval, sparse PK) | +| `config` | `src/metric-engine/src/config.rs` | `EngineConfig` (metadata flush interval) | +| `engine/options` | `src/metric-engine/src/engine/options.rs` | Physical data-region options, including the default sparse primary-key encoding | | `test_util` | `src/metric-engine/src/test_util.rs` | `TestEnv` building the Mito2 + Metric stack | ## Write path @@ -77,8 +78,9 @@ cargo nextest run -p metric-engine writes; operate on logical region ids. - TSID must be stable for the same tag set — it is a hash over sorted tag names + values and may be stored in `__tsid` or encoded into `__primary_key`. -- Metadata is cached (LRU with a TTL); after an alter, stale reads are possible - until invalidation/expiry. +- Metadata reads use an LRU cache. Metadata writes take the per-region cache + write lock and invalidate the cache after the Mito2 write succeeds; preserve + that ordering when adding a metadata mutation path. - Always convert ids via `utils::to_data_region_id` / `to_metadata_region_id`. ## Maintenance contract diff --git a/src/operator/AGENTS.md b/src/operator/AGENTS.md new file mode 100644 index 00000000000..761d8e5080a --- /dev/null +++ b/src/operator/AGENTS.md @@ -0,0 +1,40 @@ +# operator — Agent & Contributor Guide + +Navigation map for `src/operator`. Repo-wide invariants: +[`.agents/architecture-invariants.md`](../../.agents/architecture-invariants.md). + +`operator` connects frontend statements and writes to query, catalog, metasrv, +and datanodes. It owns statement dispatch, DDL task conversion, and region-level +insert/delete routing. + +## Module map + +| Area | Path | Entry point | +| --- | --- | --- | +| Statements | `src/operator/src/statement.rs`, `src/operator/src/statement/` | `StatementExecutor` and per-statement handlers | +| Inserts | `src/operator/src/insert.rs`, `src/operator/src/bulk_insert.rs` | Insert, auto-create/alter, routing, dispatch | +| Deletes | `src/operator/src/delete.rs` | Delete conversion and dispatch | +| Request conversion | `src/operator/src/req_convert.rs`, `src/operator/src/req_convert/` | Table/row/column requests to region requests | +| Region requests | `src/operator/src/region_req_factory.rs` | Region-level DDL/DML request construction | +| Procedures | `src/operator/src/procedure.rs` | Metasrv procedure administration | +| Flow | `src/operator/src/flow.rs`, `src/operator/src/statement/ddl.rs` | Flow requests and internal mode selection | + +## Change coupling + +- New statement kinds usually touch parser/AST handling, frontend permission + and interception, executor dispatch, and sqlness cases. +- DDL task changes require matching `common-meta`, metasrv, recovery, and + compatibility updates. +- Insert/delete conversion changes must cover both request formats, partition + routing, and schema-on-write behavior. +- Flow creation changes must stay aligned with `src/flow/AGENTS.md` and + `common-meta` flow metadata. + +## Testing + +```bash +cargo nextest run -p operator +cargo sqlness bare -t +``` + +Check standalone and distributed routing when changing region dispatch. diff --git a/src/query/AGENTS.md b/src/query/AGENTS.md new file mode 100644 index 00000000000..8dc03a22877 --- /dev/null +++ b/src/query/AGENTS.md @@ -0,0 +1,39 @@ +# query — Agent & Contributor Guide + +Navigation map for `src/query`. Repo-wide invariants: +[`.agents/architecture-invariants.md`](../../.agents/architecture-invariants.md). + +`query` owns SQL/PromQL/log planning, GreptimeDB optimizer rules, DataFusion +execution, and distributed query plans. Protocols live in `servers`; statement +side effects live in `operator`. + +## Module map + +| Area | Path | Entry point | +| --- | --- | --- | +| Logical planning | `src/query/src/planner.rs`, `src/query/src/parser.rs` | `LogicalPlanner`, `QueryStatement` | +| Query engine | `src/query/src/query_engine.rs`, `src/query/src/query_engine/` | Engine traits, factory, state, serialization | +| DataFusion execution | `src/query/src/datafusion.rs`, `src/query/src/datafusion/` | Physical planning and execution | +| Optimizers | `src/query/src/optimizer.rs`, `src/query/src/optimizer/` | GreptimeDB logical/physical rules | +| Distributed plans | `src/query/src/dist_plan.rs`, `src/query/src/dist_plan/` | `MergeScan`, pruning, merge/sort, remote filters | +| Remote reads | `src/query/src/region_query.rs` | `RegionQueryHandler` boundary | +| Query languages | `src/query/src/promql/`, `src/query/src/log_query/` | PromQL and log planners | + +## Change coupling + +- Optimizer or `MergeScan` changes usually touch rule ordering, physical + planning, serialization, and distributed-plan tests together. +- User-visible SQL/PromQL/log behavior needs planner coverage and a sqlness case. +- `RegionQueryHandler` changes require the frontend implementation to move with + it. +- DataFusion dependency changes follow the pinned-fork rule in the repo-wide + invariants. + +## Testing + +```bash +cargo nextest run -p query +cargo sqlness bare -t +``` + +Check both standalone and distributed plans when changing `src/query/src/dist_plan/`. diff --git a/src/servers/AGENTS.md b/src/servers/AGENTS.md new file mode 100644 index 00000000000..2a446033f62 --- /dev/null +++ b/src/servers/AGENTS.md @@ -0,0 +1,37 @@ +# servers — Agent & Contributor Guide + +Navigation map for `src/servers`. Repo-wide invariants: +[`.agents/architecture-invariants.md`](../../.agents/architecture-invariants.md). + +`servers` owns network listeners and protocol translation. Database behavior is +provided through handler traits, mostly implemented by `frontend`. + +## Module map + +| Area | Path | Entry point | +| --- | --- | --- | +| Lifecycle | `src/servers/src/server.rs` | Shared `Server` trait | +| Handler boundary | `src/servers/src/query_handler.rs`, `src/servers/src/query_handler/` | Protocol-independent handler traits | +| gRPC / Flight | `src/servers/src/grpc.rs`, `src/servers/src/grpc/` | gRPC services, auth, cancellation, builders | +| HTTP | `src/servers/src/http.rs`, `src/servers/src/http/` | Routes, middleware, protocol endpoints, output formats | +| SQL protocols | `src/servers/src/mysql.rs`, `src/servers/src/mysql/`, `src/servers/src/postgres.rs`, `src/servers/src/postgres/` | MySQL/PostgreSQL servers and type conversion | +| Observability protocols | `src/servers/src/prom_remote_write/`, `src/servers/src/prometheus.rs`, `src/servers/src/otlp/`, `src/servers/src/otel_arrow.rs` | Prometheus and OpenTelemetry conversion | +| Resources / TLS | `src/servers/src/request_memory_limiter.rs`, `src/servers/src/tls.rs`, `src/servers/src/addrs.rs` | Memory admission, TLS, listener addresses | + +## Change coupling + +- Handler trait changes require matching frontend and mock implementations. +- gRPC wire changes start in `greptime-proto`; update builders and handlers here + after bumping it. +- Auth or `QueryContext` changes must cover HTTP, gRPC, MySQL, and PostgreSQL + entry points. +- New protocols or externally visible routes also require frontend service and + configuration wiring. + +## Testing + +```bash +cargo nextest run -p servers +``` + +Keep protocol translation here and permissions/database behavior in frontend. diff --git a/tests/perf/AGENTS.md b/tests/perf/AGENTS.md index e7c24f16e62..1f13dea58dc 100644 --- a/tests/perf/AGENTS.md +++ b/tests/perf/AGENTS.md @@ -19,9 +19,10 @@ - Keep the direct-SST generator generic. Issue-specific behavior belongs in case files and thresholds, not in Rust generator logic. - Before pushing perf harness changes, run at least: - - `uv run --no-project python -m py_compile .github/scripts/query-regression-run.py .github/scripts/query-regression-summary.py .github/scripts/query-regression-pr-metadata.py tests/perf/test_query_regression_runner_compaction_toctou.py tests/perf/test_query_regression_runner_otlp_trace_load.py` - - `uv run --no-project python tests/perf/test_query_regression_runner_compaction_toctou.py && uv run --no-project python tests/perf/test_query_regression_runner_otlp_trace_load.py` + - the Python tests in the `Test query regression tooling` step of + `.github/workflows/query-regression.yml` - `cargo fmt --all -- --check` - `cargo build -p cmd --bin query_perf_fixture --features dev-tools` + - `cargo build -p cmd --bin query_regression_runner --features dev-tools` - exercise the outer lifecycle script and Rust fixture generator against all built-in cases when the DSL or workflow case selection changes. From 606e5fd6a0c2389321eab6e06ef6abd98aeff683 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Mon, 10 Aug 2026 11:42:53 +0800 Subject: [PATCH 02/92] feat: update opentelemetry family to 0.32 series (#8776) * feat: update openetelemtry family to 0.32 series * chore: resolve warning * fix: update tests --- Cargo.lock | 175 ++++++------------ Cargo.toml | 4 +- src/common/telemetry/Cargo.toml | 6 +- src/common/telemetry/src/tracing_context.rs | 2 +- src/common/telemetry/src/tracing_sampler.rs | 6 +- src/frontend/src/instance/otlp.rs | 1 + .../src/instance/otlp/trace_ingest/tests.rs | 1 + src/mito2/src/gc.rs | 2 +- src/servers/src/otlp/logs.rs | 12 ++ src/servers/src/otlp/trace/attributes.rs | 12 ++ src/servers/src/otlp/trace/span.rs | 1 + src/servers/src/otlp/trace/v1.rs | 7 + src/servers/src/otlp/utils.rs | 11 ++ 13 files changed, 112 insertions(+), 128 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 476133db9e7..841cc035f77 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2262,7 +2262,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" dependencies = [ "lazy_static", - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] @@ -2897,10 +2897,10 @@ dependencies = [ "humantime-serde", "lazy_static", "once_cell", - "opentelemetry 0.30.0", + "opentelemetry", "opentelemetry-otlp", "opentelemetry-semantic-conventions 0.30.0", - "opentelemetry_sdk 0.30.0", + "opentelemetry_sdk", "parking_lot 0.12.4", "prometheus 0.14.0", "serde", @@ -5599,7 +5599,7 @@ dependencies = [ "meta-client", "meta-srv", "num_cpus", - "opentelemetry-proto 0.31.0", + "opentelemetry-proto", "opentelemetry-semantic-conventions 0.31.0", "operator", "otel-arrow-rust", @@ -6662,7 +6662,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core 0.57.0", + "windows-core 0.61.2", ] [[package]] @@ -7831,7 +7831,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667" dependencies = [ "cfg-if", - "windows-targets 0.48.5", + "windows-targets 0.53.5", ] [[package]] @@ -9795,23 +9795,9 @@ checksum = "d05e27ee213611ffe7d6348b942e8f942b37114c00cc03cec254295a4a17852e" [[package]] name = "opentelemetry" -version = "0.30.0" +version = "0.32.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aaf416e4cb72756655126f7dd7bb0af49c674f4c1b9903e80c009e0c37e552e6" -dependencies = [ - "futures-core", - "futures-sink", - "js-sys", - "pin-project-lite", - "thiserror 2.0.17", - "tracing", -] - -[[package]] -name = "opentelemetry" -version = "0.31.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b84bcd6ae87133e903af7ef497404dda70c60d0ea14895fc8a5e6722754fc2a0" +checksum = "b0142c63252a9e054e68a4c61a5778f7b14f576274d593f8ce883d191a099682" dependencies = [ "futures-core", "futures-sink", @@ -9823,61 +9809,48 @@ dependencies = [ [[package]] name = "opentelemetry-http" -version = "0.30.0" +version = "0.32.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "50f6639e842a97dbea8886e3439710ae463120091e2e064518ba8e716e6ac36d" +checksum = "5683015d09e2df236ef005b17f6f196f0d5f6313c4fa43a7b6a53b52776e4331" dependencies = [ "async-trait", "bytes", "http 1.3.1", - "opentelemetry 0.30.0", - "reqwest 0.12.24", + "opentelemetry", + "reqwest 0.13.4", ] [[package]] name = "opentelemetry-otlp" -version = "0.30.0" +version = "0.32.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dbee664a43e07615731afc539ca60c6d9f1a9425e25ca09c57bc36c87c55852b" +checksum = "9966929966d17620d7c316c643ba62631826e10021409357772d5eea84f62c35" dependencies = [ "http 1.3.1", - "opentelemetry 0.30.0", + "opentelemetry", "opentelemetry-http", - "opentelemetry-proto 0.30.0", - "opentelemetry_sdk 0.30.0", - "prost 0.13.5", - "reqwest 0.12.24", + "opentelemetry-proto", + "opentelemetry_sdk", + "prost 0.14.1", + "reqwest 0.13.4", "thiserror 2.0.17", "tokio", - "tonic 0.13.1", - "tracing", + "tonic 0.14.2", + "tonic-types", ] [[package]] name = "opentelemetry-proto" -version = "0.30.0" +version = "0.32.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e046fd7660710fe5a05e8748e70d9058dc15c94ba914e7c4faa7c728f0e8ddc" -dependencies = [ - "opentelemetry 0.30.0", - "opentelemetry_sdk 0.30.0", - "prost 0.13.5", - "tonic 0.13.1", -] - -[[package]] -name = "opentelemetry-proto" -version = "0.31.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7175df06de5eaee9909d4805a3d07e28bb752c34cab57fa9cff549da596b30f" +checksum = "56d658ba1faf63f7b9c492cfbe6e0ec365440a16132d3270c1065f7b33f1b638" dependencies = [ "base64 0.22.1", "const-hex", - "opentelemetry 0.31.0", - "opentelemetry_sdk 0.31.0", + "opentelemetry", + "opentelemetry_sdk", "prost 0.14.1", "serde", - "serde_json", "tonic 0.14.2", "tonic-prost", ] @@ -9896,37 +9869,22 @@ checksum = "e62e29dfe041afb8ed2a6c9737ab57db4907285d999ef8ad3a59092a36bdc846" [[package]] name = "opentelemetry_sdk" -version = "0.30.0" +version = "0.32.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11f644aa9e5e31d11896e024305d7e3c98a88884d9f8919dbf37a9991bc47a4b" +checksum = "9b59f80e1ac4d5ff7a2db8fb6c80badb7f0f3f858211fba08dd9aaec750894f9" dependencies = [ "futures-channel", "futures-executor", "futures-util", - "opentelemetry 0.30.0", + "opentelemetry", "percent-encoding", + "portable-atomic", "rand 0.9.4", - "serde_json", "thiserror 2.0.17", "tokio", "tokio-stream", ] -[[package]] -name = "opentelemetry_sdk" -version = "0.31.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e14ae4f5991976fd48df6d843de219ca6d31b01daaab2dad5af2badeded372bd" -dependencies = [ - "futures-channel", - "futures-executor", - "futures-util", - "opentelemetry 0.31.0", - "percent-encoding", - "rand 0.9.4", - "thiserror 2.0.17", -] - [[package]] name = "operator" version = "1.3.0" @@ -11287,7 +11245,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ "heck 0.5.0", - "itertools 0.11.0", + "itertools 0.14.0", "log", "multimap", "once_cell", @@ -11335,7 +11293,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.11.0", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -11348,7 +11306,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9120690fafc389a67ba3803df527d0ec9cbbc9cc45e4cc20b332996dfb672425" dependencies = [ "anyhow", - "itertools 0.11.0", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -12247,7 +12205,6 @@ checksum = "9d0946410b9f7b082a427e4ef5c8ff541a88b357bc6c637c40db3a68ac70a36f" dependencies = [ "base64 0.22.1", "bytes", - "futures-channel", "futures-core", "futures-util", "h2 0.4.11", @@ -12291,6 +12248,7 @@ checksum = "219c5811de6525e5416c7d5d53bb656d3afdbc6c5af816e0802bcfa42dbdc1c3" dependencies = [ "base64 0.22.1", "bytes", + "futures-channel", "futures-core", "futures-util", "http 1.3.1", @@ -13349,7 +13307,7 @@ dependencies = [ "once_cell", "openmetrics-parser", "opensrv-mysql", - "opentelemetry-proto 0.31.0", + "opentelemetry-proto", "operator", "otel-arrow-rust", "parking_lot 0.12.4", @@ -14857,7 +14815,7 @@ dependencies = [ "mito2", "mysql_async", "object-store", - "opentelemetry-proto 0.31.0", + "opentelemetry-proto", "operator", "otel-arrow-rust", "partition", @@ -15318,32 +15276,6 @@ dependencies = [ "tracing", ] -[[package]] -name = "tonic" -version = "0.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e581ba15a835f4d9ea06c55ab1bd4dce26fc53752c69a04aac00703bfb49ba9" -dependencies = [ - "async-trait", - "base64 0.22.1", - "bytes", - "http 1.3.1", - "http-body 1.0.1", - "http-body-util", - "hyper 1.6.0", - "hyper-timeout 0.5.2", - "hyper-util", - "percent-encoding", - "pin-project", - "prost 0.13.5", - "tokio", - "tokio-stream", - "tower 0.5.2", - "tower-layer", - "tower-service", - "tracing", -] - [[package]] name = "tonic" version = "0.14.2" @@ -15430,6 +15362,17 @@ dependencies = [ "tonic-prost", ] +[[package]] +name = "tonic-types" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a875a902255423d34c1f20838ab374126db8eb41625b7947a1d54113b0b7399" +dependencies = [ + "prost 0.14.1", + "prost-types 0.14.1", + "tonic 0.14.2", +] + [[package]] name = "tonic-web" version = "0.14.5" @@ -15552,9 +15495,9 @@ checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" [[package]] name = "tracing" -version = "0.1.41" +version = "0.1.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "784e0ac535deb450455cbfa28a6f0df145ea1bb7ae51b821cf5e7927fdcfbdd0" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" dependencies = [ "log", "pin-project-lite", @@ -15576,9 +15519,9 @@ dependencies = [ [[package]] name = "tracing-attributes" -version = "0.1.28" +version = "0.1.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "395ae124c09f9e6918a2310af6038fba074bcf474ac352496d5910dd59a2226d" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", @@ -15587,9 +15530,9 @@ dependencies = [ [[package]] name = "tracing-core" -version = "0.1.34" +version = "0.1.36" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9d12581f227e93f094d3af2ae690a574abb8a2b9b7a96e7cfe9647b2b617678" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" dependencies = [ "once_cell", "valuable", @@ -15608,14 +15551,12 @@ dependencies = [ [[package]] name = "tracing-opentelemetry" -version = "0.31.0" +version = "0.33.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ddcf5959f39507d0d04d6413119c04f33b623f4f951ebcbdddddfad2d0623a9c" +checksum = "adbc64cba7137545b8044cb1fe9814f7aacf3c6b5f9b45be8bb5db538befdb26" dependencies = [ "js-sys", - "once_cell", - "opentelemetry 0.30.0", - "opentelemetry_sdk 0.30.0", + "opentelemetry", "smallvec", "tracing", "tracing-core", @@ -15636,9 +15577,9 @@ dependencies = [ [[package]] name = "tracing-subscriber" -version = "0.3.20" +version = "0.3.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2054a14f5307d601f88daf0553e1cbf472acc4f2c51afab632431cdcd72124d5" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" dependencies = [ "matchers", "nu-ansi-term", @@ -16462,7 +16403,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 9582b12bd85..c9b199b3673 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -182,7 +182,7 @@ notify = "8.0" num_cpus = "1.16" object_store_opendal = "0.57" once_cell = "1.18" -opentelemetry-proto = { version = "0.31", features = [ +opentelemetry-proto = { version = "0.32", features = [ "gen-tonic", "metrics", "trace", @@ -260,7 +260,7 @@ tower = "0.5" tower-http = "0.6" tracing = "0.1" tracing-appender = "0.2" -tracing-opentelemetry = "0.31.0" +tracing-opentelemetry = "0.33.0" tracing-subscriber = { version = "0.3", features = ["env-filter", "json", "fmt"] } typetag = "0.2" uuid = { version = "1.17", features = ["serde", "v4", "v7", "fast-rng"] } diff --git a/src/common/telemetry/Cargo.toml b/src/common/telemetry/Cargo.toml index 92c3304d53f..ae25a43c5dc 100644 --- a/src/common/telemetry/Cargo.toml +++ b/src/common/telemetry/Cargo.toml @@ -21,12 +21,12 @@ greptime-proto.workspace = true humantime-serde.workspace = true lazy_static.workspace = true once_cell.workspace = true -opentelemetry = { version = "0.30.0", default-features = false, features = [ +opentelemetry = { version = "0.32.0", default-features = false, features = [ "trace", ] } -opentelemetry-otlp = { version = "0.30.0", features = ["trace", "grpc-tonic", "http-proto"] } +opentelemetry-otlp = { version = "0.32.0", features = ["trace", "grpc-tonic", "http-proto"] } opentelemetry-semantic-conventions = { version = "0.30.0", features = ["semconv_experimental"] } -opentelemetry_sdk = { version = "0.30.0", features = ["rt-tokio", "trace"] } +opentelemetry_sdk = { version = "0.32.1", features = ["rt-tokio", "trace"] } parking_lot.workspace = true prometheus.workspace = true serde.workspace = true diff --git a/src/common/telemetry/src/tracing_context.rs b/src/common/telemetry/src/tracing_context.rs index fd389bccf59..e427b40effd 100644 --- a/src/common/telemetry/src/tracing_context.rs +++ b/src/common/telemetry/src/tracing_context.rs @@ -73,7 +73,7 @@ impl TracingContext { /// Attach the given span as a child of the context. Returns the attached span. pub fn attach(&self, span: tracing::Span) -> tracing::Span { - span.set_parent(self.0.clone()); + let _ = span.set_parent(self.0.clone()); span } diff --git a/src/common/telemetry/src/tracing_sampler.rs b/src/common/telemetry/src/tracing_sampler.rs index ccd029b781b..facc0e5d40f 100644 --- a/src/common/telemetry/src/tracing_sampler.rs +++ b/src/common/telemetry/src/tracing_sampler.rs @@ -15,10 +15,8 @@ use std::collections::HashSet; use opentelemetry::KeyValue; -use opentelemetry::trace::{ - Link, SamplingDecision, SamplingResult, SpanKind, TraceContextExt, TraceId, TraceState, -}; -use opentelemetry_sdk::trace::{Sampler, ShouldSample}; +use opentelemetry::trace::{Link, SpanKind, TraceContextExt, TraceId, TraceState}; +use opentelemetry_sdk::trace::{Sampler, SamplingDecision, SamplingResult, ShouldSample}; use serde::{Deserialize, Serialize}; #[derive(Clone, Debug, Serialize, Deserialize)] diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index dc31416dfe4..c4616ae578d 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -291,6 +291,7 @@ mod tests { value: Some(AnyValue { value: Some(any_value::Value::StringValue("frontend".to_string())), }), + ..Default::default() }], ..Default::default() }), diff --git a/src/frontend/src/instance/otlp/trace_ingest/tests.rs b/src/frontend/src/instance/otlp/trace_ingest/tests.rs index 949ab1535ac..0f18a7194be 100644 --- a/src/frontend/src/instance/otlp/trace_ingest/tests.rs +++ b/src/frontend/src/instance/otlp/trace_ingest/tests.rs @@ -234,6 +234,7 @@ fn trace_attribute(key: &str, value: OtlpValue) -> KeyValue { KeyValue { key: key.to_string(), value: Some(AnyValue { value: Some(value) }), + ..Default::default() } } diff --git a/src/mito2/src/gc.rs b/src/mito2/src/gc.rs index babebabb741..a5474099cbf 100644 --- a/src/mito2/src/gc.rs +++ b/src/mito2/src/gc.rs @@ -410,7 +410,7 @@ impl LocalGcWorker { region_present = region.is_some() ) )] - pub async fn do_region_gc( + pub(crate) async fn do_region_gc( &self, region_id: RegionId, region: Option, diff --git a/src/servers/src/otlp/logs.rs b/src/servers/src/otlp/logs.rs index a956c7977c8..33d4f4eb5e5 100644 --- a/src/servers/src/otlp/logs.rs +++ b/src/servers/src/otlp/logs.rs @@ -24,6 +24,7 @@ use api::v1::{ RowInsertRequest, Rows, SemanticType, Value as GreptimeValue, }; use bytes::Bytes; +use common_telemetry::warn; use common_time::Timestamp; use common_time::timestamp::TimeUnit; use jsonb::{Number as JsonbNumber, Value as JsonbValue}; @@ -1031,6 +1032,16 @@ fn any_value_to_vrl_value(value: any_value::Value) -> VrlValue { VrlValue::Object(key_value_to_map(key_value_list.values)) } any_value::Value::BytesValue(items) => VrlValue::Bytes(Bytes::from(items)), + // `StringValueStrindex` references the Profiling signal's + // `ProfilesDictionary.string_table`, which is unavailable to logs/traces. + // Per the OTLP spec, non-Profiling receivers must treat it as a non-fatal + // issue and process the value as if it were absent. + any_value::Value::StringValueStrindex(_) => { + warn!( + "encountered a profiling-only `StringValueStrindex` value in a non-profiling signal; ignoring" + ); + VrlValue::Null + } } } @@ -1082,6 +1093,7 @@ mod tests { KeyValue { key: key.to_string(), value: Some(AnyValue { value: Some(value) }), + ..Default::default() } } diff --git a/src/servers/src/otlp/trace/attributes.rs b/src/servers/src/otlp/trace/attributes.rs index 08554863432..471832b8920 100644 --- a/src/servers/src/otlp/trace/attributes.rs +++ b/src/servers/src/otlp/trace/attributes.rs @@ -16,6 +16,7 @@ use std::fmt::Display; use opentelemetry_proto::tonic::common::v1::any_value::Value::{ ArrayValue, BoolValue, BytesValue, DoubleValue, IntValue, KvlistValue, StringValue, + StringValueStrindex, }; use opentelemetry_proto::tonic::common::v1::{AnyValue, KeyValue}; use serde::Serialize; @@ -78,6 +79,11 @@ impl Serialize for OtlpAnyValue<'_> { map.end() } BytesValue(v) => zer.serialize_bytes(v), + // `StringValueStrindex` is profiling-signal-only and references the + // Profiling `ProfilesDictionary.string_table`, which is unavailable + // here. Per the OTLP spec, non-Profiling receivers must treat it as + // a non-fatal issue and process the value as if it were absent. + StringValueStrindex(_) => zer.serialize_none(), }, None => zer.serialize_none(), } @@ -293,6 +299,7 @@ mod tests { vec![KeyValue { key: "key1".into(), value: None, + ..Default::default() }], ), ( @@ -300,6 +307,7 @@ mod tests { vec![KeyValue { key: "key1".into(), value: Some(AnyValue { value: None }), + ..Default::default() }], ), ( @@ -309,6 +317,7 @@ mod tests { value: Some(AnyValue { value: Some(Value::StringValue(String::from("val1"))), }), + ..Default::default() }], ), ]; @@ -331,6 +340,7 @@ mod tests { vec![KeyValue { key: "key1".into(), value: None, + ..Default::default() }], ), ( @@ -338,6 +348,7 @@ mod tests { vec![KeyValue { key: "key1".into(), value: Some(AnyValue { value: None }), + ..Default::default() }], ), ( @@ -347,6 +358,7 @@ mod tests { value: Some(AnyValue { value: Some(Value::StringValue(String::from("val1"))), }), + ..Default::default() }], ), ]; diff --git a/src/servers/src/otlp/trace/span.rs b/src/servers/src/otlp/trace/span.rs index 19103240f62..fe8b37fed0a 100644 --- a/src/servers/src/otlp/trace/span.rs +++ b/src/servers/src/otlp/trace/span.rs @@ -319,6 +319,7 @@ mod tests { value: Some(AnyValue { value: Some(any_value::Value::StringValue(value.to_string())), }), + ..Default::default() } } diff --git a/src/servers/src/otlp/trace/v1.rs b/src/servers/src/otlp/trace/v1.rs index eb1612eada6..8ddeac24f33 100644 --- a/src/servers/src/otlp/trace/v1.rs +++ b/src/servers/src/otlp/trace/v1.rs @@ -633,6 +633,12 @@ fn write_attributes_with_schema( row, ); } + // `StringValueStrindex` is profiling-signal-only and references the + // Profiling `ProfilesDictionary.string_table`, which is unavailable to + // traces. Per the OTLP spec, non-Profiling receivers must treat it as a + // non-fatal issue and process the value as if it were absent. Like the + // `None` arm, no field is written for the attribute. + Some(OtlpValue::StringValueStrindex(_)) => {} None => {} } } @@ -656,6 +662,7 @@ mod tests { KeyValue { key: key.to_string(), value: Some(AnyValue { value: Some(value) }), + ..Default::default() } } diff --git a/src/servers/src/otlp/utils.rs b/src/servers/src/otlp/utils.rs index d69f8896d38..5ae5fe6ef50 100644 --- a/src/servers/src/otlp/utils.rs +++ b/src/servers/src/otlp/utils.rs @@ -14,6 +14,7 @@ use api::v1::ColumnDataType; use api::v1::value::ValueData; +use common_telemetry::warn; use jsonb::{Number as JsonbNumber, Value as JsonbValue}; use opentelemetry_proto::tonic::common::v1::{KeyValue, any_value}; @@ -40,6 +41,16 @@ pub fn any_value_to_jsonb(value: any_value::Value) -> JsonbValue<'static> { } any_value::Value::KvlistValue(kv) => key_value_to_jsonb(kv.values), any_value::Value::BytesValue(b) => JsonbValue::String(bytes_to_hex_string(&b).into()), + // `StringValueStrindex` references the Profiling signal's + // `ProfilesDictionary.string_table`, which is unavailable to logs/traces. + // Per the OTLP spec, non-Profiling receivers must treat it as a non-fatal + // issue and process the value as if it were absent. + any_value::Value::StringValueStrindex(_) => { + warn!( + "encountered a profiling-only `StringValueStrindex` value in a non-profiling signal; ignoring" + ); + JsonbValue::Null + } } } From 42c94211a2878c402d98bafd2304596a9d848ac7 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 10 Aug 2026 12:35:20 +0800 Subject: [PATCH 03/92] chore(ci): warn when a member has too many open pull requests (#8813) * chore(ci): warn when a member has too many open pull requests Review capacity is the bottleneck. Add a `pull_request_target` workflow that counts an org member's open pull requests (drafts included) on open/reopen and posts a warning comment when the count exceeds the limit. Advisory only for now: nothing is closed and no check fails. Signed-off-by: Dennis Zhuang * chore(ci): address review on the pr-open-limit workflow - Only match marker comments authored by the Actions bot; a marker pasted by anyone else would otherwise be picked up and fail the edit with 403. - Validate MAX_OPEN_PRS and fall back to 5 on a non-numeric variable. - Drop pull-requests write permission; commenting goes through the issues API. - continue-on-error so a script failure never marks the pull request red. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- .github/scripts/pr-open-limit.js | 117 ++++++++++++++++++++++++++++ .github/workflows/pr-open-limit.yml | 48 ++++++++++++ 2 files changed, 165 insertions(+) create mode 100644 .github/scripts/pr-open-limit.js create mode 100644 .github/workflows/pr-open-limit.yml diff --git a/.github/scripts/pr-open-limit.js b/.github/scripts/pr-open-limit.js new file mode 100644 index 00000000000..1ba908b7450 --- /dev/null +++ b/.github/scripts/pr-open-limit.js @@ -0,0 +1,117 @@ +// Warns when a repository member has too many open pull requests. +// +// Review capacity is the bottleneck, so authors are expected to land or close +// existing work before opening more. Drafts count: they still occupy attention. +// This is advisory only — it never closes a pull request or fails the job. + +import { appendFileSync } from "node:fs"; + +const MARKER = ""; +// GITHUB_TOKEN can only edit comments authored by the Actions bot itself. +const COMMENT_AUTHOR = "github-actions[bot]"; +// Associations that identify an organization member. Outside collaborators and +// community contributors are out of scope. +const MEMBER_ASSOCIATIONS = ["OWNER", "MEMBER"]; +// Keep the comment readable for authors who are far over the limit. +const MAX_LISTED_PRS = 10; + +function summary(text) { + if (process.env.GITHUB_STEP_SUMMARY) { + appendFileSync(process.env.GITHUB_STEP_SUMMARY, `${text}\n`); + } + console.log(text); +} + +function buildComment(author, total, limit, openPrs) { + const listed = openPrs + .slice(0, MAX_LISTED_PRS) + .map((pr) => `- #${pr.number} ${pr.title}${pr.draft ? " _(draft)_" : ""}`); + const hidden = openPrs.length - listed.length; + const list = hidden > 0 ? `${listed.join("\n")}\n- ...and ${hidden} more` : listed.join("\n"); + + return `${MARKER} +> [!WARNING] +> @${author} has **${total}** open pull requests in this repository, over the limit of **${limit}**. + +Review is the scarcest resource here. Please land or close some of these before +pushing this one forward: + +${list} + +This check is advisory for now and blocks nothing.`; +} + +async function upsertComment(octokit, params, body) { + const comments = await octokit.paginate(octokit.issues.listComments, { + ...params, + per_page: 100, + }); + const existing = comments.find( + (comment) => + comment.user?.login === COMMENT_AUTHOR && comment.body?.includes(MARKER) + ); + + if (existing) { + await octokit.issues.updateComment({ + owner: params.owner, + repo: params.repo, + comment_id: existing.id, + body, + }); + } else { + await octokit.issues.createComment({ ...params, body }); + } +} + +(async () => { + const { Octokit } = await import("@octokit/rest"); + + const [owner, repo] = process.env.GITHUB_REPOSITORY.split("/"); + const prNumber = Number(process.env.PR_NUMBER); + const author = process.env.PR_AUTHOR; + const association = process.env.PR_AUTHOR_ASSOCIATION; + const rawLimit = Number(process.env.MAX_OPEN_PRS); + const limit = Number.isInteger(rawLimit) && rawLimit >= 0 ? rawLimit : 5; + + if (author.endsWith("[bot]")) { + summary(`Skipping bot author \`${author}\`.`); + return; + } + + if (!MEMBER_ASSOCIATIONS.includes(association)) { + summary(`Skipping \`${author}\`: association is \`${association}\`, not an org member.`); + return; + } + + const octokit = new Octokit({ auth: process.env.GITHUB_TOKEN }); + + const allOpen = await octokit.paginate(octokit.pulls.list, { + owner, + repo, + state: "open", + sort: "created", + direction: "asc", + per_page: 100, + }); + // The pull request that triggered this run is already open, so exclude it and + // report the total separately. + const others = allOpen.filter( + (pr) => pr.user?.login === author && pr.number !== prNumber + ); + const total = others.length + 1; + + if (others.length < limit) { + summary(`\`${author}\` has ${total} open pull requests (limit ${limit}). OK.`); + return; + } + + summary(`\`${author}\` has ${total} open pull requests, over the limit of ${limit}.`); + await upsertComment( + octokit, + { owner, repo, issue_number: prNumber }, + buildComment(author, total, limit, others) + ); +})().catch((error) => { + console.error(error); + process.exit(1); +}); diff --git a/.github/workflows/pr-open-limit.yml b/.github/workflows/pr-open-limit.yml new file mode 100644 index 00000000000..d22929911a8 --- /dev/null +++ b/.github/workflows/pr-open-limit.yml @@ -0,0 +1,48 @@ +name: PR Open Limit + +on: + pull_request_target: + types: + - opened + - reopened + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number }} + cancel-in-progress: true + +# The warning comment goes through the issues API, hence `issues: write`. +permissions: + contents: read + pull-requests: read + issues: write + +jobs: + check-open-limit: + name: Check Open Pull Request Limit + runs-on: ubuntu-latest + timeout-minutes: 10 + # Advisory: a script failure must not put a red cross on the pull request. + continue-on-error: true + if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Setup Node.js + uses: actions/setup-node@v4 + with: + node-version: '20' + + - name: Install dependencies + working-directory: .github/scripts + run: npm ci + + - name: Check open pull request limit + working-directory: .github/scripts + run: node pr-open-limit.js + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + PR_NUMBER: ${{ github.event.pull_request.number }} + PR_AUTHOR: ${{ github.event.pull_request.user.login }} + PR_AUTHOR_ASSOCIATION: ${{ github.event.pull_request.author_association }} + MAX_OPEN_PRS: ${{ vars.MAX_OPEN_PRS_PER_AUTHOR || 5 }} From 0a2b10e476635ea9ffd2ac7fd21e45d48a1ecb24 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 10 Aug 2026 13:30:13 +0800 Subject: [PATCH 04/92] fix(mito2): avoid region worker panic when building a WAL entry fails (#8810) When `add_wal_entry` fails for a region, the worker only sets the error on the write context, which stays in `region_ctxs`. The region's entry is not in the batch, so a successful `write_to_wal` returns no last entry id for it and the success branch panics on `response.last_entry_ids.get(region_id).unwrap()`, killing the region worker. When the failed region is the only one in the batch, the batch is empty and `append_batch` always returns an empty response, so the panic is guaranteed. No in-tree log store can fail to build an entry at runtime today (the provider/log store combination is validated when the region opens), so this is a latent panic rather than a reachable crash. Skip contexts already marked as failed when updating next entry ids; their waiters are already notified with the error. Extract the WAL phase of `handle_write_requests` into `write_wal` and cover the failure paths with unit tests. Signed-off-by: jeremyhi --- src/mito2/src/region_write_ctx.rs | 11 + src/mito2/src/worker/handle_write.rs | 287 ++++++++++++++++++++++++--- 2 files changed, 269 insertions(+), 29 deletions(-) diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index c6abdb0872c..b71e3e5ab17 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -218,11 +218,22 @@ impl RegionWriteCtx { self.failed = true; } + /// Returns whether the write operation is already marked as failed. + pub(crate) fn is_failed(&self) -> bool { + self.failed + } + /// Updates next entry id. pub(crate) fn set_next_entry_id(&mut self, next_entry_id: EntryId) { self.next_entry_id = next_entry_id } + /// Returns the next entry id to write. + #[cfg(test)] + pub(crate) fn next_entry_id(&self) -> EntryId { + self.next_entry_id + } + /// Consumes mutations and writes them into mutable memtable. pub(crate) async fn write_memtable(&mut self) { debug_assert_eq!(self.notifiers.len(), self.wal_entry.mutations.len()); diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index c52572df114..3be40553145 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -35,6 +35,7 @@ use crate::metrics::{ use crate::region::{RegionLeaderState, RegionRoleState}; use crate::region_write_ctx::RegionWriteCtx; use crate::request::{SenderBulkRequest, SenderWriteRequest, WriteRequest}; +use crate::wal::Wal; use crate::worker::RegionWorkerLoop; impl RegionWorkerLoop { @@ -110,35 +111,9 @@ impl RegionWorkerLoop { let _timer = WRITE_STAGE_ELAPSED .with_label_values(&["write_wal"]) .start_timer(); - let mut wal_writer = self.wal.writer(); - for region_ctx in region_ctxs.values_mut() { - if region_ctx.skip_wal() { - continue; - } - if let Err(e) = region_ctx.add_wal_entry(&mut wal_writer).map_err(Arc::new) { - region_ctx.set_error(e); - } - } - match wal_writer.write_to_wal().await.map_err(Arc::new) { - Ok(response) => { - for (region_id, region_ctx) in region_ctxs.iter_mut() { - if region_ctx.skip_wal() { - continue; - } - - // Safety: the log store implementation ensures that either the `write_to_wal` fails and no - // response is returned or the last entry ids for each region do exist. - let last_entry_id = response.last_entry_ids.get(region_id).unwrap(); - region_ctx.set_next_entry_id(last_entry_id + 1); - } - } - Err(e) => { - // Failed to write wal. - for mut region_ctx in region_ctxs.into_values() { - region_ctx.set_error(e.clone()); - } - return; - } + if !write_wal(&self.wal, &mut region_ctxs).await { + // Failed to write to the WAL, all waiters are notified with the error. + return; } } @@ -625,6 +600,55 @@ impl RegionWorkerLoop { } } +/// Writes WAL entries of all region contexts to the WAL in one batch and updates +/// the next entry id of each region on success. +/// +/// Returns `false` if the batch fails to be written to the WAL. In this case all +/// contexts are consumed and their waiters are notified with the error, so the +/// caller should skip the memtable phase. +async fn write_wal( + wal: &Wal, + region_ctxs: &mut HashMap, +) -> bool { + let mut wal_writer = wal.writer(); + for region_ctx in region_ctxs.values_mut() { + if region_ctx.skip_wal() { + continue; + } + if let Err(e) = region_ctx.add_wal_entry(&mut wal_writer).map_err(Arc::new) { + region_ctx.set_error(e); + } + } + match wal_writer.write_to_wal().await.map_err(Arc::new) { + Ok(response) => { + for (region_id, region_ctx) in region_ctxs.iter_mut() { + if region_ctx.skip_wal() { + continue; + } + // The entry of a failed region (e.g. failed to build its WAL entry) is + // not in the batch so the response has no last entry id for it. Its + // waiters are already notified with the error. + if region_ctx.is_failed() { + continue; + } + + // Safety: the log store implementation ensures that either the `write_to_wal` fails and no + // response is returned or the last entry ids for each region in the batch do exist. + let last_entry_id = response.last_entry_ids.get(region_id).unwrap(); + region_ctx.set_next_entry_id(last_entry_id + 1); + } + true + } + Err(e) => { + // Failed to write wal. + for (_, mut region_ctx) in region_ctxs.drain() { + region_ctx.set_error(e.clone()); + } + false + } + } +} + /// Send rejected error to all `write_requests`. fn reject_write_requests( write_requests: &mut Vec, @@ -707,3 +731,208 @@ fn check_partition_expr_version( } Ok(()) } + +#[cfg(test)] +mod tests { + use api::v1::{Row, Rows}; + use futures::stream; + use log_store::error::{ + Error as LogStoreError, IllegalStateSnafu, InvalidProviderSnafu, Result as LogStoreResult, + }; + use store_api::logstore::entry::{Entry, NaiveEntry}; + use store_api::logstore::provider::Provider; + use store_api::logstore::{AppendBatchResponse, EntryId, SendableEntryStream, WalIndex}; + use store_api::region_request::AffectedRows; + use tokio::sync::oneshot; + + use super::*; + use crate::request::OptionOutputTx; + use crate::test_util::version_util::VersionControlBuilder; + + /// A log store that fails to build entries for `failing_region` and fails the + /// whole batch when `fail_append` is true. + #[derive(Debug, Default)] + struct MockLogStore { + failing_region: Option, + fail_append: bool, + } + + #[async_trait::async_trait] + impl LogStore for MockLogStore { + type Error = LogStoreError; + + async fn stop(&self) -> LogStoreResult<()> { + Ok(()) + } + + async fn append_batch(&self, entries: Vec) -> LogStoreResult { + if self.fail_append { + return IllegalStateSnafu {}.fail(); + } + let mut last_entry_ids = HashMap::new(); + for entry in &entries { + let last_entry_id = last_entry_ids.entry(entry.region_id()).or_insert(0); + *last_entry_id = entry.entry_id().max(*last_entry_id); + } + Ok(AppendBatchResponse { last_entry_ids }) + } + + async fn read( + &self, + _provider: &Provider, + _entry_id: EntryId, + _index: Option, + ) -> LogStoreResult> { + Ok(Box::pin(stream::empty())) + } + + async fn create_namespace(&self, _ns: &Provider) -> LogStoreResult<()> { + Ok(()) + } + + async fn delete_namespace(&self, _ns: &Provider) -> LogStoreResult<()> { + Ok(()) + } + + async fn list_namespaces(&self) -> LogStoreResult> { + Ok(vec![]) + } + + async fn obsolete( + &self, + _provider: &Provider, + _region_id: RegionId, + _entry_id: EntryId, + ) -> LogStoreResult<()> { + Ok(()) + } + + async fn obsolete_all( + &self, + _provider: &Provider, + _region_id: RegionId, + ) -> LogStoreResult<()> { + Ok(()) + } + + fn entry( + &self, + data: Vec, + entry_id: EntryId, + region_id: RegionId, + provider: &Provider, + ) -> LogStoreResult { + if self.failing_region == Some(region_id) { + return InvalidProviderSnafu { + expected: "raft_engine", + actual: "mock", + } + .fail(); + } + Ok(Entry::Naive(NaiveEntry { + provider: provider.clone(), + region_id, + entry_id, + data, + })) + } + + fn latest_entry_id(&self, _provider: &Provider) -> LogStoreResult { + Ok(0) + } + } + + /// Creates a write context for `region_id` with one pending mutation of one row. + fn new_region_ctx( + region_id: RegionId, + ) -> (RegionWriteCtx, oneshot::Receiver>) { + let version_control = Arc::new(VersionControlBuilder::new().build()); + let mut ctx = RegionWriteCtx::new( + region_id, + &version_control, + Provider::raft_engine_provider(region_id.as_u64()), + None, + ); + let (tx, rx) = oneshot::channel(); + ctx.push_mutation( + OpType::Put as i32, + Some(Rows { + schema: vec![], + rows: vec![Row { values: vec![] }], + }), + None, + OptionOutputTx::from(tx), + None, + ); + (ctx, rx) + } + + #[tokio::test] + async fn test_write_wal_skips_region_failed_to_build_entry() { + let failing_region = RegionId::new(1, 1); + let ok_region = RegionId::new(1, 2); + let wal = Wal::new(Arc::new(MockLogStore { + failing_region: Some(failing_region), + ..Default::default() + })); + + let mut region_ctxs = HashMap::new(); + let (ctx, failing_rx) = new_region_ctx(failing_region); + region_ctxs.insert(failing_region, ctx); + let (ctx, ok_rx) = new_region_ctx(ok_region); + region_ctxs.insert(ok_region, ctx); + let entry_id = region_ctxs[&ok_region].next_entry_id(); + + // The failed region must not fail the batch or panic the worker. + assert!(write_wal(&wal, &mut region_ctxs).await); + + assert!(region_ctxs[&failing_region].is_failed()); + assert!(!region_ctxs[&ok_region].is_failed()); + assert_eq!(entry_id + 1, region_ctxs[&ok_region].next_entry_id()); + + // Waiters of the failed region get the error while others get the result. + drop(region_ctxs); + assert!(failing_rx.await.unwrap().is_err()); + assert_eq!(1, ok_rx.await.unwrap().unwrap()); + } + + #[tokio::test] + async fn test_write_wal_all_regions_failed_to_build_entries() { + let failing_region = RegionId::new(1, 1); + let wal = Wal::new(Arc::new(MockLogStore { + failing_region: Some(failing_region), + ..Default::default() + })); + + let mut region_ctxs = HashMap::new(); + let (ctx, rx) = new_region_ctx(failing_region); + region_ctxs.insert(failing_region, ctx); + + // Writing an empty batch to the WAL succeeds, the failed region must not panic + // the worker. + assert!(write_wal(&wal, &mut region_ctxs).await); + + assert!(region_ctxs[&failing_region].is_failed()); + drop(region_ctxs); + assert!(rx.await.unwrap().is_err()); + } + + #[tokio::test] + async fn test_write_wal_append_batch_failure() { + let region_id = RegionId::new(1, 1); + let wal = Wal::new(Arc::new(MockLogStore { + fail_append: true, + ..Default::default() + })); + + let mut region_ctxs = HashMap::new(); + let (ctx, rx) = new_region_ctx(region_id); + region_ctxs.insert(region_id, ctx); + + assert!(!write_wal(&wal, &mut region_ctxs).await); + + // All contexts are consumed and waiters are notified with the error. + assert!(region_ctxs.is_empty()); + assert!(rx.await.unwrap().is_err()); + } +} From d4af650ec01c335b6ff392d6c78018ae73c8883b Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:18:16 +0800 Subject: [PATCH 05/92] perf: reduce cold workspace compile time (#8801) * refactor: remove datanode and meta-srv dep from frontend Signed-off-by: shuiyisong * refactor: use on-device protoc if possible Signed-off-by: shuiyisong * refactor: remove unused dep Signed-off-by: shuiyisong * fix: CR issue Signed-off-by: shuiyisong * fix: version and docs Signed-off-by: shuiyisong * chore: update logs Co-authored-by: fys <40801205+fengys1996@users.noreply.github.com> --------- Signed-off-by: shuiyisong Co-authored-by: fys <40801205+fengys1996@users.noreply.github.com> --- .agents/architecture-invariants.md | 4 +- Cargo.lock | 138 +++++++++++++----- Cargo.toml | 2 + src/cmd/src/standalone.rs | 6 +- src/common/stat/Cargo.toml | 4 +- src/common/test-util/Cargo.toml | 1 - src/common/test-util/src/recordbatch.rs | 28 ---- src/file-engine/Cargo.toml | 1 - src/frontend/AGENTS.md | 12 +- src/frontend/Cargo.toml | 3 - src/frontend/src/error.rs | 14 -- src/frontend/src/frontend.rs | 6 +- src/frontend/src/instance.rs | 2 - src/index/Cargo.toml | 4 +- src/log-store/Cargo.toml | 5 +- src/log-store/build.rs | 95 +++++++++++- src/standalone/Cargo.toml | 2 + .../src/datanode_manager.rs} | 4 +- src/standalone/src/error.rs | 21 ++- src/standalone/src/lib.rs | 2 + tests-integration/src/standalone.rs | 4 +- 21 files changed, 244 insertions(+), 114 deletions(-) rename src/{frontend/src/instance/standalone.rs => standalone/src/datanode_manager.rs} (97%) diff --git a/.agents/architecture-invariants.md b/.agents/architecture-invariants.md index 78c78f3a976..f290e1ce386 100644 --- a/.agents/architecture-invariants.md +++ b/.agents/architecture-invariants.md @@ -46,8 +46,8 @@ The workspace is layered; dependencies point downward only. `file-engine`) implement it; `datanode` drives engines **through the trait**, not through engine internals. - `frontend` reaches storage through `operator` / `query` / `catalog`, not by - depending on `datanode` internals. Standalone mode is the one bridge, via a - `RegionServer` wrapper (`src/frontend/src/instance/standalone.rs`). + depending on `datanode` internals. Standalone mode is the one bridge, via the + `RegionServer` adapter in `src/standalone/src/datanode_manager.rs`. - Do not introduce circular dependencies. New deps go through `[workspace.dependencies]` in the root `Cargo.toml`, not per-crate version literals. diff --git a/Cargo.lock b/Cargo.lock index 841cc035f77..4ba0a7f7e27 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -989,15 +989,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" -[[package]] -name = "autotools" -version = "0.2.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef941527c41b0fc0dd48511a8154cd5fc7e29200a0ff8b7203c5d777dbc795cf" -dependencies = [ - "cc", -] - [[package]] name = "aws-lc-rs" version = "1.17.0" @@ -2917,7 +2908,6 @@ dependencies = [ name = "common-test-util" version = "1.3.0" dependencies = [ - "client", "common-grpc", "common-query", "common-recordbatch", @@ -5274,7 +5264,6 @@ dependencies = [ "common-datasource", "common-error", "common-macro", - "common-procedure", "common-query", "common-recordbatch", "common-telemetry", @@ -5587,7 +5576,6 @@ dependencies = [ "dashmap", "datafusion", "datafusion-expr", - "datanode", "datatypes", "futures", "hostname 0.4.1", @@ -5597,7 +5585,6 @@ dependencies = [ "lazy_static", "log-query", "meta-client", - "meta-srv", "num_cpus", "opentelemetry-proto", "opentelemetry-semantic-conventions 0.31.0", @@ -8013,7 +8000,8 @@ dependencies = [ "pin-project", "prometheus 0.14.0", "protobuf 2.28.0", - "protobuf-build", + "protoc-bin-vendored", + "protoc-rust", "raft-engine", "rand 0.9.4", "rskafka", @@ -11379,19 +11367,6 @@ dependencies = [ "thiserror 1.0.69", ] -[[package]] -name = "protobuf-build" -version = "0.15.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c852d9625b912c3e50480cdc701f60f49890b5d7ad46198dd583600f15e7c6ec" -dependencies = [ - "bitflags 1.3.2", - "protobuf 2.28.0", - "protobuf-codegen", - "protobuf-src", - "regex", -] - [[package]] name = "protobuf-codegen" version = "2.28.0" @@ -11401,15 +11376,6 @@ dependencies = [ "protobuf 2.28.0", ] -[[package]] -name = "protobuf-src" -version = "1.1.0+21.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7ac8852baeb3cc6fb83b93646fb93c0ffe5d14bf138c945ceb4b9948ee0e3c1" -dependencies = [ - "autotools", -] - [[package]] name = "protobuf-support" version = "3.7.2" @@ -11419,6 +11385,92 @@ dependencies = [ "thiserror 1.0.69", ] +[[package]] +name = "protoc" +version = "2.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a0218039c514f9e14a5060742ecd50427f8ac4f85a6dc58f2ddb806e318c55ee" +dependencies = [ + "log", + "which", +] + +[[package]] +name = "protoc-bin-vendored" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1c381df33c98266b5f08186583660090a4ffa0889e76c7e9a5e175f645a67fa" +dependencies = [ + "protoc-bin-vendored-linux-aarch_64", + "protoc-bin-vendored-linux-ppcle_64", + "protoc-bin-vendored-linux-s390_64", + "protoc-bin-vendored-linux-x86_32", + "protoc-bin-vendored-linux-x86_64", + "protoc-bin-vendored-macos-aarch_64", + "protoc-bin-vendored-macos-x86_64", + "protoc-bin-vendored-win32", +] + +[[package]] +name = "protoc-bin-vendored-linux-aarch_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c350df4d49b5b9e3ca79f7e646fde2377b199e13cfa87320308397e1f37e1a4c" + +[[package]] +name = "protoc-bin-vendored-linux-ppcle_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a55a63e6c7244f19b5c6393f025017eb5d793fd5467823a099740a7a4222440c" + +[[package]] +name = "protoc-bin-vendored-linux-s390_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1dba5565db4288e935d5330a07c264a4ee8e4a5b4a4e6f4e83fad824cc32f3b0" + +[[package]] +name = "protoc-bin-vendored-linux-x86_32" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8854774b24ee28b7868cd71dccaae8e02a2365e67a4a87a6cd11ee6cdbdf9cf5" + +[[package]] +name = "protoc-bin-vendored-linux-x86_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b38b07546580df720fa464ce124c4b03630a6fb83e05c336fea2a241df7e5d78" + +[[package]] +name = "protoc-bin-vendored-macos-aarch_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89278a9926ce312e51f1d999fee8825d324d603213344a9a706daa009f1d8092" + +[[package]] +name = "protoc-bin-vendored-macos-x86_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81745feda7ccfb9471d7a4de888f0652e806d5795b61480605d4943176299756" + +[[package]] +name = "protoc-bin-vendored-win32" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95067976aca6421a523e491fce939a3e65249bac4b977adee0ee9771568e8aa3" + +[[package]] +name = "protoc-rust" +version = "2.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22f8a182bb17c485f20bdc4274a8c39000a61024cfe461c799b50fec77267838" +dependencies = [ + "protobuf 2.28.0", + "protobuf-codegen", + "protoc", + "tempfile", +] + [[package]] name = "psl" version = "2.1.123" @@ -14066,6 +14118,7 @@ dependencies = [ name = "standalone" version = "1.3.0" dependencies = [ + "api", "async-trait", "catalog", "client", @@ -14079,6 +14132,7 @@ dependencies = [ "common-options", "common-procedure", "common-query", + "common-recordbatch", "common-stat", "common-telemetry", "common-time", @@ -16360,6 +16414,18 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "which" +version = "4.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7" +dependencies = [ + "either", + "home", + "once_cell", + "rustix 0.38.44", +] + [[package]] name = "whoami" version = "1.6.0" diff --git a/Cargo.toml b/Cargo.toml index c9b199b3673..502e652ffc8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -202,6 +202,8 @@ prometheus = { version = "0.14", features = ["process"] } promql-parser = { version = "0.10.0", features = ["ser"] } prost = { version = "0.14", features = ["no-recursion-limit"] } prost-types = "0.14" +protoc-bin-vendored = "3.2" +protoc-rust = "2.28" raft-engine = { version = "0.4.1", default-features = false } rand = "0.9" ratelimit = "0.10" diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index aec373612a4..9f8155b9cc2 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -60,7 +60,6 @@ use flow::{ FrontendInvoker, GrpcQueryHandlerWithBoxedError, }; use frontend::frontend::Frontend; -use frontend::instance::StandaloneDatanodeManager; use frontend::instance::builder::FrontendBuilder; use frontend::server::Services; use meta_srv::metasrv::{FLOW_ID_SEQ, TABLE_ID_SEQ}; @@ -72,7 +71,10 @@ use plugins::standalone::context::DdlManagerConfigureContext; use servers::tls::{TlsMode, TlsOption, merge_tls_option}; use snafu::{OptionExt, ResultExt}; use standalone::options::StandaloneOptions; -use standalone::{StandaloneInformationExtension, StandaloneRepartitionProcedureFactory}; +use standalone::{ + StandaloneDatanodeManager, StandaloneInformationExtension, + StandaloneRepartitionProcedureFactory, +}; use tracing_appender::non_blocking::WorkerGuard; use crate::error::{OtherSnafu, Result, StartFlownodeSnafu}; diff --git a/src/common/stat/Cargo.toml b/src/common/stat/Cargo.toml index d0e8b5448ff..2a99107c909 100644 --- a/src/common/stat/Cargo.toml +++ b/src/common/stat/Cargo.toml @@ -9,11 +9,13 @@ common-base.workspace = true common-runtime.workspace = true common-telemetry.workspace = true lazy_static.workspace = true -nix.workspace = true num_cpus.workspace = true prometheus.workspace = true sysinfo.workspace = true tokio.workspace = true +[target.'cfg(target_os = "linux")'.dependencies] +nix.workspace = true + [lints] workspace = true diff --git a/src/common/test-util/Cargo.toml b/src/common/test-util/Cargo.toml index 18dc70e100b..46b374746ab 100644 --- a/src/common/test-util/Cargo.toml +++ b/src/common/test-util/Cargo.toml @@ -8,7 +8,6 @@ license.workspace = true workspace = true [dependencies] -client = { workspace = true, features = ["testing"] } common-grpc.workspace = true common-query.workspace = true common-recordbatch.workspace = true diff --git a/src/common/test-util/src/recordbatch.rs b/src/common/test-util/src/recordbatch.rs index aa68f793566..15e754ca7a6 100644 --- a/src/common/test-util/src/recordbatch.rs +++ b/src/common/test-util/src/recordbatch.rs @@ -12,15 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. -use client::Database; use common_query::OutputData; use common_recordbatch::util; -pub enum ExpectedOutput<'a> { - AffectedRows(usize), - QueryResult(&'a str), -} - pub async fn check_output_stream(output: OutputData, expected: &str) { let recordbatches = match output { OutputData::Stream(stream) => util::collect_batches(stream).await.unwrap(), @@ -30,25 +24,3 @@ pub async fn check_output_stream(output: OutputData, expected: &str) { let pretty_print = recordbatches.pretty_print().unwrap(); assert_eq!(pretty_print, expected.trim(), "actual: \n{}", pretty_print); } - -pub async fn execute_and_check_output(db: &Database, sql: &str, expected: ExpectedOutput<'_>) { - let output = db.sql(sql).await.unwrap(); - let output = output.data; - - match (&output, expected) { - (OutputData::AffectedRows(x), ExpectedOutput::AffectedRows(y)) => { - assert_eq!( - *x, y, - r#" -expected: {y} -actual: {x} -"# - ) - } - (OutputData::RecordBatches(_), ExpectedOutput::QueryResult(x)) - | (OutputData::Stream(_), ExpectedOutput::QueryResult(x)) => { - check_output_stream(output, x).await - } - _ => panic!(), - } -} diff --git a/src/file-engine/Cargo.toml b/src/file-engine/Cargo.toml index 9d031cb279f..ec904addf1e 100644 --- a/src/file-engine/Cargo.toml +++ b/src/file-engine/Cargo.toml @@ -18,7 +18,6 @@ common-catalog.workspace = true common-datasource.workspace = true common-error.workspace = true common-macro.workspace = true -common-procedure.workspace = true common-query.workspace = true common-recordbatch.workspace = true common-telemetry.workspace = true diff --git a/src/frontend/AGENTS.md b/src/frontend/AGENTS.md index 3afd75634d0..b5c0871fc11 100644 --- a/src/frontend/AGENTS.md +++ b/src/frontend/AGENTS.md @@ -16,8 +16,9 @@ reads go to the query engine (`query` crate), writes go to the inserter/deleter Boundary with `servers`: the `servers` crate implements the wire protocols and network I/O; `frontend` provides the business logic by implementing handler traits (`SqlQueryHandler`, `GrpcQueryHandler`, `InfluxdbLineProtocolHandler`, -...). In standalone mode the frontend embeds a datanode `RegionServer`; in -distributed mode it talks to remote datanodes via `operator`/`client`. +...). In standalone mode the `standalone` crate bridges the frontend to an +embedded datanode `RegionServer`; in distributed mode the frontend talks to +remote datanodes via `operator`/`client`. ## Module map @@ -26,7 +27,6 @@ distributed mode it talks to remote datanodes via `operator`/`client`. | `instance` | `src/frontend/src/instance.rs` | `Instance`: the core handler; implements `SqlQueryHandler`, `PrometheusHandler`, etc. | | `instance/builder` | `src/frontend/src/instance/builder.rs` | `FrontendBuilder` assembles `Instance` from its dependencies | | `instance/grpc` | `src/frontend/src/instance/grpc.rs` | `GrpcQueryHandler`: insert/delete/query/promql over gRPC | -| `instance/standalone` | `src/frontend/src/instance/standalone.rs` | Calls the local `RegionServer` instead of RPC | | `instance/region_query` | `src/frontend/src/instance/region_query.rs` | Routes distributed region reads to datanodes | | `instance/*` | `src/frontend/src/instance/` | Per-protocol handlers (`influxdb.rs`, `promql.rs`, `otlp/`, `jaeger.rs`, `logs.rs`, `prom_store.rs`, ...) | | `frontend` | `src/frontend/src/frontend.rs` | `Frontend` lifecycle wrapper (`FrontendOptions`, start/shutdown) | @@ -70,9 +70,9 @@ cargo nextest run -p frontend - Keep the frontend/servers split straight: wire format and network live in `servers`; permissions, planning, and routing live here. -- Standalone vs distributed diverge in datanode access (local `RegionServer` vs - `NodeClients` RPC), MetaClient usage, and whether heartbeat matters. In - standalone, the cache invalidator is a no-op. +- Standalone vs distributed diverge in datanode access (the `standalone` crate's + local `RegionServer` adapter vs `NodeClients` RPC), MetaClient usage, and + whether heartbeat matters. In standalone, the cache invalidator is a no-op. ## Maintenance contract diff --git a/src/frontend/Cargo.toml b/src/frontend/Cargo.toml index a5a589e2ac7..29ca224a55a 100644 --- a/src/frontend/Cargo.toml +++ b/src/frontend/Cargo.toml @@ -52,7 +52,6 @@ common-version.workspace = true dashmap.workspace = true datafusion.workspace = true datafusion-expr.workspace = true -datanode.workspace = true datatypes.workspace = true futures.workspace = true hostname.workspace = true @@ -91,11 +90,9 @@ tracing.workspace = true [dev-dependencies] catalog = { workspace = true, features = ["testing"] } common-test-util.workspace = true -datanode.workspace = true datatypes.workspace = true futures.workspace = true hyper-util = { workspace = true, features = ["tokio"] } -meta-srv.workspace = true reqwest.workspace = true serde_json.workspace = true strfmt = "0.2" diff --git a/src/frontend/src/error.rs b/src/frontend/src/error.rs index 9a70b716619..c0f9e2fe6bd 100644 --- a/src/frontend/src/error.rs +++ b/src/frontend/src/error.rs @@ -203,13 +203,6 @@ pub enum Error { source: query::error::Error, }, - #[snafu(display("Operation to region server failed"))] - InvokeRegionServer { - #[snafu(implicit)] - location: Location, - source: servers::error::Error, - }, - #[snafu(display("Not supported: {}", feat))] NotSupported { feat: String }, @@ -296,9 +289,6 @@ pub enum Error { location: Location, }, - #[snafu(display("Invalid region request, reason: {}", reason))] - InvalidRegionRequest { reason: String }, - #[snafu(display("Table operation error"))] TableOperation { source: operator::error::Error, @@ -426,8 +416,6 @@ impl ErrorExt for Error { Error::CacheRequired { .. } => StatusCode::Internal, - Error::InvalidRegionRequest { .. } => StatusCode::IllegalState, - Error::TableNotFound { .. } => StatusCode::TableNotFound, Error::Catalog { source, .. } => source.status_code(), @@ -438,7 +426,6 @@ impl ErrorExt for Error { | Error::ReadTable { source, .. } | Error::ExecLogicalPlan { source, .. } => source.status_code(), - Error::InvokeRegionServer { source, .. } => source.status_code(), Error::External { source, .. } | Error::InitPlugin { source, .. } => { source.status_code() } @@ -472,7 +459,6 @@ impl ErrorExt for Error { Error::StartServer { source, .. } | Error::ShutdownServer { source, .. } - | Error::InvokeRegionServer { source, .. } | Error::ExecutePromql { source, .. } | Error::PromStoreRemoteQueryPlan { source, .. } | Error::PrometheusMetricNamesQueryPlan { source, .. } => source.retry_hint(), diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index 30fcd394922..7a09b40e736 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -165,6 +165,7 @@ impl Frontend { mod tests { use std::any::Any; use std::net::SocketAddr; + use std::pin::Pin; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::time::Duration; @@ -186,9 +187,9 @@ mod tests { use common_meta::heartbeat::handler::suspend::SuspendHandler; use common_meta::instruction::Instruction; use common_stat::ResourceStatImpl; + use futures::Stream; use meta_client::MetaClientRef; use meta_client::client::MetaClientBuilder; - use meta_srv::service::GrpcStream; use servers::grpc::{FlightCompression, GRPC_SERVER}; use servers::http::HTTP_SERVER; use servers::http::result::greptime_result_v1::GreptimedbV1Response; @@ -206,6 +207,9 @@ mod tests { use crate::instance::builder::FrontendBuilder; use crate::server::Services; + type GrpcStream = + Pin> + Send + Sync + 'static>>; + #[test] fn test_toml() { let opts = FrontendOptions::default(); diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 23090ef932b..bdf4d66c7e9 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -25,7 +25,6 @@ mod otlp; pub mod prom_store; mod promql; mod region_query; -pub mod standalone; use std::collections::HashSet; use std::pin::Pin; @@ -100,7 +99,6 @@ use sql::statements::statement::Statement; use sql::statements::tql::Tql; use sql::util::{extract_tables_from_prom_expr_checked, extract_tables_from_statement_checked}; use sqlparser::ast::{AnalyzeFormat, ObjectName}; -pub use standalone::StandaloneDatanodeManager; use table::requests::{OTLP_METRIC_COMPAT_KEY, OTLP_METRIC_COMPAT_PROM}; use tracing::Span; diff --git a/src/index/Cargo.toml b/src/index/Cargo.toml index 167c1c0df14..b8518112610 100644 --- a/src/index/Cargo.toml +++ b/src/index/Cargo.toml @@ -8,7 +8,7 @@ license.workspace = true workspace = true [features] -vector_index = ["dep:usearch"] +vector_index = ["dep:nalgebra", "dep:usearch"] [dependencies] async-trait.workspace = true @@ -29,7 +29,7 @@ itertools.workspace = true jieba-rs = "0.10" lazy_static.workspace = true mockall.workspace = true -nalgebra.workspace = true +nalgebra = { workspace = true, optional = true } pin-project.workspace = true prost.workspace = true puffin.workspace = true diff --git a/src/log-store/Cargo.toml b/src/log-store/Cargo.toml index b98657f5609..c8372c578bd 100644 --- a/src/log-store/Cargo.toml +++ b/src/log-store/Cargo.toml @@ -5,9 +5,8 @@ edition.workspace = true license.workspace = true [build-dependencies] -protobuf-build = { version = "0.15", default-features = false, features = [ - "protobuf-codec", -] } +protoc-bin-vendored.workspace = true +protoc-rust.workspace = true [lints] workspace = true diff --git a/src/log-store/build.rs b/src/log-store/build.rs index c6c15cdfe2c..a8b6f791487 100644 --- a/src/log-store/build.rs +++ b/src/log-store/build.rs @@ -12,13 +12,94 @@ // See the License for the specific language governing permissions and // limitations under the License. -use protobuf_build::Builder; +//! Generates Rust protobuf bindings for `proto/logstore.proto`. +//! +//! The protoc compiler is selected in the following order: +//! +//! 1. Use the compiler specified by the `PROTOC` environment variable. +//! If the configured compiler is unusable or does not meet the minimum +//! supported version requirement, the build will panic instead of falling +//! back to another compiler. +//! +//! 2. Use a compatible `protoc` compiler available from `PATH`. +//! If it is unavailable or does not meet the minimum supported version +//! requirement, continue with the vendored compiler fallback. +//! +//! 3. Use a prebuilt vendored `protoc` compiler. +//! This provides a portable fallback when no supported system compiler is +//! available. + +use std::env; +use std::path::{Path, PathBuf}; +use std::process::Command; + +use protoc_rust::{Codegen, Customize}; + +/// Returns whether `protoc` runs and meets GreptimeDB's minimum supported version. +fn usable_protoc(protoc: &Path) -> bool { + let Ok(output) = Command::new(protoc).arg("--version").output() else { + return false; + }; + if !output.status.success() { + return false; + } + + let Some(version) = std::str::from_utf8(&output.stdout) + .ok() + .and_then(|output| output.trim().strip_prefix("libprotoc ")) + else { + return false; + }; + let mut components = version.split('.'); + matches!( + ( + components.next().and_then(|part| part.parse::().ok()), + components.next().and_then(|part| part.parse::().ok()), + ), + (Some(major), Some(minor)) if (major, minor) >= (3, 15) + ) +} + +/// Selects a compiler from `PROTOC`, `PATH`, or the vendored fallback, in that order. +/// +/// An invalid explicit `PROTOC` is rejected so configuration errors are not +/// silently masked. +fn protoc_path() -> PathBuf { + if let Some(protoc) = env::var_os("PROTOC") { + let protoc = PathBuf::from(protoc); + assert!(usable_protoc(&protoc), "PROTOC version not usable"); + return protoc; + } + + let protoc = PathBuf::from("protoc"); + if usable_protoc(&protoc) { + return protoc; + } + + protoc_bin_vendored::protoc_bin_path().expect("no bundled protoc for this platform") +} fn main() { - let base = std::env::var("CARGO_MANIFEST_DIR").unwrap_or_else(|_| ".".to_string()); - Builder::new() - .search_dir_for_protos(&format!("{base}/proto")) - .includes(&[format!("{base}/include"), format!("{base}/proto")]) - .include_google_protos() - .generate() + let base = + PathBuf::from(env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not set")); + let proto_dir = base.join("proto"); + let proto = proto_dir.join("logstore.proto"); + let out_dir = PathBuf::from(env::var_os("OUT_DIR").expect("OUT_DIR not set")).join("protos"); + + println!("cargo:rerun-if-env-changed=PROTOC"); + println!("cargo:rerun-if-env-changed=PATH"); + println!("cargo:rerun-if-changed={}", proto.display()); + std::fs::create_dir_all(&out_dir).expect("failed to create protobuf output directory"); + + Codegen::new() + .protoc_path(protoc_path()) + .out_dir(out_dir) + .input(proto) + .include(proto_dir) + .customize(Customize { + gen_mod_rs: Some(true), + ..Default::default() + }) + .run() + .expect("failed to generate log-store protobuf bindings") } diff --git a/src/standalone/Cargo.toml b/src/standalone/Cargo.toml index 16479991f41..836f891fa49 100644 --- a/src/standalone/Cargo.toml +++ b/src/standalone/Cargo.toml @@ -8,6 +8,7 @@ license.workspace = true workspace = true [dependencies] +api.workspace = true async-trait.workspace = true catalog.workspace = true client.workspace = true @@ -21,6 +22,7 @@ common-meta.workspace = true common-options.workspace = true common-procedure.workspace = true common-query.workspace = true +common-recordbatch.workspace = true common-stat.workspace = true common-telemetry.workspace = true common-time.workspace = true diff --git a/src/frontend/src/instance/standalone.rs b/src/standalone/src/datanode_manager.rs similarity index 97% rename from src/frontend/src/instance/standalone.rs rename to src/standalone/src/datanode_manager.rs index d34922956b3..13a17ea2b43 100644 --- a/src/frontend/src/instance/standalone.rs +++ b/src/standalone/src/datanode_manager.rs @@ -54,12 +54,12 @@ impl FlownodeManager for StandaloneDatanodeManager { } /// Relative to [client::region::RegionRequester] -pub struct RegionInvoker { +struct RegionInvoker { region_server: RegionServer, } impl RegionInvoker { - pub fn arc(region_server: RegionServer) -> Arc { + fn arc(region_server: RegionServer) -> Arc { Arc::new(Self { region_server }) } diff --git a/src/standalone/src/error.rs b/src/standalone/src/error.rs index 37859490392..8301c104b59 100644 --- a/src/standalone/src/error.rs +++ b/src/standalone/src/error.rs @@ -14,7 +14,7 @@ use std::any::Any; -use common_error::ext::{BoxedError, ErrorExt}; +use common_error::ext::{BoxedError, ErrorExt, RetryHint}; use common_error::status_code::StatusCode; use common_macro::stack_trace_debug; use snafu::{Location, Snafu}; @@ -30,6 +30,13 @@ pub enum Error { source: BoxedError, }, + #[snafu(display("Operation to region server failed"))] + InvokeRegionServer { + #[snafu(implicit)] + location: Location, + source: servers::error::Error, + }, + #[snafu(display("External error"))] External { source: BoxedError, @@ -49,6 +56,9 @@ pub enum Error { #[snafu(display("Invalid url scheme: {}", scheme))] InvalidUrlScheme { scheme: String }, + #[snafu(display("Invalid region request, reason: {}", reason))] + InvalidRegionRequest { reason: String }, + #[snafu(display("Repartition procedure is not supported in standalone mode"))] NoSupportRepartitionProcedure { #[snafu(implicit)] @@ -62,9 +72,11 @@ impl ErrorExt for Error { fn status_code(&self) -> StatusCode { match self { Error::OpenMetadataKvBackend { source, .. } => source.status_code(), + Error::InvokeRegionServer { source, .. } => source.status_code(), Error::External { source, .. } => source.status_code(), Error::ParseUrl { .. } => StatusCode::InvalidArguments, Error::InvalidUrlScheme { .. } => StatusCode::InvalidArguments, + Error::InvalidRegionRequest { .. } => StatusCode::IllegalState, Error::NoSupportRepartitionProcedure { .. } => StatusCode::Unsupported, } } @@ -72,4 +84,11 @@ impl ErrorExt for Error { fn as_any(&self) -> &dyn Any { self } + + fn retry_hint(&self) -> RetryHint { + match self { + Error::InvokeRegionServer { source, .. } => source.retry_hint(), + _ => RetryHint::NonRetryable, + } + } } diff --git a/src/standalone/src/lib.rs b/src/standalone/src/lib.rs index 5f3d2a2de4d..57c4efee249 100644 --- a/src/standalone/src/lib.rs +++ b/src/standalone/src/lib.rs @@ -12,12 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod datanode_manager; pub mod error; pub mod information_extension; pub mod metadata; pub mod options; pub mod procedure; +pub use datanode_manager::StandaloneDatanodeManager; pub use information_extension::StandaloneInformationExtension; pub use metadata::{build_metadata_kv_from_url, build_metadata_kvbackend}; pub use procedure::{StandaloneRepartitionProcedureFactory, build_procedure_manager}; diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index f93d7da804e..1085c21983e 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -49,14 +49,14 @@ use common_wal::config::{DatanodeWalConfig, MetasrvWalConfig}; use datanode::datanode::DatanodeBuilder; use flow::{FlownodeBuilder, FrontendClient, GrpcQueryHandlerWithBoxedError}; use frontend::frontend::Frontend; +use frontend::instance::Instance; use frontend::instance::builder::FrontendBuilder; -use frontend::instance::{Instance, StandaloneDatanodeManager}; use frontend::server::Services; use meta_srv::metasrv::{FLOW_ID_SEQ, TABLE_ID_SEQ}; use servers::grpc::GrpcOptions; use snafu::ResultExt; -use standalone::StandaloneRepartitionProcedureFactory; use standalone::options::StandaloneOptions; +use standalone::{StandaloneDatanodeManager, StandaloneRepartitionProcedureFactory}; use crate::test_util::{self, StorageType, TestGuard, create_tmp_dir_and_datanode_opts}; From 7fa6f5f98e0811f2f16bba686ed5ab2b2f9286ee Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:11:21 +0800 Subject: [PATCH 06/92] feat: add riscv64 cross-build support (#8820) * feat: add riscv64 cross-build support Add the missing build infrastructure for riscv64gc-unknown-linux-gnu. The codebase itself already compiles cleanly for riscv64 (verified with `cargo check --workspace --target riscv64gc-unknown-linux-gnu`): all architecture-sensitive dependencies support it (tikv-jemalloc-sys, aws-lc-sys, ring, pprof, simd-json). - .cargo/config.toml: set riscv64-linux-gnu-gcc as the linker for the riscv64gc-unknown-linux-gnu target - rust.yml: add a check-riscv64 CI job that cross-checks the whole workspace to prevent regressions from future dependency changes - docker/dev-builder/riscv64/Dockerfile: new cross dev-builder image with gcc/g++-riscv64-linux-gnu and the riscv64 rust target - Makefile: add dev-builder-riscv64 and build-riscv64-bin targets Verified end-to-end: the produced riscv64 binary starts in standalone mode under qemu and serves SQL (create/insert/select) over HTTP. Signed-off-by: Lei, HUANG * ci: build riscv64 artifacts in the release workflow - release.yml: add build-linux-riscv64-artifacts job that cross-compiles greptime for riscv64gc-unknown-linux-gnu on the amd64 runner with the dev-builder-riscv64 image, and uploads greptime-linux-riscv64-* artifacts. The job is wired into the needs of publish-github-release, release-cn-artifacts and stop-linux-amd64-runner. Integration tests are skipped since the cross-compiled binary cannot run on the host. - release-dev-builder-images.yaml + build-dev-builder-images action: build and push the dev-builder-riscv64 image to DockerHub, and sync it to ECR and ACR via skopeo like the other dev-builder images. - Makefile: DEV_BUILDER_RISCV64_IMAGE_TAG now defaults to DEV_BUILDER_IMAGE_TAG so the existing tag-bump automation keeps the riscv64 image tag in sync. Signed-off-by: Lei, HUANG * fix: forward cargo extension in riscv64 build Pass CARGO_EXTENSION through build-riscv64-bin just like the existing build-by-dev-builder target, so wrappers such as sccache are preserved inside the cross-build container. Signed-off-by: Lei, HUANG * ci: check riscv64 release feature graph Check all workspace targets with the servers/dashboard feature enabled so the riscv64 CI job covers the same optional dependency graph used by the release artifact build. Signed-off-by: Lei, HUANG * fix: gate riscv64 latest tags to main pushes Manual dev-builder workflow dispatches now publish only their immutable version tag. Update DockerHub and ECR latest tags only for the workflow's main-branch push event, preventing feature-branch builds from replacing the shared latest image. Signed-off-by: Lei, HUANG * docs: include riscv64 in release input description Update the build_linux_artifacts workflow input description to reflect that it now triggers amd64, arm64, and riscv64 artifact builds. Signed-off-by: Lei, HUANG * fix: fall back when riscv64 builder is unpublished Before building a release artifact, pull the pinned RISC-V dev-builder from ECR. If the image has not been published yet, build the same tag locally from the current Dockerfile so releases remain unblocked during the builder-tag update window. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- .cargo/config.toml | 3 + .../build-dev-builder-images/action.yml | 19 +++++++ .../workflows/release-dev-builder-images.yaml | 50 ++++++++++++++++- .github/workflows/release.yml | 52 +++++++++++++++++- .github/workflows/rust.yml | 26 +++++++++ Makefile | 26 +++++++++ docker/dev-builder/riscv64/Dockerfile | 55 +++++++++++++++++++ 7 files changed, 229 insertions(+), 2 deletions(-) create mode 100644 docker/dev-builder/riscv64/Dockerfile diff --git a/.cargo/config.toml b/.cargo/config.toml index 53dcb05f335..9ffc1961b09 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -1,6 +1,9 @@ [target.aarch64-unknown-linux-gnu] linker = "aarch64-linux-gnu-gcc" +[target.riscv64gc-unknown-linux-gnu] +linker = "riscv64-linux-gnu-gcc" + [alias] sqlness = "run --bin sqlness-runner --target-dir target/sqlness --" diff --git a/.github/actions/build-dev-builder-images/action.yml b/.github/actions/build-dev-builder-images/action.yml index 41f69193969..1ca37101302 100644 --- a/.github/actions/build-dev-builder-images/action.yml +++ b/.github/actions/build-dev-builder-images/action.yml @@ -31,6 +31,10 @@ inputs: description: Build dev-builder-android image required: false default: "true" + build-dev-builder-riscv64: + description: Build dev-builder-riscv64 image + required: false + default: "true" runs: using: composite steps: @@ -70,6 +74,21 @@ runs: IMAGE_NAMESPACE=${{ inputs.dockerhub-image-namespace }} \ DEV_BUILDER_IMAGE_TAG=${{ inputs.version }} + - name: Build and push dev-builder-riscv64 image # Only build image for amd64 platform (it cross-compiles riscv64 binaries). + shell: bash + if: ${{ inputs.build-dev-builder-riscv64 == 'true' }} + run: | + make dev-builder-riscv64 \ + IMAGE_REGISTRY=${{ inputs.dockerhub-image-registry }} \ + IMAGE_NAMESPACE=${{ inputs.dockerhub-image-namespace }} \ + DEV_BUILDER_RISCV64_IMAGE_TAG=${{ inputs.version }} + docker push ${{ inputs.dockerhub-image-registry }}/${{ inputs.dockerhub-image-namespace }}/dev-builder-riscv64:${{ inputs.version }} + if [[ "${{ github.event_name }}" == "push" ]]; then + docker tag ${{ inputs.dockerhub-image-registry }}/${{ inputs.dockerhub-image-namespace }}/dev-builder-riscv64:${{ inputs.version }} \ + ${{ inputs.dockerhub-image-registry }}/${{ inputs.dockerhub-image-namespace }}/dev-builder-riscv64:latest + docker push ${{ inputs.dockerhub-image-registry }}/${{ inputs.dockerhub-image-namespace }}/dev-builder-riscv64:latest + fi + - name: Build and push dev-builder-android image # Only build image for amd64 platform. shell: bash if: ${{ inputs.build-dev-builder-android == 'true' }} diff --git a/.github/workflows/release-dev-builder-images.yaml b/.github/workflows/release-dev-builder-images.yaml index e6a33dcceb0..bebd5ed61e7 100644 --- a/.github/workflows/release-dev-builder-images.yaml +++ b/.github/workflows/release-dev-builder-images.yaml @@ -24,6 +24,11 @@ on: description: Release dev-builder-android image required: false default: false + release_dev_builder_riscv64_image: + type: boolean + description: Release dev-builder-riscv64 image + required: false + default: false update_dev_builder_image_tag: type: boolean description: Update the DEV_BUILDER_IMAGE_TAG in Makefile and create a PR @@ -36,7 +41,7 @@ jobs: # The jobs are triggered by the following events: # 1. Manually triggered workflow_dispatch event # 2. Push event when the PR that modifies the `rust-toolchain.toml` or `docker/dev-builder/**` is merged to main - if: ${{ github.event_name == 'push' || inputs.release_dev_builder_ubuntu_image || inputs.release_dev_builder_centos_image || inputs.release_dev_builder_android_image }} + if: ${{ github.event_name == 'push' || inputs.release_dev_builder_ubuntu_image || inputs.release_dev_builder_centos_image || inputs.release_dev_builder_android_image || inputs.release_dev_builder_riscv64_image }} runs-on: ubuntu-latest outputs: version: ${{ steps.set-version.outputs.version }} @@ -68,6 +73,7 @@ jobs: build-dev-builder-ubuntu: ${{ inputs.release_dev_builder_ubuntu_image || github.event_name == 'push' }} build-dev-builder-centos: ${{ inputs.release_dev_builder_centos_image || github.event_name == 'push' }} build-dev-builder-android: ${{ inputs.release_dev_builder_android_image || github.event_name == 'push' }} + build-dev-builder-riscv64: ${{ inputs.release_dev_builder_riscv64_image || github.event_name == 'push' }} release-dev-builder-images-ecr: name: Release dev builder images to AWS ECR @@ -154,6 +160,35 @@ jobs: copy -a docker://docker.io/$IMAGE_NAMESPACE/dev-builder-android:latest \ docker://$ECR_IMAGE_REGISTRY/$ECR_IMAGE_NAMESPACE/dev-builder-android:latest + - name: Push dev-builder-riscv64 image + shell: bash + if: ${{ inputs.release_dev_builder_riscv64_image || github.event_name == 'push' }} + env: + IMAGE_VERSION: ${{ needs.release-dev-builder-images.outputs.version }} + IMAGE_NAMESPACE: ${{ vars.IMAGE_NAMESPACE }} + ECR_IMAGE_REGISTRY: ${{ vars.ECR_IMAGE_REGISTRY }} + ECR_IMAGE_NAMESPACE: ${{ vars.ECR_IMAGE_NAMESPACE }} + run: | + docker run -v "${DOCKER_CONFIG:-$HOME/.docker}:/root/.docker:ro" \ + -e "REGISTRY_AUTH_FILE=/root/.docker/config.json" \ + quay.io/skopeo/stable:latest \ + copy -a docker://docker.io/$IMAGE_NAMESPACE/dev-builder-riscv64:$IMAGE_VERSION \ + docker://$ECR_IMAGE_REGISTRY/$ECR_IMAGE_NAMESPACE/dev-builder-riscv64:$IMAGE_VERSION + + - name: Push latest dev-builder-riscv64 image + shell: bash + if: ${{ github.event_name == 'push' }} + env: + IMAGE_NAMESPACE: ${{ vars.IMAGE_NAMESPACE }} + ECR_IMAGE_REGISTRY: ${{ vars.ECR_IMAGE_REGISTRY }} + ECR_IMAGE_NAMESPACE: ${{ vars.ECR_IMAGE_NAMESPACE }} + run: | + docker run -v "${DOCKER_CONFIG:-$HOME/.docker}:/root/.docker:ro" \ + -e "REGISTRY_AUTH_FILE=/root/.docker/config.json" \ + quay.io/skopeo/stable:latest \ + copy -a docker://docker.io/$IMAGE_NAMESPACE/dev-builder-riscv64:latest \ + docker://$ECR_IMAGE_REGISTRY/$ECR_IMAGE_NAMESPACE/dev-builder-riscv64:latest + release-dev-builder-images-cn: # Note: Be careful issue: https://github.com/containers/skopeo/issues/1874 and we decide to use the latest stable skopeo container. name: Release dev builder images to CN region runs-on: ubuntu-latest @@ -209,6 +244,19 @@ jobs: quay.io/skopeo/stable:latest \ copy -a docker://docker.io/$IMAGE_NAMESPACE/dev-builder-android:$IMAGE_VERSION \ docker://$ACR_IMAGE_REGISTRY/$IMAGE_NAMESPACE/dev-builder-android:$IMAGE_VERSION + - name: Push dev-builder-riscv64 image + shell: bash + if: ${{ inputs.release_dev_builder_riscv64_image || github.event_name == 'push' }} + env: + IMAGE_VERSION: ${{ needs.release-dev-builder-images.outputs.version }} + IMAGE_NAMESPACE: ${{ vars.IMAGE_NAMESPACE }} + ACR_IMAGE_REGISTRY: ${{ vars.ACR_IMAGE_REGISTRY }} + run: | + docker run -v "${DOCKER_CONFIG:-$HOME/.docker}:/root/.docker:ro" \ + -e "REGISTRY_AUTH_FILE=/root/.docker/config.json" \ + quay.io/skopeo/stable:latest \ + copy -a docker://docker.io/$IMAGE_NAMESPACE/dev-builder-riscv64:$IMAGE_VERSION \ + docker://$ACR_IMAGE_REGISTRY/$IMAGE_NAMESPACE/dev-builder-riscv64:$IMAGE_VERSION update-dev-builder-image-tag: name: Update dev-builder image tag diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index dbb0f3c3146..52d62299a07 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -52,7 +52,7 @@ on: default: true build_linux_artifacts: type: boolean - description: Build linux artifacts (both amd64 and arm64) + description: Build linux artifacts (amd64, arm64, and riscv64) required: false default: false build_macos_artifacts: @@ -225,6 +225,52 @@ jobs: image-namespace: ${{ vars.ECR_IMAGE_NAMESPACE }} large-page-size: true + build-linux-riscv64-artifacts: + name: Build linux-riscv64 artifacts + if: ${{ inputs.build_linux_artifacts || inputs.publish_github_release || inputs.release_images || github.event_name == 'push' || github.event_name == 'schedule' }} + needs: [ + allocate-runners, + ] + runs-on: ${{ needs.allocate-runners.outputs.linux-amd64-runner }} + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + persist-credentials: false + + - name: Prepare riscv64 dev-builder image + shell: bash + env: + IMAGE_REGISTRY: ${{ vars.ECR_IMAGE_REGISTRY }} + IMAGE_NAMESPACE: ${{ vars.ECR_IMAGE_NAMESPACE }} + run: | + builder_tag=$(grep '^DEV_BUILDER_IMAGE_TAG ?=' Makefile | awk '{print $3}') + builder_image="${IMAGE_REGISTRY}/${IMAGE_NAMESPACE}/dev-builder-riscv64:${builder_tag}" + if ! docker pull "${builder_image}"; then + make dev-builder-riscv64 \ + IMAGE_REGISTRY="${IMAGE_REGISTRY}" \ + IMAGE_NAMESPACE="${IMAGE_NAMESPACE}" + fi + + # Cross-compile greptime for riscv64gc-unknown-linux-gnu on the amd64 runner + # with the riscv64 cross dev-builder image. Integration tests are skipped + # because the cross-compiled binary cannot run on the amd64 host. + - name: Build greptime riscv64 binary + shell: bash + run: | + make build-riscv64-bin \ + CARGO_PROFILE=${{ env.CARGO_PROFILE }} \ + FEATURES=servers/dashboard \ + IMAGE_REGISTRY=${{ vars.ECR_IMAGE_REGISTRY }} \ + IMAGE_NAMESPACE=${{ vars.ECR_IMAGE_NAMESPACE }} + + - name: Upload artifacts + uses: ./.github/actions/upload-artifacts + with: + artifacts-dir: greptime-linux-riscv64-${{ needs.allocate-runners.outputs.version }} + target-files: ./target/riscv64gc-unknown-linux-gnu/${{ env.CARGO_PROFILE }}/greptime + version: ${{ needs.allocate-runners.outputs.version }} + build-macos-artifacts: name: Build macOS artifacts strategy: @@ -520,6 +566,7 @@ jobs: build-linux-amd64-artifacts, build-linux-arm64-artifacts, build-macos-artifacts, + build-linux-riscv64-artifacts, build-windows-artifacts, release-images-to-dockerhub, ] @@ -568,6 +615,7 @@ jobs: (needs.prepare-release-validation.outputs.run-query-regression == 'false' && needs.query-regression-release.result == 'skipped')) && needs.build-linux-amd64-artifacts.result == 'success' && needs.build-linux-arm64-artifacts.result == 'success' && + needs.build-linux-riscv64-artifacts.result == 'success' && (needs.build-macos-artifacts.result == 'success' || needs.build-macos-artifacts.result == 'skipped') && (needs.build-windows-artifacts.result == 'success' || needs.build-windows-artifacts.result == 'skipped') && (needs.release-images-to-dockerhub.result == 'success' || @@ -581,6 +629,7 @@ jobs: build-linux-amd64-artifacts, build-linux-arm64-artifacts, build-macos-artifacts, + build-linux-riscv64-artifacts, build-windows-artifacts, release-images-to-dockerhub, prepare-release-validation, @@ -619,6 +668,7 @@ jobs: needs: [ allocate-runners, build-linux-amd64-artifacts, + build-linux-riscv64-artifacts, ] steps: - name: Checkout diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index b5bfa1a2968..6d12566b684 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -96,6 +96,32 @@ jobs: - name: Check unused dependencies run: make check-udeps + check-riscv64: + if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} + name: Check riscv64 build + runs-on: ubuntu-latest + timeout-minutes: 90 + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: arduino/setup-protoc@v3 + with: + repo-token: ${{ secrets.GITHUB_TOKEN }} + - name: Install riscv64 cross toolchain + run: sudo apt-get update && sudo apt-get install -y gcc-riscv64-linux-gnu g++-riscv64-linux-gnu + - uses: actions-rust-lang/setup-rust-toolchain@v1 + - name: Add riscv64 target + run: rustup target add riscv64gc-unknown-linux-gnu + - name: Rust Cache + uses: Swatinem/rust-cache@v2 + with: + shared-key: "check-riscv64" + cache-all-crates: "true" + save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Run cargo check for riscv64 + run: cargo check --locked --workspace --all-targets --features servers/dashboard --target riscv64gc-unknown-linux-gnu + conflict-check: if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} name: Check for conflict diff --git a/Makefile b/Makefile index de93a2f113a..5bdb63d3520 100644 --- a/Makefile +++ b/Makefile @@ -9,6 +9,9 @@ IMAGE_REGISTRY ?= docker.io IMAGE_NAMESPACE ?= greptime IMAGE_TAG ?= latest DEV_BUILDER_IMAGE_TAG ?= 2026-03-21-9c9d9e9e-20260331090344 +DEV_BUILDER_RISCV64_IMAGE_TAG ?= $(DEV_BUILDER_IMAGE_TAG) +DEV_BUILDER_RISCV64_BASE_IMAGE ?= ubuntu:22.04 +RISCV64_TARGET ?= riscv64gc-unknown-linux-gnu BUILDX_MULTI_PLATFORM_BUILD ?= false BUILDX_BUILDER_NAME ?= gtbuilder BASE_IMAGE ?= ubuntu @@ -98,6 +101,21 @@ build-by-dev-builder: ## Build greptime by dev-builder. RELEASE=${RELEASE} \ CARGO_BUILD_EXTRA_OPTS="${CARGO_BUILD_EXTRA_OPTS}" +.PHONY: build-riscv64-bin +build-riscv64-bin: ## Build greptime binary for riscv64 (linux-gnu) by the riscv64 cross dev-builder image. + docker run --network=host \ + ${ASSEMBLED_EXTRA_BUILD_ENV} \ + -v ${PWD}:/greptimedb -v ${CARGO_REGISTRY_CACHE}:/root/.cargo/registry -v ${CARGO_GIT_CACHE}:/root/.cargo/git \ + -w /greptimedb ${IMAGE_REGISTRY}/${IMAGE_NAMESPACE}/dev-builder-riscv64:${DEV_BUILDER_RISCV64_IMAGE_TAG} \ + make build \ + CARGO_EXTENSION="${CARGO_EXTENSION}" \ + CARGO_PROFILE=${CARGO_PROFILE} \ + FEATURES=${FEATURES} \ + TARGET_DIR=${TARGET_DIR} \ + TARGET=${RISCV64_TARGET} \ + RELEASE=${RELEASE} \ + CARGO_BUILD_EXTRA_OPTS="${CARGO_BUILD_EXTRA_OPTS}" + .PHONY: build-android-bin build-android-bin: ## Build greptime binary for android. docker run --network=host \ @@ -158,6 +176,14 @@ dev-builder: multi-platform-buildx ## Build dev-builder image. -f docker/dev-builder/${BASE_IMAGE}/Dockerfile \ -t ${IMAGE_REGISTRY}/${IMAGE_NAMESPACE}/dev-builder-${BASE_IMAGE}:${DEV_BUILDER_IMAGE_TAG} ${BUILDX_MULTI_PLATFORM_BUILD_OPTS} . +.PHONY: dev-builder-riscv64 +dev-builder-riscv64: ## Build the riscv64 cross dev-builder image. + docker build --network=host \ + --build-arg="BASE_IMAGE=${DEV_BUILDER_RISCV64_BASE_IMAGE}" \ + --build-arg="RUST_TOOLCHAIN=${RUST_TOOLCHAIN}" \ + -f docker/dev-builder/riscv64/Dockerfile \ + -t ${IMAGE_REGISTRY}/${IMAGE_NAMESPACE}/dev-builder-riscv64:${DEV_BUILDER_RISCV64_IMAGE_TAG} . + .PHONY: multi-platform-buildx multi-platform-buildx: ## Create buildx multi-platform builder. docker buildx inspect ${BUILDX_BUILDER_NAME} || docker buildx create --name ${BUILDX_BUILDER_NAME} --driver docker-container --bootstrap --use diff --git a/docker/dev-builder/riscv64/Dockerfile b/docker/dev-builder/riscv64/Dockerfile new file mode 100644 index 00000000000..8505fa3b49a --- /dev/null +++ b/docker/dev-builder/riscv64/Dockerfile @@ -0,0 +1,55 @@ +# Cross-build environment for producing riscv64 (riscv64gc-unknown-linux-gnu) +# greptime binaries on x86_64/aarch64 hosts. +ARG BASE_IMAGE=ubuntu:22.04 +FROM ${BASE_IMAGE} + +ENV LANG=en_US.utf8 +WORKDIR /greptimedb + +RUN apt-get update && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y software-properties-common +# Install dependencies, including the riscv64 cross toolchain. +RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \ + libssl-dev \ + tzdata \ + curl \ + unzip \ + ca-certificates \ + git \ + build-essential \ + pkg-config \ + gcc-riscv64-linux-gnu \ + g++-riscv64-linux-gnu + +ARG TARGETPLATFORM +RUN echo "host platform: $TARGETPLATFORM" + +ARG PROTOBUF_VERSION=29.3 + +# Install protobuf for the host arch (build scripts and protoc run on the host), +# because the one in the apt is too old (v3.12). +RUN if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \ + curl -OL https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOBUF_VERSION}/protoc-${PROTOBUF_VERSION}-linux-aarch_64.zip && \ + unzip -o protoc-${PROTOBUF_VERSION}-linux-aarch_64.zip -d protoc3; \ +else \ + curl -OL https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOBUF_VERSION}/protoc-${PROTOBUF_VERSION}-linux-x86_64.zip && \ + unzip -o protoc-${PROTOBUF_VERSION}-linux-x86_64.zip -d protoc3; \ +fi +RUN mv -f protoc3/bin/* /usr/local/bin/ +# Use cp -r instead of mv so that it also works when the base image already +# contains protobuf headers under /usr/local/include. +RUN cp -r protoc3/include/* /usr/local/include/ + +# Silence all `safe.directory` warnings, to avoid the "detect dubious repository" +# error when building with the repository mounted from a different user. +RUN git config --global --add safe.directory '*' + +# Install Rust. +SHELL ["/bin/bash", "-c"] +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- --no-modify-path --default-toolchain none -y +ENV PATH=/root/.cargo/bin/:$PATH + +# Install Rust toolchains and the riscv64 target std. +ARG RUST_TOOLCHAIN +RUN rustup toolchain install ${RUST_TOOLCHAIN} +RUN rustup target add --toolchain ${RUST_TOOLCHAIN} riscv64gc-unknown-linux-gnu From 9fcfeff9ce876522fe2c47d2c47d7ea5fffa3c89 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Mon, 10 Aug 2026 16:28:11 +0800 Subject: [PATCH 07/92] feat(auth): support HTTP bearer-token authentication (#8719) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(auth): support HTTP bearer-token authentication (#8718) Adds an opt-in bearer-token (JWT / OAuth2) authentication path to the HTTP layer, so clients can authenticate with `Authorization: Bearer ` against any `/v1/` interface. Today such requests are rejected with `UnsupportedAuthScheme("bearer")` -> 401 before any handler runs. The token is treated as opaque by the server; validation and identity derivation stay in the UserProvider, so JWT/JWKS/OIDC policy remains pluggable and out of core. Changes: - `auth::UserProvider` gains `auth_token(token, catalog, schema) -> Result` with a default that rejects (`Error::UnsupportedAuthMethod`), so password-only providers keep today's behavior. A provider that supports token auth overrides it to validate the token, resolve it to a user, and authorize the connection. - `auth::Error::UnsupportedAuthMethod` for the default-reject case. - `servers::http::authorize::inner_auth` extracts a bearer token (`extract_bearer_token`) and, when present, authenticates via `UserProvider::auth_token`; otherwise it falls through unchanged to the username/password path (Basic / influxdb / splunk). Basic and bearer coexist on the same server. Backward compatible: the default impl preserves existing behavior, and non-bearer requests take the exact same path as before. Tests: - `extract_bearer_token` recognizes `Bearer` (either header) and ignores Basic/Token/Splunk/empty. - `inner_auth` dispatches a bearer token to `auth_token` and populates the QueryContext user on success; rejects on failure. - A password-only provider (default `auth_token`) rejects bearer tokens. Refs: #8718 * chore: fmt * refactor(auth): address bearer-auth review feedback (#8719) Address the review comments on the HTTP bearer-token authentication PR: - Match the `Bearer` scheme case-insensitively (RFC 9110 §11.1) via `eq_ignore_ascii_case`. `extract_bearer_token` previously only accepted `Bearer`/`bearer`, so a valid `BEARER ` fell through to `UnsupportedAuthScheme` and never reached the provider. The opaque token itself is deliberately not lowercased. - Return `Option<&str>` (borrowing the request headers) instead of `Option`, avoiding an allocation per bearer request. - Rename `UserProvider::auth_token` -> `auth_bearer_token` for clarity. - Route bearer-auth failures on Splunk HEC requests through `splunk_hec_err` (FORBIDDEN, code 4) instead of the generic 401 `ErrorResponse`, so HEC clients retain their `{"text":"Invalid token","code":4}` endpoint contract. Adds test coverage for case-insensitive scheme parsing (token preserved verbatim) and a regression test for the bearer/splunk routing path. Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- src/auth/src/error.rs | 8 ++ src/auth/src/user_provider.rs | 29 +++- src/servers/src/http/authorize.rs | 226 +++++++++++++++++++++++++++++- 3 files changed, 259 insertions(+), 4 deletions(-) diff --git a/src/auth/src/error.rs b/src/auth/src/error.rs index 2b339b86554..f99a485db4e 100644 --- a/src/auth/src/error.rs +++ b/src/auth/src/error.rs @@ -88,6 +88,13 @@ pub enum Error { #[snafu(implicit)] location: Location, }, + + #[snafu(display("Unsupported authentication method: {}", method))] + UnsupportedAuthMethod { + method: String, + #[snafu(implicit)] + location: Location, + }, } impl ErrorExt for Error { @@ -105,6 +112,7 @@ impl ErrorExt for Error { Error::UserPasswordMismatch { .. } => StatusCode::UserPasswordMismatch, Error::AccessDenied { .. } => StatusCode::AccessDenied, Error::PermissionDenied { .. } => StatusCode::PermissionDenied, + Error::UnsupportedAuthMethod { .. } => StatusCode::UserPasswordMismatch, } } diff --git a/src/auth/src/user_provider.rs b/src/auth/src/user_provider.rs index 60c178fb598..985e25978e8 100644 --- a/src/auth/src/user_provider.rs +++ b/src/auth/src/user_provider.rs @@ -34,8 +34,8 @@ use crate::common::{ PgScramSha256Verifier, auth_mysql_with_hash_stage_2, }; use crate::error::{ - IllegalParamSnafu, InvalidConfigSnafu, IoSnafu, Result, UnsupportedPasswordTypeSnafu, - UserNotFoundSnafu, UserPasswordMismatchSnafu, + IllegalParamSnafu, InvalidConfigSnafu, IoSnafu, Result, UnsupportedAuthMethodSnafu, + UnsupportedPasswordTypeSnafu, UserNotFoundSnafu, UserPasswordMismatchSnafu, }; use crate::user_info::{DefaultUserInfo, PermissionMode}; use crate::{UserInfoRef, auth_mysql}; @@ -66,6 +66,31 @@ pub trait UserProvider: Send + Sync { Ok(user_info) } + /// Authenticates and authorizes an opaque bearer token (e.g. a JWT or an + /// OAuth2 access token). + /// + /// Unlike [auth()](Self::auth), the caller has no `Identity`/`Password` — + /// the provider validates the token and *derives* the identity from it. + /// The token is opaque to the server, so JWT/JWKS/OIDC validation policy + /// stays pluggable and out of core. + /// + /// The default rejects token auth with + /// [`Error::UnsupportedAuthMethod`], so password-only providers keep + /// today's behavior. Providers that support token auth override this to + /// validate the token, resolve it to a user, and + /// [`authorize`](Self::authorize) the connection. + async fn auth_bearer_token( + &self, + _token: &str, + _catalog: &str, + _schema: &str, + ) -> Result { + UnsupportedAuthMethodSnafu { + method: "bearer token", + } + .fail() + } + async fn postgres_auth_info(&self, _id: Identity<'_>) -> Result { Ok(PgAuthInfo::Cleartext) } diff --git a/src/servers/src/http/authorize.rs b/src/servers/src/http/authorize.rs index a6d8ce2438e..755bd4d4c24 100644 --- a/src/servers/src/http/authorize.rs +++ b/src/servers/src/http/authorize.rs @@ -80,7 +80,36 @@ pub async fn inner_auth( return Ok(req); }; - // 3. get username and pwd + // 3. bearer token auth (JWT / OAuth2). When an `Authorization: Bearer + // ` header is present, authenticate via + // [`UserProvider::auth_bearer_token`]; otherwise fall through to the + // username/password path (Basic / influxdb / splunk) below. + if let Some(token) = extract_bearer_token(&req) { + match user_provider + .auth_bearer_token(token, &catalog, &schema) + .await + { + Ok(userinfo) => { + query_ctx.set_current_user(userinfo); + let _ = req.extensions_mut().insert(query_ctx); + return Ok(req); + } + Err(e) => { + warn!(e; "bearer token authentication failed"); + crate::metrics::METRIC_AUTH_FAILURE + .with_label_values(&[e.status_code().as_ref()]) + .inc(); + // Splunk HEC clients expect `{"text":"Invalid token","code":4}` + // (FORBIDDEN), not the generic 401 `ErrorResponse`. + if is_splunk_request(&req) { + return Err(splunk_hec_err(StatusCode::FORBIDDEN, 4)); + } + return Err(err_response(e)); + } + } + } + + // 4. get username and pwd let (username, password) = match extract_username_and_password(&req) { Ok((username, password)) => (username, password), Err(e) => { @@ -100,7 +129,7 @@ pub async fn inner_auth( } }; - // 4. auth + // 5. auth match user_provider .auth( auth::Identity::UserId(&username, None), @@ -307,6 +336,29 @@ impl From for api::v1::auth_header::AuthScheme { type Credential<'a> = &'a str; +/// Extracts an opaque bearer token from an `Authorization: Bearer ` +/// header (the standard or `x-greptime-auth` header). +/// +/// Returns `None` for any other scheme (`Basic`, influxdb `Token`, splunk +/// `Splunk`, …) so the caller can fall through to the username/password path. +fn extract_bearer_token(req: &Request) -> Option<&str> { + let header = req + .headers() + .get(AUTHORIZATION_HEADER) + .or_else(|| req.headers().get(http::header::AUTHORIZATION))?; + let value = header.to_str().ok()?; + // HTTP authentication schemes are case-insensitive (RFC 9110 §11.1), so + // match the scheme with `eq_ignore_ascii_case` — but never lowercase the + // *token* itself, which is opaque and case-sensitive. Returns a borrow + // into the request's headers, so there is no allocation. + let (scheme, token) = value.split_once(' ')?; + if !scheme.eq_ignore_ascii_case("bearer") { + return None; + } + let token = token.trim_start(); + (!token.is_empty()).then_some(token) +} + fn auth_header(req: &Request) -> Result { let auth_header = req .headers() @@ -610,4 +662,174 @@ mod tests { (Some("123"), Some("4")) ); } + + #[test] + fn test_extract_bearer_token() { + let bearer = |scheme: &str, val: &str| { + mock_http_request(Some(&format!("{scheme} {val}")), None).unwrap() + }; + + // Standard bearer scheme, on either header. + assert_eq!( + extract_bearer_token(&bearer("Bearer", "abc.def.ghi")), + Some("abc.def.ghi") + ); + assert_eq!(extract_bearer_token(&bearer("bearer", "tok")), Some("tok")); + // HTTP schemes are case-insensitive (RFC 9110 §11.1); the token + // itself is opaque and must NOT be lowercased. + assert_eq!( + extract_bearer_token(&bearer("BEARER", "ABC.DEF.GHI")), + Some("ABC.DEF.GHI") + ); + assert_eq!(extract_bearer_token(&bearer("BeArEr", "tok")), Some("tok")); + let mut req = mock_http_request(Some("Bearer xyz"), None).unwrap(); + req.headers_mut().insert( + AUTHORIZATION_HEADER, + "Bearer from-x-greptime".parse().unwrap(), + ); + assert_eq!(extract_bearer_token(&req), Some("from-x-greptime")); + + // Non-bearer schemes are ignored so the caller falls through to Basic. + assert_eq!(extract_bearer_token(&bearer("Basic", "dXNlcjpwYXNz")), None); + assert_eq!(extract_bearer_token(&bearer("Token", "u:p")), None); + assert_eq!(extract_bearer_token(&bearer("Splunk", "u:p")), None); + + // No header, empty token. + assert_eq!( + extract_bearer_token(&mock_http_request(None, None).unwrap()), + None + ); + assert_eq!(extract_bearer_token(&bearer("Bearer", "")), None); + } + + /// A `UserProvider` that resolves exactly one bearer token to a known user + /// and rejects everything else (including password auth). + struct TokenUserProvider { + token: String, + user: auth::UserInfoRef, + } + + #[async_trait::async_trait] + impl auth::UserProvider for TokenUserProvider { + fn name(&self) -> &str { + "token-test" + } + + async fn authenticate( + &self, + _: auth::Identity<'_>, + _: auth::Password<'_>, + ) -> auth::error::Result { + unreachable!("password auth should not be reached for a bearer request") + } + + async fn authorize( + &self, + _: &str, + _: &str, + _: &auth::UserInfoRef, + ) -> auth::error::Result<()> { + Ok(()) + } + + async fn auth_bearer_token( + &self, + token: &str, + _: &str, + _: &str, + ) -> auth::error::Result { + if token == self.token { + Ok(self.user.clone()) + } else { + auth::error::UnsupportedAuthMethodSnafu { + method: "bearer token", + } + .fail() + } + } + } + + #[tokio::test] + async fn test_bearer_token_dispatches_to_auth_bearer_token() { + let provider = TokenUserProvider { + token: "good-token".to_string(), + user: auth::userinfo_by_name(Some("alice".into())), + }; + let req = mock_http_request(Some("Bearer good-token"), None).unwrap(); + + let req = inner_auth::<()>( + Some(std::sync::Arc::new(provider) as auth::UserProviderRef), + req, + ) + .await + .expect("valid bearer token authenticates"); + + let user = req + .extensions() + .get::() + .expect("query context is populated") + .current_user(); + assert_eq!(user.username(), "alice"); + } + + #[tokio::test] + async fn test_bearer_token_failure_rejects() { + let provider = TokenUserProvider { + token: "good-token".to_string(), + user: auth::userinfo_by_name(Some("alice".into())), + }; + let req = mock_http_request(Some("Bearer bad-token"), None).unwrap(); + + let result = inner_auth::<()>( + Some(std::sync::Arc::new(provider) as auth::UserProviderRef), + req, + ) + .await; + assert!(result.is_err(), "an invalid bearer token is rejected"); + } + + #[tokio::test] + async fn test_default_provider_rejects_bearer() { + // A password-only provider uses the default `auth_bearer_token`, which rejects. + let provider = + auth::user_provider_from_option("static_user_provider:cmd:alice=s3cret").unwrap(); + let req = mock_http_request(Some("Bearer some-jwt"), None).unwrap(); + let result = inner_auth::<()>(Some(provider), req).await; + assert!( + result.is_err(), + "password-only providers reject bearer tokens" + ); + } + + /// A bearer-token failure on a Splunk HEC request must keep the HEC + /// contract — `{"text":"Invalid token","code":4}` with FORBIDDEN — + /// instead of falling through to the generic 401 `ErrorResponse`. + /// Regression for the bearer/splunk routing gap. + #[tokio::test] + async fn test_bearer_failure_on_splunk_keeps_hec_contract() { + let provider = TokenUserProvider { + token: "good-token".to_string(), + user: auth::userinfo_by_name(Some("alice".into())), + }; + let req = mock_http_request( + Some("Bearer bad-token"), + Some("http://127.0.0.1/v1/splunk/services/collector/event"), + ) + .unwrap(); + + let resp = inner_auth::<()>( + Some(std::sync::Arc::new(provider) as auth::UserProviderRef), + req, + ) + .await + .expect_err("an invalid bearer token is rejected"); + + assert_eq!(resp.status(), StatusCode::FORBIDDEN); + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let payload: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(payload["code"], 4); + assert_eq!(payload["text"], "Invalid token"); + } } From 4ce49ec33eb0b416cbd86f737631a0c37cf12774 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:06:20 +0800 Subject: [PATCH 08/92] feat(promql): support mixed sample ranges (#8784) * feat(promql): support mixed sample ranges Signed-off-by: shuiyisong * chore: update test Signed-off-by: shuiyisong * chore: add comments Signed-off-by: shuiyisong * fix: CR issue Signed-off-by: shuiyisong * fix: return error than panic Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/promql/src/functions.rs | 2 +- src/promql/src/functions/native_histogram.rs | 997 +++++++++++++++++- src/query/src/promql/planner.rs | 343 +++++- .../src/query_engine/default_serializer.rs | 109 +- 4 files changed, 1419 insertions(+), 32 deletions(-) diff --git a/src/promql/src/functions.rs b/src/promql/src/functions.rs index b534bb06928..b34c21c5569 100644 --- a/src/promql/src/functions.rs +++ b/src/promql/src/functions.rs @@ -45,7 +45,7 @@ pub use double_exponential_smoothing::DoubleExponentialSmoothing; pub use extrapolate_rate::{Delta, Increase, Rate}; pub use idelta::IDelta; pub use native_histogram::{ - NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg, + MixedRange, NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg, NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges, NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta, NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, diff --git a/src/promql/src/functions/native_histogram.rs b/src/promql/src/functions/native_histogram.rs index 4a4e3adad5c..1863e663b82 100644 --- a/src/promql/src/functions/native_histogram.rs +++ b/src/promql/src/functions/native_histogram.rs @@ -25,7 +25,8 @@ use datafusion::arrow::array::{ Array, ArrayRef, BooleanArray, Float64Array, Float64Builder, Int64Array, StringBuilder, StructArray, TimestampMillisecondArray, UInt64Array, }; -use datafusion::arrow::datatypes::{DataType, TimeUnit}; +use datafusion::arrow::compute::filter; +use datafusion::arrow::datatypes::{DataType, Field, TimeUnit}; use datafusion::common::{DataFusionError, Result as DfResult}; use datafusion::logical_expr::{Accumulator as DfAccumulator, AggregateUDF, ScalarUDF, Volatility}; use datafusion::physical_plan::ColumnarValue; @@ -33,7 +34,11 @@ use datafusion_common::ScalarValue; use datafusion_expr::function::AccumulatorArgs; use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature, create_udaf, create_udf}; -use crate::functions::{extract_array, extract_range_dict}; +use crate::functions::{ + AvgOverTime, Deriv, DoubleExponentialSmoothing, IDelta, Increase, LastOverTime, MaxOverTime, + MinOverTime, PredictLinear, QuantileOverTime, Rate, StddevOverTime, StdvarOverTime, + SumOverTime, extract_array, extract_range_dict, +}; use crate::range_array::{RangeArray, unpack}; fn extract_histogram_array(value: &ColumnarValue, func_name: &str) -> DfResult { @@ -1264,6 +1269,783 @@ impl NativeHistogramResets { } } +/// Coordinated float/native-histogram range evaluation. +/// +/// The function name is passed as the first scalar argument so these two UDF names are enough for +/// distributed plan decoding. The remaining leading arguments are timestamp, float, and histogram +/// ranges, followed by the ordinary function arguments. +pub struct MixedRange; + +impl MixedRange { + const fn float_name() -> &'static str { + "prom_mixed_range_float" + } + + const fn histogram_name() -> &'static str { + "prom_mixed_range_histogram" + } + + pub fn float_udf(collector: Option) -> ScalarUDF { + ScalarUDF::new_from_impl(MixedRangeUdf::new(MixedRangeOutput::Float, collector)) + } + + pub fn histogram_udf(collector: Option) -> ScalarUDF { + ScalarUDF::new_from_impl(MixedRangeUdf::new(MixedRangeOutput::Histogram, collector)) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +enum MixedRangeOutput { + Float, + Histogram, +} + +impl MixedRangeOutput { + fn name(self) -> &'static str { + match self { + Self::Float => MixedRange::float_name(), + Self::Histogram => MixedRange::histogram_name(), + } + } + + fn data_type(self) -> DataType { + match self { + Self::Float => DataType::Float64, + Self::Histogram => native_histogram_arrow_type(), + } + } +} + +#[derive(Debug, Clone)] +struct MixedRangeUdf { + output: MixedRangeOutput, + signature: Signature, + collector: Option, +} + +impl MixedRangeUdf { + fn new(output: MixedRangeOutput, collector: Option) -> Self { + Self { + output, + signature: Signature::variadic_any(Volatility::Volatile), + collector, + } + } +} + +impl PartialEq for MixedRangeUdf { + fn eq(&self, other: &Self) -> bool { + self.output == other.output + } +} + +impl Eq for MixedRangeUdf {} + +impl Hash for MixedRangeUdf { + fn hash(&self, state: &mut H) { + self.output.hash(state); + } +} + +impl ScalarUDFImpl for MixedRangeUdf { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + self.output.name() + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> DfResult { + Ok(self.output.data_type()) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DfResult { + let collector = args + .config_options + .extensions + .get::() + .cloned() + .or_else(|| self.collector.clone()); + mixed_range(&args, self.output, collector) + } +} + +#[derive(Debug, Clone, Copy)] +enum MixedRangeFunction { + Rate, + Increase, + Delta, + IDelta, + IRate, + Changes, + Resets, + AvgOverTime, + MinOverTime, + MaxOverTime, + SumOverTime, + CountOverTime, + LastOverTime, + AbsentOverTime, + PresentOverTime, + StddevOverTime, + StdvarOverTime, + QuantileOverTime, + Deriv, + PredictLinear, + DoubleExponentialSmoothing, + HoltWinters, +} + +impl MixedRangeFunction { + fn parse(value: &ColumnarValue) -> DfResult { + let name = match value { + ColumnarValue::Scalar(ScalarValue::Utf8(Some(name))) + | ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(name))) + | ColumnarValue::Scalar(ScalarValue::Utf8View(Some(name))) => name.as_str(), + other => { + return Err(DataFusionError::Execution(format!( + "mixed range function name must be a non-null string scalar, found {}", + other.data_type() + ))); + } + }; + + match name { + "rate" => Ok(Self::Rate), + "increase" => Ok(Self::Increase), + "delta" => Ok(Self::Delta), + "idelta" => Ok(Self::IDelta), + "irate" => Ok(Self::IRate), + "changes" => Ok(Self::Changes), + "resets" => Ok(Self::Resets), + "avg_over_time" => Ok(Self::AvgOverTime), + "min_over_time" => Ok(Self::MinOverTime), + "max_over_time" => Ok(Self::MaxOverTime), + "sum_over_time" => Ok(Self::SumOverTime), + "count_over_time" => Ok(Self::CountOverTime), + "last_over_time" => Ok(Self::LastOverTime), + "absent_over_time" => Ok(Self::AbsentOverTime), + "present_over_time" => Ok(Self::PresentOverTime), + "stddev_over_time" => Ok(Self::StddevOverTime), + "stdvar_over_time" => Ok(Self::StdvarOverTime), + "quantile_over_time" => Ok(Self::QuantileOverTime), + "deriv" => Ok(Self::Deriv), + "predict_linear" => Ok(Self::PredictLinear), + "double_exponential_smoothing" => Ok(Self::DoubleExponentialSmoothing), + "holt_winters" => Ok(Self::HoltWinters), + _ => Err(DataFusionError::Execution(format!( + "unsupported mixed range function: {name}" + ))), + } + } + + fn name(self) -> &'static str { + match self { + Self::Rate => "rate", + Self::Increase => "increase", + Self::Delta => "delta", + Self::IDelta => "idelta", + Self::IRate => "irate", + Self::Changes => "changes", + Self::Resets => "resets", + Self::AvgOverTime => "avg_over_time", + Self::MinOverTime => "min_over_time", + Self::MaxOverTime => "max_over_time", + Self::SumOverTime => "sum_over_time", + Self::CountOverTime => "count_over_time", + Self::LastOverTime => "last_over_time", + Self::AbsentOverTime => "absent_over_time", + Self::PresentOverTime => "present_over_time", + Self::StddevOverTime => "stddev_over_time", + Self::StdvarOverTime => "stdvar_over_time", + Self::QuantileOverTime => "quantile_over_time", + Self::Deriv => "deriv", + Self::PredictLinear => "predict_linear", + Self::DoubleExponentialSmoothing => "double_exponential_smoothing", + Self::HoltWinters => "holt_winters", + } + } + + fn policy(self) -> MixedRangePolicy { + match self { + Self::Rate | Self::Increase | Self::Delta | Self::AvgOverTime | Self::SumOverTime => { + MixedRangePolicy::DropMixed + } + Self::IDelta | Self::IRate => MixedRangePolicy::LastTwo, + Self::LastOverTime => MixedRangePolicy::Last, + Self::Changes + | Self::Resets + | Self::CountOverTime + | Self::AbsentOverTime + | Self::PresentOverTime => MixedRangePolicy::Combined, + Self::MinOverTime + | Self::MaxOverTime + | Self::StddevOverTime + | Self::StdvarOverTime + | Self::QuantileOverTime + | Self::Deriv + | Self::PredictLinear + | Self::DoubleExponentialSmoothing + | Self::HoltWinters => MixedRangePolicy::FloatOnly, + } + } + + fn float_udf(self) -> Option { + match self { + Self::Rate => Some(Rate::scalar_udf()), + Self::Increase => Some(Increase::scalar_udf()), + Self::Delta => Some(crate::functions::Delta::scalar_udf()), + Self::IDelta => Some(IDelta::::scalar_udf()), + Self::IRate => Some(IDelta::::scalar_udf()), + Self::AvgOverTime => Some(AvgOverTime::scalar_udf()), + Self::MinOverTime => Some(MinOverTime::scalar_udf()), + Self::MaxOverTime => Some(MaxOverTime::scalar_udf()), + Self::SumOverTime => Some(SumOverTime::scalar_udf()), + Self::LastOverTime => Some(LastOverTime::scalar_udf()), + Self::StddevOverTime => Some(StddevOverTime::scalar_udf()), + Self::StdvarOverTime => Some(StdvarOverTime::scalar_udf()), + Self::QuantileOverTime => Some(QuantileOverTime::scalar_udf()), + Self::Deriv => Some(Deriv::scalar_udf()), + Self::PredictLinear => Some(PredictLinear::scalar_udf()), + Self::DoubleExponentialSmoothing | Self::HoltWinters => { + Some(DoubleExponentialSmoothing::scalar_udf()) + } + Self::Changes + | Self::Resets + | Self::CountOverTime + | Self::AbsentOverTime + | Self::PresentOverTime => None, + } + } + + fn histogram_udf(self, collector: Option) -> Option { + match self { + Self::Rate => Some(NativeHistogramRate::scalar_udf_with_collector(collector)), + Self::Increase => Some(NativeHistogramIncrease::scalar_udf_with_collector( + collector, + )), + Self::Delta => Some(NativeHistogramDelta::scalar_udf_with_collector(collector)), + Self::IDelta => Some(NativeHistogramIDelta::scalar_udf_with_collector(collector)), + Self::IRate => Some(NativeHistogramIRate::scalar_udf_with_collector(collector)), + Self::AvgOverTime => Some(NativeHistogramAvgOverTime::scalar_udf_with_collector( + collector, + )), + Self::SumOverTime => Some(NativeHistogramSumOverTime::scalar_udf_with_collector( + collector, + )), + Self::LastOverTime => Some(NativeHistogramLastOverTime::scalar_udf()), + _ => None, + } + } +} + +#[derive(Debug, Clone, Copy)] +enum MixedRangePolicy { + DropMixed, + LastTwo, + Last, + Combined, + FloatOnly, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SampleLane { + Float, + Histogram, +} + +fn mixed_range( + args: &ScalarFunctionArgs, + output: MixedRangeOutput, + collector: Option, +) -> DfResult { + if args.args.len() < 4 { + return Err(DataFusionError::Plan(format!( + "{} function should have at least 4 inputs", + output.name() + ))); + } + let function = MixedRangeFunction::parse(&args.args[0])?; + let name = function.name(); + let ts_range = extract_range_dict( + &args.args[1], + name, + "timestamp range vector", + &DataType::Timestamp(TimeUnit::Millisecond, None), + )?; + let float_range = extract_range_dict( + &args.args[2], + name, + "float range vector", + &DataType::Float64, + )?; + let histogram_range = extract_range_dict( + &args.args[3], + name, + "native histogram range vector", + &native_histogram_arrow_type(), + )?; + let keys = ts_range.keys().values(); + if float_range.keys().values() != keys || histogram_range.keys().values() != keys { + return Err(DataFusionError::Execution(format!( + "{name}: timestamp, float, and native histogram ranges should have the same layout" + ))); + } + if args.number_rows != keys.len() { + return Err(DataFusionError::Execution(format!( + "{name}: range inputs have {} windows but the batch has {} rows", + keys.len(), + args.number_rows + ))); + } + + let timestamps = ts_range + .values() + .as_any() + .downcast_ref::() + .expect("validated timestamp range"); + let floats = float_range + .values() + .as_any() + .downcast_ref::() + .expect("validated float range"); + let histograms = histogram_range + .values() + .as_any() + .downcast_ref::() + .expect("validated native histogram range"); + if floats.len() != timestamps.len() || histograms.len() != timestamps.len() { + return Err(DataFusionError::Execution(format!( + "{name}: timestamp, float, and native histogram values should be row-aligned" + ))); + } + + let bounds = checked_window_bounds(keys, timestamps.len(), name)?; + let float_valid = (0..floats.len()) + .map(|row| floats.is_valid(row)) + .collect::>(); + let histogram_valid = (0..histograms.len()) + .map(|row| histograms.is_valid(row)) + .collect::>(); + let float_prefix = validity_prefix(&float_valid, name)?; + let histogram_prefix = validity_prefix(&histogram_valid, name)?; + + if matches!(function.policy(), MixedRangePolicy::Combined) { + if output != MixedRangeOutput::Float { + return Err(DataFusionError::Execution(format!( + "{name} does not return native histograms" + ))); + } + return combined_range_float( + function, + &bounds, + timestamps, + floats, + histograms, + &float_valid, + &histogram_valid, + &float_prefix, + &histogram_prefix, + ); + } + + let selections = select_mixed_range_lanes( + function, + &bounds, + timestamps, + &float_valid, + &histogram_valid, + &float_prefix, + &histogram_prefix, + &collector, + ); + let (lane, values, valid, prefix) = match output { + MixedRangeOutput::Float => ( + SampleLane::Float, + floats as &dyn Array, + float_valid.as_slice(), + float_prefix.as_slice(), + ), + MixedRangeOutput::Histogram => ( + SampleLane::Histogram, + histograms as &dyn Array, + histogram_valid.as_slice(), + histogram_prefix.as_slice(), + ), + }; + let input = compact_lane_input( + timestamps, + values, + valid, + prefix, + &bounds, + &selections, + lane, + name, + )?; + + let udf = match output { + MixedRangeOutput::Float => function.float_udf(), + MixedRangeOutput::Histogram => function.histogram_udf(collector), + } + .ok_or_else(|| { + DataFusionError::Execution(format!( + "{name} does not return {} values", + match output { + MixedRangeOutput::Float => "float", + MixedRangeOutput::Histogram => "native histogram", + } + )) + })?; + let mut input = input; + input.extend_from_slice(&args.args[4..]); + invoke_range_udf(udf, input, args, bounds.len()) +} + +/// Decodes packed range keys into validated half-open `(offset, end)` bounds. +fn checked_window_bounds( + keys: &[i64], + value_len: usize, + name: &str, +) -> DfResult> { + keys.iter() + .map(|key| { + let (offset, length) = unpack(*key); + let offset = offset as usize; + let end = offset + .checked_add(length as usize) + .filter(|end| *end <= value_len) + .ok_or_else(|| { + DataFusionError::Execution(format!( + "{name}: invalid range ({offset}, {length}) for {value_len} values" + )) + })?; + Ok((offset, end)) + }) + .collect() +} + +/// Builds prefix counts where `prefix[i]` is the number of valid samples in `valid[..i]`. +fn validity_prefix(valid: &[bool], name: &str) -> DfResult> { + let capacity = valid.len().checked_add(1).ok_or_else(|| { + DataFusionError::Execution(format!("{name}: sample validity length overflow")) + })?; + let mut prefix = Vec::with_capacity(capacity); + prefix.push(0usize); + for is_valid in valid { + prefix.push( + prefix + .last() + .copied() + .unwrap() + .checked_add(usize::from(*is_valid)) + .ok_or_else(|| { + DataFusionError::Execution(format!("{name}: sample count overflow")) + })?, + ); + } + Ok(prefix) +} + +/// Returns the number of valid samples in the half-open window `[offset, end)`. +fn valid_count(prefix: &[usize], offset: usize, end: usize) -> usize { + prefix[end] + .checked_sub(prefix[offset]) + .expect("validity prefix is monotonic") +} + +/// Selects the sample lane for each window and records policy-required annotations. +#[allow(clippy::too_many_arguments)] +fn select_mixed_range_lanes( + function: MixedRangeFunction, + bounds: &[(usize, usize)], + timestamps: &TimestampMillisecondArray, + float_valid: &[bool], + histogram_valid: &[bool], + float_prefix: &[usize], + histogram_prefix: &[usize], + collector: &Option, +) -> Vec> { + bounds + .iter() + .map(|(offset, end)| { + let float_count = valid_count(float_prefix, *offset, *end); + let histogram_count = valid_count(histogram_prefix, *offset, *end); + match function.policy() { + MixedRangePolicy::DropMixed => match (float_count > 0, histogram_count > 0) { + (true, true) => { + record_warning( + collector, + format!( + "{}: encountered a mix of float and native histogram samples", + function.name() + ), + ); + None + } + (true, false) => Some(SampleLane::Float), + (false, true) => Some(SampleLane::Histogram), + (false, false) => None, + }, + MixedRangePolicy::FloatOnly => { + if float_count > 0 && histogram_count > 0 { + record_info( + collector, + format!( + "{}: ignored native histogram samples", + function.name() + ), + ); + } + (float_count > 0).then_some(SampleLane::Float) + } + MixedRangePolicy::Last => (*offset..*end).rev().find_map(|row| { + if histogram_valid[row] { + Some(SampleLane::Histogram) + } else if float_valid[row] { + Some(SampleLane::Float) + } else { + None + } + }), + MixedRangePolicy::LastTwo => { + let mut last_two = Vec::with_capacity(2); + for row in (*offset..*end).rev() { + // Prometheus keeps a float as the newest sample if both lanes have the + // same timestamp, while the histogram becomes the preceding sample. + if float_valid[row] { + last_two.push((SampleLane::Float, timestamps.value(row))); + } + if last_two.len() < 2 && histogram_valid[row] { + last_two.push((SampleLane::Histogram, timestamps.value(row))); + } + if last_two.len() == 2 { + break; + } + } + if last_two.len() < 2 || last_two[0].1 == last_two[1].1 { + None + } else if last_two[0].0 == last_two[1].0 { + Some(last_two[0].0) + } else { + record_warning( + collector, + format!( + "{}: encountered a mix of float and native histogram samples in the last two points", + function.name() + ), + ); + None + } + } + MixedRangePolicy::Combined => unreachable!(), + } + }) + .collect() +} + +/// Filters null placeholders from one sample lane and remaps its selected windows. +/// Windows assigned to the other lane become empty. +#[allow(clippy::too_many_arguments)] +fn compact_lane_input( + timestamps: &TimestampMillisecondArray, + values: &dyn Array, + valid: &[bool], + prefix: &[usize], + bounds: &[(usize, usize)], + selections: &[Option], + lane: SampleLane, + name: &str, +) -> DfResult> { + let mask = BooleanArray::from(valid.to_vec()); + let filtered_timestamps = filter(timestamps, &mask)?; + let filtered_values = filter(values, &mask)?; + let ranges = bounds + .iter() + .zip(selections) + .map(|((offset, end), selected)| { + let compact_offset = prefix[*offset]; + let compact_length = if *selected == Some(lane) { + valid_count(prefix, *offset, *end) + } else { + 0 + }; + Ok(( + u32::try_from(compact_offset).map_err(|_| { + DataFusionError::Execution(format!( + "{name}: compacted range offset exceeds u32" + )) + })?, + u32::try_from(compact_length).map_err(|_| { + DataFusionError::Execution(format!( + "{name}: compacted range length exceeds u32" + )) + })?, + )) + }) + .collect::>>()?; + let timestamp_range = RangeArray::from_ranges(filtered_timestamps, ranges.clone()) + .map_err(DataFusionError::from)?; + let value_range = + RangeArray::from_ranges(filtered_values, ranges).map_err(DataFusionError::from)?; + Ok(vec![ + ColumnarValue::Array(Arc::new(timestamp_range.into_dict())), + ColumnarValue::Array(Arc::new(value_range.into_dict())), + ]) +} + +fn invoke_range_udf( + udf: ScalarUDF, + input: Vec, + outer_args: &ScalarFunctionArgs, + number_rows: usize, +) -> DfResult { + let arg_fields = input + .iter() + .enumerate() + .map(|(index, value)| Arc::new(Field::new(format!("arg_{index}"), value.data_type(), true))) + .collect(); + udf.invoke_with_args(ScalarFunctionArgs { + args: input, + arg_fields, + number_rows, + return_field: outer_args.return_field.clone(), + config_options: outer_args.config_options.clone(), + }) +} + +enum MixedSample { + Float(f64), + Histogram(NativeHistogram), +} + +#[allow(clippy::too_many_arguments)] +fn combined_range_float( + function: MixedRangeFunction, + bounds: &[(usize, usize)], + timestamps: &TimestampMillisecondArray, + floats: &Float64Array, + histograms: &StructArray, + float_valid: &[bool], + histogram_valid: &[bool], + float_prefix: &[usize], + histogram_prefix: &[usize], +) -> DfResult { + let mut result = Float64Builder::with_capacity(bounds.len()); + for (offset, end) in bounds { + let sample_count = valid_count(float_prefix, *offset, *end) + .checked_add(valid_count(histogram_prefix, *offset, *end)) + .ok_or_else(|| { + DataFusionError::Execution(format!("{}: sample count overflow", function.name())) + })?; + match function { + MixedRangeFunction::CountOverTime => { + if sample_count == 0 { + result.append_null(); + } else { + result.append_value(sample_count as f64); + } + } + MixedRangeFunction::AbsentOverTime => { + if sample_count == 0 { + result.append_value(1.0); + } else { + result.append_null(); + } + } + MixedRangeFunction::PresentOverTime => { + if sample_count == 0 { + result.append_null(); + } else { + result.append_value(1.0); + } + } + MixedRangeFunction::Changes | MixedRangeFunction::Resets => { + if sample_count == 0 { + result.append_null(); + continue; + } + let mut count = 0usize; + let mut previous = None; + for row in *offset..*end { + if float_valid[row] { + count += mixed_transition( + function, + &mut previous, + timestamps.value(row), + MixedSample::Float(floats.value(row)), + ); + } + if histogram_valid[row] { + let histogram = read_histogram(histograms, row)? + .expect("validated native histogram sample"); + count += mixed_transition( + function, + &mut previous, + timestamps.value(row), + MixedSample::Histogram(histogram), + ); + } + } + result.append_value(count as f64); + } + _ => { + return Err(DataFusionError::Internal(format!( + "{} does not support combined range evaluation", + function.name() + ))); + } + } + } + Ok(ColumnarValue::Array(Arc::new(result.finish()))) +} + +fn mixed_transition( + function: MixedRangeFunction, + previous: &mut Option<(i64, MixedSample)>, + timestamp: i64, + current: MixedSample, +) -> usize { + let changed = previous.as_ref().is_some_and(|(previous_ts, previous)| { + match (function, previous, ¤t) { + ( + MixedRangeFunction::Changes, + MixedSample::Float(previous), + MixedSample::Float(current), + ) => current != previous && !(current.is_nan() && previous.is_nan()), + ( + MixedRangeFunction::Changes, + MixedSample::Histogram(previous), + MixedSample::Histogram(current), + ) => !current.promql_eq(previous), + (MixedRangeFunction::Changes, _, _) => true, + ( + MixedRangeFunction::Resets, + MixedSample::Float(previous), + MixedSample::Float(current), + ) => current < previous, + ( + MixedRangeFunction::Resets, + MixedSample::Histogram(previous), + MixedSample::Histogram(current), + ) => { + (previous.reset_hint == GAUGE_RESET_HINT) + != (current.reset_hint == GAUGE_RESET_HINT) + || current.detect_counter_reset(previous, *previous_ts, timestamp) + } + (MixedRangeFunction::Resets, _, _) => true, + _ => unreachable!(), + } + }); + *previous = Some((timestamp, current)); + usize::from(changed) +} + fn native_histogram_scalar(histogram: Option) -> ScalarValue { let array = build_histogram_array(&[histogram]); let histogram = array @@ -2024,6 +2806,53 @@ mod tests { ] } + fn mixed_range_input( + name: &str, + floats: Vec>, + histograms: Vec>, + ) -> Vec { + assert_eq!(floats.len(), histograms.len()); + let timestamps = Arc::new(TimestampMillisecondArray::from_iter( + (0..floats.len()).map(|idx| Some((idx as i64 + 1) * 1000)), + )); + let floats = Arc::new(Float64Array::from(floats)); + let histograms = build_histogram_array(&histograms); + let range = [(0, u32::try_from(floats.len()).unwrap())]; + vec![ + ColumnarValue::Scalar(ScalarValue::Utf8(Some(name.to_string()))), + ColumnarValue::Array(Arc::new( + RangeArray::from_ranges(timestamps, range) + .unwrap() + .into_dict(), + )), + ColumnarValue::Array(Arc::new( + RangeArray::from_ranges(floats, range).unwrap().into_dict(), + )), + ColumnarValue::Array(Arc::new( + RangeArray::from_ranges(histograms, range) + .unwrap() + .into_dict(), + )), + ] + } + + fn mixed_float_result(udf: ScalarUDF, input: Vec) -> Option { + let result = run_udf(udf, input, DataType::Float64); + let result = extract_array(&result).unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + result.is_valid(0).then(|| result.value(0)) + } + + fn mixed_histogram_result( + udf: ScalarUDF, + input: Vec, + ) -> Option { + let result = run_udf(udf, input, native_histogram_arrow_type()); + let result = extract_array(&result).unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + read_histogram(result, 0).unwrap() + } + fn run_histogram_range_udf( udf: ScalarUDF, histograms: Vec, @@ -2098,6 +2927,170 @@ mod tests { warnings } + fn collected_infos(collector: &PromqlAnnotationCollector) -> Vec { + let mut infos = Vec::new(); + collector.append_to(&mut Vec::new(), &mut infos); + infos + } + + #[test] + fn mixed_ranges_follow_prometheus_sample_type_semantics() { + let first = sample_histogram(1.0, 1.0, vec![1.0]); + let second = sample_histogram(3.0, 3.0, vec![3.0]); + let collector = PromqlAnnotationCollector::default(); + + let mut rate = mixed_range_input( + "rate", + vec![Some(1.0), None, Some(3.0)], + vec![None, Some(first.clone()), None], + ); + rate.push(ColumnarValue::Array(Arc::new( + TimestampMillisecondArray::from(vec![3000]), + ))); + rate.push(ColumnarValue::Array(Arc::new(Int64Array::from(vec![3000])))); + assert_eq!( + mixed_float_result(MixedRange::float_udf(Some(collector.clone())), rate.clone()), + None + ); + assert_eq!( + mixed_histogram_result(MixedRange::histogram_udf(Some(collector.clone())), rate), + None + ); + assert!( + collected_warnings(&collector) + .iter() + .any(|warning| warning.contains("mix of float and native histogram")) + ); + + let pure_rate = |floats, histograms| { + let mut input = mixed_range_input("rate", floats, histograms); + input.push(ColumnarValue::Array(Arc::new( + TimestampMillisecondArray::from(vec![3000]), + ))); + input.push(ColumnarValue::Array(Arc::new(Int64Array::from(vec![3000])))); + input + }; + let pure_float = pure_rate( + vec![Some(1.0), Some(2.0), Some(3.0)], + vec![None, None, None], + ); + assert_eq!( + mixed_float_result(MixedRange::float_udf(None), pure_float.clone()), + Some(1.0) + ); + assert_eq!( + mixed_histogram_result(MixedRange::histogram_udf(None), pure_float), + None + ); + let pure_histogram = pure_rate( + vec![None, None, None], + vec![ + Some(sample_histogram(1.0, 1.0, vec![1.0])), + Some(sample_histogram(2.0, 2.0, vec![2.0])), + Some(sample_histogram(3.0, 3.0, vec![3.0])), + ], + ); + assert_eq!( + mixed_float_result(MixedRange::float_udf(None), pure_histogram.clone()), + None + ); + assert_eq!( + mixed_histogram_result(MixedRange::histogram_udf(None), pure_histogram) + .unwrap() + .count, + 1.0 + ); + + let idelta = mixed_range_input( + "idelta", + vec![Some(10.0), None, None], + vec![None, Some(first.clone()), Some(second.clone())], + ); + assert_eq!( + mixed_float_result(MixedRange::float_udf(None), idelta.clone()), + None + ); + assert_eq!( + mixed_histogram_result(MixedRange::histogram_udf(None), idelta) + .unwrap() + .count, + 2.0 + ); + + let alternating = || { + mixed_range_input( + "changes", + vec![Some(1.0), None, None, Some(1.0)], + vec![None, Some(first.clone()), Some(first.clone()), None], + ) + }; + assert_eq!( + mixed_float_result(MixedRange::float_udf(None), alternating()), + Some(2.0) + ); + let mut resets = alternating(); + resets[0] = ColumnarValue::Scalar(ScalarValue::Utf8(Some("resets".to_string()))); + assert_eq!( + mixed_float_result(MixedRange::float_udf(None), resets), + Some(2.0) + ); + + for (name, expected) in [ + ("count_over_time", Some(4.0)), + ("present_over_time", Some(1.0)), + ("absent_over_time", None), + ] { + let mut input = alternating(); + input[0] = ColumnarValue::Scalar(ScalarValue::Utf8(Some(name.to_string()))); + assert_eq!( + mixed_float_result(MixedRange::float_udf(None), input), + expected, + "{name}" + ); + } + + let last = mixed_range_input( + "last_over_time", + vec![None, Some(4.0)], + vec![Some(first.clone()), None], + ); + assert_eq!( + mixed_float_result(MixedRange::float_udf(None), last.clone()), + Some(4.0) + ); + assert_eq!( + mixed_histogram_result(MixedRange::histogram_udf(None), last), + None + ); + + let collector = PromqlAnnotationCollector::default(); + let histogram_only_min = + mixed_range_input("min_over_time", vec![None], vec![Some(first.clone())]); + assert_eq!( + mixed_float_result( + MixedRange::float_udf(Some(collector.clone())), + histogram_only_min, + ), + None + ); + assert!(collected_infos(&collector).is_empty()); + + let min = mixed_range_input( + "min_over_time", + vec![Some(3.0), None, Some(1.0)], + vec![None, Some(first), None], + ); + assert_eq!( + mixed_float_result(MixedRange::float_udf(Some(collector.clone())), min), + Some(1.0) + ); + assert!( + collected_infos(&collector) + .iter() + .any(|info| info.contains("ignored native histogram")) + ); + } + #[test] fn count_sum_and_avg_read_struct() { let histograms = vec![Some(sample_histogram(6.0, 10.0, vec![2.0, 4.0]))]; diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index a463f176c71..d21a90bf123 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -63,15 +63,15 @@ use promql::extension_plan::{ }; use promql::functions::{ AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, - IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, NativeHistogramAbsentOverTime, - NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges, NativeHistogramCount, - NativeHistogramCountOverTime, NativeHistogramDelta, NativeHistogramDrop, - NativeHistogramFraction, NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease, - NativeHistogramLastOverTime, NativeHistogramPresentOverTime, NativeHistogramQuantile, - NativeHistogramRate, NativeHistogramResets, NativeHistogramStddev, NativeHistogramStdvar, - NativeHistogramSum, NativeHistogramSumOverTime, PredictLinear, PresentOverTime, - QuantileOverTime, Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, - quantile_udaf, + IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, MixedRange, + NativeHistogramAbsentOverTime, NativeHistogramAvg, NativeHistogramAvgOverTime, + NativeHistogramChanges, NativeHistogramCount, NativeHistogramCountOverTime, + NativeHistogramDelta, NativeHistogramDrop, NativeHistogramFraction, NativeHistogramIDelta, + NativeHistogramIRate, NativeHistogramIncrease, NativeHistogramLastOverTime, + NativeHistogramPresentOverTime, NativeHistogramQuantile, NativeHistogramRate, + NativeHistogramResets, NativeHistogramStddev, NativeHistogramStdvar, NativeHistogramSum, + NativeHistogramSumOverTime, PredictLinear, PresentOverTime, QuantileOverTime, Rate, Resets, + Round, StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf, }; use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers}; use promql_parser::parser::token::TokenType; @@ -2833,6 +2833,90 @@ impl PromPlanner { Ok(result) } + fn create_mixed_range_function_exprs( + &mut self, + func: &Function, + mut other_input_exprs: VecDeque, + float_field: &str, + histogram_field: &str, + ) -> Result>> { + let returns_histogram = matches!( + func.name, + "rate" + | "increase" + | "delta" + | "idelta" + | "irate" + | "avg_over_time" + | "sum_over_time" + | "last_over_time" + ); + if !returns_histogram + && !matches!( + func.name, + "changes" + | "resets" + | "deriv" + | "min_over_time" + | "max_over_time" + | "count_over_time" + | "absent_over_time" + | "present_over_time" + | "stddev_over_time" + | "stdvar_over_time" + | "quantile_over_time" + | "predict_linear" + | "double_exponential_smoothing" + | "holt_winters" + ) + { + return Ok(None); + } + + if func.name == "predict_linear" { + other_input_exprs[0] = DfExpr::Cast(Cast { + expr: Box::new(other_input_exprs[0].clone()), + data_type: ArrowDataType::Int64, + }); + } + + let mut args = Vec::with_capacity(other_input_exprs.len() + 6); + args.push(lit(func.name)); + args.push(DfExpr::Column(Column::from_name( + RangeManipulate::build_timestamp_range_name( + self.ctx.time_index_column.as_ref().unwrap(), + ), + ))); + args.push(DfExpr::Column(Column::from_name(float_field))); + args.push(DfExpr::Column(Column::from_name(histogram_field))); + args.extend(other_input_exprs); + if matches!(func.name, "rate" | "increase" | "delta") { + args.push(self.create_time_index_column_expr()?); + args.push(lit(self.ctx.range.context(ExpectRangeSelectorSnafu)?)); + } + + let float_expr = DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(MixedRange::float_udf(self.promql_annotations.clone())), + args: args.clone(), + }); + let exprs = if returns_histogram { + self.ctx.field_columns = vec![float_field.to_string(), histogram_field.to_string()]; + vec![ + float_expr.alias(float_field), + DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(MixedRange::histogram_udf(self.promql_annotations.clone())), + args, + }) + .alias(histogram_field), + ] + } else { + let display_name = float_expr.schema_name().to_string(); + self.ctx.field_columns = vec![display_name.clone()]; + vec![float_expr.alias(display_name)] + }; + Ok(Some(exprs)) + } + /// Creates function expressions for projection and returns the expressions and new tags. /// /// # Side Effects @@ -2847,6 +2931,18 @@ impl PromPlanner { ) -> Result<(Vec, Vec)> { // TODO(ruihang): check function args list let mut other_input_exprs: VecDeque = other_input_exprs.into(); + if let Some((float_field, histogram_field)) = + Self::alternative_sample_range_columns(input_schema, &self.ctx.field_columns) + .map(|(float, histogram)| (float.to_string(), histogram.to_string())) + && let Some(exprs) = self.create_mixed_range_function_exprs( + func, + other_input_exprs.clone(), + &float_field, + &histogram_field, + )? + { + return Ok((exprs, vec![])); + } let alternative_samples = Self::field_columns_are_alternative_samples(input_schema, &self.ctx.field_columns); let all_field_columns_are_native_histogram_ranges = @@ -4444,6 +4540,17 @@ impl PromPlanner { .is_some_and(|data_type| data_type == &Self::native_histogram_arrow_type()) } + fn field_column_is_float_range(schema: &DFSchemaRef, field_column: &str) -> bool { + Self::field_column_type(schema, field_column).is_some_and(|data_type| { + matches!( + data_type, + ArrowDataType::Dictionary(key_type, value_type) + if key_type.as_ref() == &ArrowDataType::Int64 + && value_type.as_ref() == &ArrowDataType::Float64 + ) + }) + } + fn field_columns_are_alternative_samples( schema: &DFSchemaRef, field_columns: &[String], @@ -4461,11 +4568,13 @@ impl PromPlanner { let canonical_float = field_columns.iter().find(|field| { field.as_str() == greptime_value() - && Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64) + && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64) + || Self::field_column_is_float_range(schema, field)) }); let canonical_histogram = field_columns.iter().find(|field| { field.as_str() == greptime_native_histogram() - && Self::field_column_is_native_histogram(schema, field) + && (Self::field_column_is_native_histogram(schema, field) + || Self::field_column_is_native_histogram_range(schema, field)) }); if let (Some(float), Some(histogram)) = (canonical_float, canonical_histogram) { return Some((float, histogram)); @@ -4473,21 +4582,34 @@ impl PromPlanner { let float = field_columns.iter().find(|field| { field.starts_with(OR_FLOAT_FIELD_PREFIX) - && Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64) + && (Self::field_column_type(schema, field) == Some(&ArrowDataType::Float64) + || Self::field_column_is_float_range(schema, field)) })?; let histogram = field_columns.iter().find(|field| { field.starts_with(OR_HISTOGRAM_FIELD_PREFIX) - && Self::field_column_is_native_histogram(schema, field) + && (Self::field_column_is_native_histogram(schema, field) + || Self::field_column_is_native_histogram_range(schema, field)) })?; Some((float, histogram)) } + fn alternative_sample_range_columns<'a>( + schema: &DFSchemaRef, + field_columns: &'a [String], + ) -> Option<(&'a str, &'a str)> { + Self::alternative_sample_columns(schema, field_columns).filter(|(float, histogram)| { + Self::field_column_is_float_range(schema, float) + && Self::field_column_is_native_histogram_range(schema, histogram) + }) + } + fn field_column_is_native_histogram_range(schema: &DFSchemaRef, field_column: &str) -> bool { Self::field_column_type(schema, field_column).is_some_and(|data_type| { matches!( data_type, - ArrowDataType::Dictionary(_, value_type) - if value_type.as_ref() == &Self::native_histogram_arrow_type() + ArrowDataType::Dictionary(key_type, value_type) + if key_type.as_ref() == &ArrowDataType::Int64 + && value_type.as_ref() == &Self::native_histogram_arrow_type() ) }) } @@ -5617,7 +5739,10 @@ mod test { use catalog::memory::{MemoryCatalogManager, new_memory_catalog_manager}; use common_base::Plugins; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; - use common_query::prelude::{greptime_native_histogram, greptime_timestamp}; + use common_query::native_histogram::{ + CounterResetHint, NativeHistogram, build_histogram_array, + }; + use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; use common_query::test_util::DummyDecoder; use datafusion::arrow::array::{ Array, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, @@ -5791,6 +5916,23 @@ mod test { } } + fn direct_or_histogram() -> NativeHistogram { + NativeHistogram { + schema: 0, + zero_threshold: 0.0, + sum: 1.0, + reset_hint: CounterResetHint::Unknown, + start_timestamp: None, + custom_values: vec![], + positive_spans: vec![], + negative_spans: vec![], + count: 1.0, + zero_count: 1.0, + positive_buckets: vec![], + negative_buckets: vec![], + } + } + struct DirectOrSource { name: &'static str, empty: bool, @@ -8921,6 +9063,175 @@ mod test { } } + #[tokio::test] + async fn mixed_native_histogram_ranges_use_coordinated_udfs() { + let dual_output = [ + "increase(some_metric[5m])", + "rate(some_metric[5m])", + "delta(some_metric[5m])", + "idelta(some_metric[5m])", + "irate(some_metric[5m])", + "avg_over_time(some_metric[5m])", + "sum_over_time(some_metric[5m])", + "last_over_time(some_metric[5m])", + ]; + let float_output = [ + "resets(some_metric[5m])", + "changes(some_metric[5m])", + "deriv(some_metric[5m])", + "min_over_time(some_metric[5m])", + "max_over_time(some_metric[5m])", + "count_over_time(some_metric[5m])", + "absent_over_time(some_metric[5m])", + "present_over_time(some_metric[5m])", + "stddev_over_time(some_metric[5m])", + "stdvar_over_time(some_metric[5m])", + "quantile_over_time(0.9, some_metric[5m])", + "predict_linear(some_metric[5m], 60)", + "double_exponential_smoothing(some_metric[5m], 0.5, 0.5)", + ]; + + for query in dual_output.iter().chain(float_output.iter()) { + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &build_eval_stmt(query), + &build_query_engine_state(), + ) + .await + .unwrap() + .display_indent_schema() + .to_string(); + assert!(plan.contains("prom_mixed_range_float"), "{query}\n{plan}"); + assert_eq!( + plan.contains("prom_mixed_range_histogram"), + dual_output.contains(query), + "{query}\n{plan}" + ); + } + + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &build_eval_stmt("sum_over_time(rate(some_metric[5m])[10m:1m])"), + &build_query_engine_state(), + ) + .await + .unwrap() + .display_indent_schema() + .to_string(); + let expected = r#"Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_value, prom_mixed_range_histogram(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=["greptime_value", "greptime_native_histogram"] [timestamp:Timestamp(ms), greptime_value:Dictionary(Int64, Float64);N, greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, tag_0:Utf8, timestamp_range:Dictionary(Int64, Timestamp(ms))] + PromSeriesDivide: tags=["tag_0"] [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_value, prom_mixed_range_histogram(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + PromRangeManipulate: req range=[-540000..100000000], interval=[60000], eval range=[300000], time index=[timestamp], values=["greptime_native_histogram", "greptime_value"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, greptime_value:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))] + PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + Filter: some_metric.timestamp >= TimestampMillisecond(-839999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]"#; + assert_eq!(plan, expected); + } + + #[tokio::test] + async fn mixed_native_histogram_rate_executes_real_ranges() { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "timestamp", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + ), + Field::new(greptime_value(), ArrowDataType::Float64, true), + Field::new( + greptime_native_histogram(), + native_histogram_value_type().as_arrow_type(), + true, + ), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1000, 2000, 3000])), + Arc::new(Float64Array::from(vec![Some(1.0), None, Some(3.0)])), + build_histogram_array(&[None, Some(direct_or_histogram()), None]), + ], + ) + .unwrap(); + let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); + let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) + .unwrap() + .build() + .unwrap(); + let collector = PromqlAnnotationCollector::default(); + let mut planner = PromPlanner { + table_provider: build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await, + ctx: PromPlannerContext { + start: 3000, + end: 3000, + interval: 1000, + range: Some(3000), + time_index_column: Some("timestamp".to_string()), + field_columns: vec![ + greptime_native_histogram().to_string(), + greptime_value().to_string(), + ], + ..Default::default() + }, + promql_annotations: Some(collector.clone()), + }; + let input = LogicalPlan::Extension(Extension { + node: Arc::new( + RangeManipulate::new( + 3000, + 3000, + 1000, + 3000, + "timestamp".to_string(), + planner.ctx.field_columns.clone(), + input, + ) + .unwrap(), + ), + }); + let PromExpr::Call(call) = parser::parse("rate(mixed[3s])").unwrap() else { + unreachable!() + }; + let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( + input.schema(), + &planner.ctx.field_columns, + ); + let state = build_query_engine_state(); + let (mut exprs, _) = planner + .create_function_expr(&call.func, vec![], input.schema(), &state) + .unwrap(); + exprs.insert(0, planner.create_time_index_column_expr().unwrap()); + let plan = LogicalPlanBuilder::from(input) + .project(exprs) + .unwrap() + .filter( + planner + .create_empty_values_filter_expr(preserve_any_value) + .unwrap(), + ) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + let mut warnings = Vec::new(); + collector.append_to(&mut warnings, &mut Vec::new()); + assert!( + warnings + .iter() + .any(|warning| warning.contains("mix of float and native histogram")) + ); + } + #[tokio::test] async fn native_histogram_mixed_field_table_behaves() { // Metric Engine exposes float and histogram samples as alternative nullable fields. diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index ca4bdd81044..030ed27b4e3 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -30,18 +30,18 @@ use datafusion_expr::UserDefinedLogicalNode; use greptime_proto::substrait_extension::MergeScan as PbMergeScan; use promql::functions::{ AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, - IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, NativeHistogramAbsentOverTime, - NativeHistogramAdd, NativeHistogramAggAvg, NativeHistogramAggSum, NativeHistogramAvg, - NativeHistogramAvgOverTime, NativeHistogramChanges, NativeHistogramCount, - NativeHistogramCountOverTime, NativeHistogramDelta, NativeHistogramDivScalar, - NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, NativeHistogramIDelta, - NativeHistogramIRate, NativeHistogramIncrease, NativeHistogramLastOverTime, - NativeHistogramMulScalar, NativeHistogramNeg, NativeHistogramNotEq, - NativeHistogramPresentOverTime, NativeHistogramQuantile, NativeHistogramRate, - NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar, - NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, NativeHistogramToString, - PredictLinear, PresentOverTime, QuantileOverTime, Rate, Resets, Round, StddevOverTime, - StdvarOverTime, SumOverTime, quantile_udaf, + IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, MixedRange, + NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg, + NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges, + NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta, + NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, + NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease, + NativeHistogramLastOverTime, NativeHistogramMulScalar, NativeHistogramNeg, + NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile, + NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, + NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, + NativeHistogramToString, PredictLinear, PresentOverTime, QuantileOverTime, Rate, Resets, Round, + StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf, }; use prost::Message; use session::context::QueryContextRef; @@ -195,6 +195,8 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { NativeHistogramIRate::scalar_udf(), NativeHistogramIncrease::scalar_udf(), NativeHistogramLastOverTime::scalar_udf(), + MixedRange::float_udf(None), + MixedRange::histogram_udf(None), NativeHistogramMulScalar::scalar_udf(), NativeHistogramNeg::scalar_udf(), NativeHistogramNotEq::scalar_udf(), @@ -240,10 +242,14 @@ mod tests { use common_query::native_histogram::native_histogram_value_type; use datafusion::catalog::TableProvider; use datafusion::datasource::MemTable; + use datafusion::logical_expr::Extension; use datafusion_expr::expr::ScalarFunction; use datafusion_expr::{Expr, LogicalPlanBuilder, LogicalTableSource, col, lit}; - use datatypes::arrow::datatypes::{Field, Schema, SchemaRef}; + use datatypes::arrow::datatypes::{ + DataType as ArrowDataType, Field, Schema, SchemaRef, TimeUnit, + }; use datatypes::data_type::DataType; + use promql::extension_plan::RangeManipulate; use session::context::QueryContext; use super::*; @@ -374,5 +380,82 @@ mod tests { assert!(decoded.contains("prom_native_histogram_count")); assert!(decoded.contains("prom_native_histogram_drop_bool")); assert!(decoded.contains("prom_native_histogram_drop_float")); + + let schema = Arc::new(Schema::new(vec![ + Field::new( + "timestamp", + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("float", ArrowDataType::Float64, true), + Field::new( + "histogram", + native_histogram_value_type().as_arrow_type(), + true, + ), + ])); + let input = LogicalPlanBuilder::scan( + "devices", + Arc::new(LogicalTableSource::new(schema.clone())), + None, + ) + .unwrap() + .build() + .unwrap(); + let input = LogicalPlan::Extension(Extension { + node: Arc::new( + RangeManipulate::new( + 0, + 1000, + 1000, + 1000, + "timestamp".to_string(), + vec!["float".to_string(), "histogram".to_string()], + input, + ) + .unwrap(), + ), + }); + let plan = LogicalPlanBuilder::from(input) + .project(vec![ + Expr::ScalarFunction(ScalarFunction { + func: Arc::new(MixedRange::float_udf(None)), + args: vec![ + lit("last_over_time"), + col("timestamp_range"), + col("float"), + col("histogram"), + ], + }) + .alias("mixed_float"), + Expr::ScalarFunction(ScalarFunction { + func: Arc::new(MixedRange::histogram_udf(None)), + args: vec![ + lit("last_over_time"), + col("timestamp_range"), + col("float"), + col("histogram"), + ], + }) + .alias("mixed_histogram"), + ]) + .unwrap() + .build() + .unwrap(); + let bytes = DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .unwrap(); + let table_provider = Arc::new(MemTable::try_new(schema, vec![vec![]]).unwrap()); + let decoded = plan_decoder + .decode( + bytes, + Arc::new(DummyCatalogList::with_table_provider(table_provider)), + false, + ) + .await + .unwrap() + .to_string(); + assert!(decoded.contains("prom_mixed_range_float")); + assert!(decoded.contains("prom_mixed_range_histogram")); } } From 584d66e3c3c2ae878b6e7bc71910fda6db12b5ca Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 10 Aug 2026 17:49:35 +0800 Subject: [PATCH 09/92] fix(meta): validate leaders before enabling skip_wal (#8823) Signed-off-by: jeremyhi --- src/common/meta/src/ddl/alter_table.rs | 11 +++- src/common/meta/src/ddl/tests/alter_table.rs | 60 ++++++++++++++++++++ 2 files changed, 68 insertions(+), 3 deletions(-) diff --git a/src/common/meta/src/ddl/alter_table.rs b/src/common/meta/src/ddl/alter_table.rs index 9504e8672ca..acc02c4d1a8 100644 --- a/src/common/meta/src/ddl/alter_table.rs +++ b/src/common/meta/src/ddl/alter_table.rs @@ -198,6 +198,12 @@ impl AlterTableProcedure { // DDL manager can submit another procedure with locks for the current route. return Ok(Status::done_with_output(RegionRouteChanged)); } + ensure!( + !find_leaders(&physical_table_route.region_routes).is_empty(), + NoLeaderSnafu { + table_id: physical_table_id + } + ); self.data.region_distribution = Some(region_distribution(&physical_table_route.region_routes)); } @@ -264,9 +270,8 @@ impl AlterTableProcedure { MultipleResults::PartialRetryable(error) => Err(error), MultipleResults::PartialNonRetryable(error) | MultipleResults::AllNonRetryable(error) => { - // The metadata already enables skip-WAL. Keep retrying until every - // replica applies the runtime option instead of leaving the table in - // a permanently inconsistent state. + // The metadata already enables skip-WAL. Retry the idempotent request + // so later attempts can update the remaining replicas. Err(BoxedError::new(error)).context(RetryLaterSnafu { clean_poisons: true, }) diff --git a/src/common/meta/src/ddl/tests/alter_table.rs b/src/common/meta/src/ddl/tests/alter_table.rs index af72aaa65df..d68d1c96dfa 100644 --- a/src/common/meta/src/ddl/tests/alter_table.rs +++ b/src/common/meta/src/ddl/tests/alter_table.rs @@ -850,6 +850,66 @@ async fn test_skip_wal_detects_region_route_change() { assert!(!table_info.meta.options.skip_wal); } +#[tokio::test] +async fn test_skip_wal_rejects_no_leader_before_updating_metadata() { + let ddl_context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + let table_name = "foo"; + let table_id = 1024; + let task = test_create_table_task(table_name, table_id); + let mut region_routes = prepare_table_route(table_id) + .region_routes() + .unwrap() + .clone(); + for route in &mut region_routes { + route.leader_peer = None; + } + let region_locks = region_routes.iter().map(|route| route.region.id).collect(); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info, + TableRouteValue::physical(region_routes), + HashMap::new(), + ) + .await + .unwrap(); + + let alter_task = AlterTableTask { + alter_table: AlterTableExpr { + catalog_name: DEFAULT_CATALOG_NAME.to_string(), + schema_name: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + kind: Some(Kind::SetTableOptions(SetTableOptions { + table_options: vec![api::v1::Option { + key: SKIP_WAL_KEY.to_string(), + value: "true".to_string(), + }], + })), + }, + }; + let mut procedure = AlterTableProcedure::new_with_region_locks( + table_id, + alter_task, + EventContext::default(), + region_locks, + ddl_context.clone(), + ) + .unwrap(); + + let error = procedure.on_prepare().await.unwrap_err(); + assert_matches!(error, Error::NoLeader { .. }); + let table_info = ddl_context + .table_metadata_manager + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap() + .into_inner() + .table_info; + assert!(!table_info.meta.options.skip_wal); +} + #[tokio::test] async fn test_skip_wal_updates_metadata_before_all_replicas() { let (tx, mut rx) = mpsc::channel(8); From 60c7bbcaf30a0e65ad7fa2b2b8ad0c1199a662b1 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 10 Aug 2026 18:47:50 +0800 Subject: [PATCH 10/92] refactor: share query channel definition (#8825) refactor: share query channel protocol mapping Signed-off-by: WenyXu --- Cargo.lock | 3 +- src/common/base/Cargo.toml | 1 + src/common/base/src/lib.rs | 1 + src/common/base/src/protocol.rs | 106 ++++++++++++++++++++++++++++++++ src/common/meta/Cargo.toml | 1 - src/common/meta/src/rpc/ddl.rs | 16 ++--- src/common/session/src/lib.rs | 49 +-------------- src/session/Cargo.toml | 1 + src/session/src/context.rs | 71 +++------------------ src/session/src/lib.rs | 4 +- 10 files changed, 126 insertions(+), 127 deletions(-) create mode 100644 src/common/base/src/protocol.rs diff --git a/Cargo.lock b/Cargo.lock index 4ba0a7f7e27..e2210b80bfe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2292,6 +2292,7 @@ dependencies = [ "pin-project", "regex", "serde", + "strum 0.27.1", "tokio", "toml 0.8.23", "zeroize", @@ -2633,7 +2634,6 @@ dependencies = [ "common-query", "common-recordbatch", "common-runtime", - "common-session", "common-telemetry", "common-test-util", "common-time", @@ -13426,6 +13426,7 @@ dependencies = [ "api", "arc-swap", "auth", + "common-base", "common-catalog", "common-error", "common-macro", diff --git a/src/common/base/Cargo.toml b/src/common/base/Cargo.toml index 44c30cd548d..760241a3e6a 100644 --- a/src/common/base/Cargo.toml +++ b/src/common/base/Cargo.toml @@ -21,6 +21,7 @@ lazy_static.workspace = true pin-project.workspace = true regex.workspace = true serde = { version = "1.0", features = ["derive"] } +strum.workspace = true tokio.workspace = true zeroize = { version = "1.6", default-features = false, features = ["alloc"] } diff --git a/src/common/base/src/lib.rs b/src/common/base/src/lib.rs index 18e6ffb3e0f..04fc7038018 100644 --- a/src/common/base/src/lib.rs +++ b/src/common/base/src/lib.rs @@ -18,6 +18,7 @@ pub mod cancellation; pub mod hash; pub mod memory_limit; pub mod plugins; +pub mod protocol; pub mod range_read; #[allow(clippy::all)] pub mod readable_size; diff --git a/src/common/base/src/protocol.rs b/src/common/base/src/protocol.rs new file mode 100644 index 00000000000..66704dbc18b --- /dev/null +++ b/src/common/base/src/protocol.rs @@ -0,0 +1,106 @@ +// Copyright 2023 Greptime Team +// +// 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. + +use std::fmt::{Display, Formatter}; + +/// The protocol through which a query is received. +#[derive(Debug, PartialEq, Default, Clone, Copy, strum::FromRepr)] +#[repr(u8)] +pub enum Channel { + #[default] + Unknown = 0, + + Mysql = 1, + Postgres = 2, + HttpSql = 3, + Prometheus = 4, + Otlp = 5, + Grpc = 6, + Influx = 7, + Opentsdb = 8, + Loki = 9, + Elasticsearch = 10, + Jaeger = 11, + Log = 12, + Promql = 13, + Splunk = 14, +} + +impl From for Channel { + fn from(value: u32) -> Self { + u8::try_from(value) + .ok() + .and_then(Self::from_repr) + .unwrap_or_default() + } +} + +impl Display for Channel { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.as_ref()) + } +} + +impl AsRef for Channel { + fn as_ref(&self) -> &str { + match self { + Self::Unknown => "unknown", + Self::Mysql => "mysql", + Self::Postgres => "postgres", + Self::HttpSql => "httpsql", + Self::Prometheus => "prometheus", + Self::Otlp => "otlp", + Self::Grpc => "grpc", + Self::Influx => "influx", + Self::Opentsdb => "opentsdb", + Self::Loki => "loki", + Self::Elasticsearch => "elasticsearch", + Self::Jaeger => "jaeger", + Self::Log => "log", + Self::Promql => "promql", + Self::Splunk => "splunk", + } + } +} + +#[cfg(test)] +mod tests { + use super::Channel; + + #[test] + fn test_channel_name() { + let expected = [ + (1, "mysql"), + (2, "postgres"), + (3, "httpsql"), + (4, "prometheus"), + (5, "otlp"), + (6, "grpc"), + (7, "influx"), + (8, "opentsdb"), + (9, "loki"), + (10, "elasticsearch"), + (11, "jaeger"), + (12, "log"), + (13, "promql"), + (14, "splunk"), + ]; + + for (value, name) in expected { + assert_eq!(name, Channel::from(value).as_ref()); + } + assert_eq!("unknown", Channel::from(0).as_ref()); + assert_eq!("unknown", Channel::from(15).as_ref()); + } +} diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index 43714b7d599..8e0cbbdf354 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -41,7 +41,6 @@ common-procedure.workspace = true common-query.workspace = true common-recordbatch.workspace = true common-runtime.workspace = true -common-session.workspace = true common-telemetry.workspace = true common-time.workspace = true common-wal.workspace = true diff --git a/src/common/meta/src/rpc/ddl.rs b/src/common/meta/src/rpc/ddl.rs index 9cb75aa29c0..8aef3325e64 100644 --- a/src/common/meta/src/rpc/ddl.rs +++ b/src/common/meta/src/rpc/ddl.rs @@ -42,9 +42,9 @@ use api::v1::{ }; use base64::Engine as _; use base64::engine::general_purpose; +use common_base::protocol::Channel; use common_catalog::{format_full_flow_name, format_full_table_name}; use common_error::ext::BoxedError; -use common_session::channel_protocol; use common_time::{DatabaseTimeToLive, Timestamp}; use prost::Message; use serde::{Deserialize, Serialize}; @@ -1668,13 +1668,6 @@ pub struct QueryContext { pub sst_min_sequences: HashMap, } -impl QueryContext { - /// Returns the protocol name represented by the wire channel value. - pub fn channel_protocol(&self) -> Option<&'static str> { - channel_protocol(self.channel) - } -} - /// The stable context recorded for a procedure event. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct EventContext { @@ -1708,11 +1701,12 @@ impl EventContext { .get(TRIGGER_REASON_EXTENSION_KEY) .map(|reason| TriggerReason::from_extension(reason)) .unwrap_or_default(); + let channel = Channel::from(u32::from(query_context.channel)); let context = Self::new(reason); - if let Some(protocol) = query_context.channel_protocol() { - context.with_protocol(protocol) - } else { + if channel == Channel::Unknown { context + } else { + context.with_protocol(channel.as_ref()) } } } diff --git a/src/common/session/src/lib.rs b/src/common/session/src/lib.rs index 418425df15e..f09baf92866 100644 --- a/src/common/session/src/lib.rs +++ b/src/common/session/src/lib.rs @@ -15,27 +15,6 @@ use serde::{Deserialize, Serialize}; use strum::{AsRefStr, Display, EnumString}; -/// Returns the protocol name represented by a wire query channel value. -pub fn channel_protocol(channel: u8) -> Option<&'static str> { - match channel { - 1 => Some("mysql"), - 2 => Some("postgres"), - 3 => Some("httpsql"), - 4 => Some("prometheus"), - 5 => Some("otlp"), - 6 => Some("grpc"), - 7 => Some("influx"), - 8 => Some("opentsdb"), - 9 => Some("loki"), - 10 => Some("elasticsearch"), - 11 => Some("jaeger"), - 12 => Some("log"), - 13 => Some("promql"), - 14 => Some("splunk"), - _ => None, - } -} - /// Defines the read preference for frontend route operations, /// determining whether to read from the region leader or follower. #[derive( @@ -52,33 +31,7 @@ pub enum ReadPreference { mod tests { use std::str::FromStr; - use crate::{ReadPreference, channel_protocol}; - - #[test] - fn test_channel_protocol() { - let expected = [ - (1, "mysql"), - (2, "postgres"), - (3, "httpsql"), - (4, "prometheus"), - (5, "otlp"), - (6, "grpc"), - (7, "influx"), - (8, "opentsdb"), - (9, "loki"), - (10, "elasticsearch"), - (11, "jaeger"), - (12, "log"), - (13, "promql"), - (14, "splunk"), - ]; - - for (channel, protocol) in expected { - assert_eq!(Some(protocol), channel_protocol(channel)); - } - assert_eq!(None, channel_protocol(0)); - assert_eq!(None, channel_protocol(15)); - } + use crate::ReadPreference; #[test] fn test_read_preference() { diff --git a/src/session/Cargo.toml b/src/session/Cargo.toml index c7be8c2f7e8..889b911c3cc 100644 --- a/src/session/Cargo.toml +++ b/src/session/Cargo.toml @@ -15,6 +15,7 @@ ahash.workspace = true api.workspace = true arc-swap = "1.5" auth.workspace = true +common-base.workspace = true common-catalog.workspace = true common-error.workspace = true common-macro.workspace = true diff --git a/src/session/src/context.rs b/src/session/src/context.rs index 4780523e4b7..d383aef477e 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -22,6 +22,7 @@ use api::v1::ExplainOptions; use api::v1::region::RegionRequestHeader; use arc_swap::ArcSwap; use auth::UserInfoRef; +pub use common_base::protocol::Channel; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_catalog::{build_db_string, parse_catalog_and_schema_from_db_string}; use common_recordbatch::cursor::RecordBatchStreamCursor; @@ -618,70 +619,12 @@ impl ConnInfo { } } -/// The protocol mapping for these wire values lives in [`common_session::channel_protocol`]. -#[derive(Debug, PartialEq, Default, Clone, Copy)] -#[repr(u8)] -pub enum Channel { - #[default] - Unknown = 0, - - Mysql = 1, - Postgres = 2, - HttpSql = 3, - Prometheus = 4, - Otlp = 5, - Grpc = 6, - Influx = 7, - Opentsdb = 8, - Loki = 9, - Elasticsearch = 10, - Jaeger = 11, - Log = 12, - Promql = 13, - Splunk = 14, -} - -impl From for Channel { - fn from(value: u32) -> Self { - match value { - 1 => Self::Mysql, - 2 => Self::Postgres, - 3 => Self::HttpSql, - 4 => Self::Prometheus, - 5 => Self::Otlp, - 6 => Self::Grpc, - 7 => Self::Influx, - 8 => Self::Opentsdb, - 9 => Self::Loki, - 10 => Self::Elasticsearch, - 11 => Self::Jaeger, - 12 => Self::Log, - 13 => Self::Promql, - 14 => Self::Splunk, - _ => Self::Unknown, - } - } -} - -impl Channel { - pub fn dialect(&self) -> Arc { - match self { - Channel::Mysql => Arc::new(MySqlDialect {}), - Channel::Postgres => Arc::new(PostgreSqlDialect {}), - _ => Arc::new(GenericDialect {}), - } - } -} - -impl Display for Channel { - fn fmt(&self, f: &mut Formatter) -> std::fmt::Result { - write!(f, "{}", self.as_ref()) - } -} - -impl AsRef for Channel { - fn as_ref(&self) -> &str { - common_session::channel_protocol(*self as u8).unwrap_or("unknown") +/// Returns the SQL dialect for the given query channel. +pub fn dialect_for_channel(channel: Channel) -> Arc { + match channel { + Channel::Mysql => Arc::new(MySqlDialect {}), + Channel::Postgres => Arc::new(PostgreSqlDialect {}), + _ => Arc::new(GenericDialect {}), } } diff --git a/src/session/src/lib.rs b/src/session/src/lib.rs index 81ee7b0cdc2..110560b3243 100644 --- a/src/session/src/lib.rs +++ b/src/session/src/lib.rs @@ -34,7 +34,7 @@ use common_time::timezone::get_timezone; use context::{ConfigurationVariables, QueryContextBuilder}; use derive_more::Debug; -use crate::context::{Channel, ConnInfo, QueryContextRef}; +use crate::context::{Channel, ConnInfo, QueryContextRef, dialect_for_channel}; /// Maximum number of warnings to store per session (similar to MySQL's max_error_count) const MAX_WARNINGS: usize = 64; @@ -102,7 +102,7 @@ impl Session { // string here .current_catalog(self.catalog.read().unwrap().clone()) .mutable_session_data(self.mutable_inner.clone()) - .sql_dialect(self.conn_info.channel.dialect()) + .sql_dialect(dialect_for_channel(self.conn_info.channel)) .configuration_parameter(self.configuration_variables.clone()) .channel(self.conn_info.channel) .process_id(self.process_id) From 335a95a369c728c355f1c8452d8688525beb24b9 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 10 Aug 2026 20:18:28 +0800 Subject: [PATCH 11/92] feat: declared edges and the derivation contract for the entity graph (#8794) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(frontend): run entity-graph derivation as the caller The derivation contract requires the computed graph tables to run under the outer query's identity. Capture the caller's QueryContext when the computed table is resolved, thread it through EntityGraphProvider, and: - authorize every contributing source table against the caller via the new semantic_graph.query permission action, silently excluding denied sources (entities, edges and source_tables never appear); - execute the derivation plan under the caller's context so it inherits permissions, cancellation and deadline instead of a fresh default. Signed-off-by: Dennis Zhuang * feat(operator): derive the entity-graph window from the scan's time predicate Implements the RFC window contract for the computed graph tables: - table: add extract_time_range_strict, a strict variant of the lenient time-range extraction that distinguishes an absent observed_at filter from one that cannot be safely turned into a range; - operator: replace GraphWindow with GraphQueryWindow, splitting the queried observed_at range from the source-scan range widened to whole 60s buckets, so boundary buckets aggregate over their full extent; - frontend: resolve the window from ScanRequest filters — no predicate keeps the last-hour default, a missing upper bound means now, and a missing lower bound or unextractable shape is an explicit error, never a silent fallback. Signed-off-by: Dennis Zhuang * feat(operator): system-defined declared-edge table for the entity graph Reintroduces greptime_private.semantic_relationships_declared with a canonical, system-owned definition: - the CREATE TABLE expr (8-tag primary key, business validity columns, RED fields, 30d TTL); attributes is now a json column so the future union branch matches the computed table without a per-scan parse; - created on first use on every write path: SQL INSERT creates it before executing, and the gRPC row-insert auto-create substitutes the canonical expr instead of deriving a schema from the request; - user DDL (CREATE/ALTER/DROP/RENAME/TRUNCATE) and write-path auto-ALTER are rejected via the new is_ddl_reserved_table guard, while INSERT/DELETE stay allowed. Signed-off-by: Dennis Zhuang * feat(operator): union declared edges into semantic_relationships Adds the declared-edge branch to the relationship derivation (build_relationships_plan replaces build_calls_plan): - latest revision per edge key first (mito dedups on primary key plus observed_at, so a re-asserted edge stores a new revision), then the business-validity overlap against the queried window; valid_from defaults to the declaration time and a NULL valid_until means the edge holds while its row exists; - the projected observed_at is synthesized inside the queried range (Inexact pushdown re-applies the scan's filters above the computed table, which would drop rows keyed by the physical revision time); window_end/fresh_until of open-ended edges take the window's upper bound so 'fresh_until >= now() - ...' queries see them; - tag columns are cast out of dictionary encoding, and the union is re-projected to the 16-column contract; - the frontend feeds the branch only when the physical table exists, the caller may read it, and its schema still matches the canonical definition (mismatch is an explicit error, not a silent drop). Signed-off-by: Dennis Zhuang * test: cover declared edges, window contract and caller authorization - sqlness: system auto-create on first INSERT, latest-revision reads, open-ended vs retired validity, explicit/lower-only/upper-only window behavior, user-DDL rejection, rename-into rejection, DELETE cleanup; - integration: a permission checker denying one trace table excludes it from both semantic_relationships and semantic_entities. Signed-off-by: Dennis Zhuang * fix: allow DROP/TRUNCATE on the declared-edge table and fix CI lints The definition guard rejected every DDL, which left sqlness (and any shared deployment) no way to remove the table the semantic_graph case creates — its extra region then broke unrelated region/partition case expectations. Narrow the guard to what actually protects the canonical definition: user CREATE, ALTER, RENAME-into and repartition stay rejected, while DROP and TRUNCATE are allowed — dropping loses nothing structural, the next INSERT recreates the table canonically, and DROP doubles as the recovery path if the canonical definition ever changes. The sqlness case now verifies drop-then-recreate and cleans up after itself. Also: rustfmt for the catalog crate and two typo fixes. Signed-off-by: Dennis Zhuang * chore: adapt canonical declared-table create to TriggerReason Signed-off-by: Dennis Zhuang * fix: address review on the declared-edge table lifecycle and revision reads - gRPC first writes actually work now: the reserved table's creation went through the generic create_table_inner, which the definition guard itself rejects; both branches of create_or_alter_tables_on_demand route it to create_declared_relationships_table instead, and being a system action it also bypasses the auto_create_table config/hint; - revision selection is as-of the queried window: revisions recorded after the window's end, or whose validity starts after it, no longer outrank (and hide) the revision that was in effect inside it; - the canonical-schema check validates the whole definition the union semantics lean on — time index, primary key, engine, append/merge mode — not just column names and types; - UNDROP TABLE of the reserved name is rejected like CREATE: it could resurrect a pre-canonical shape, and the next INSERT recreates the table anyway. Signed-off-by: Dennis Zhuang * chore: trim over-commenting in the entity-graph code Comments that restated adjacent code or narrated justification are cut; the ones stating non-obvious contracts and gotchas stay. Signed-off-by: Dennis Zhuang * fix: reject CREATE VIEW against DDL-reserved table names A view named greptime_private.semantic_relationships_declared would squat the reserved name: the first INSERT then skips the canonical create (an object already exists) and graph reads fail on the schema mismatch. CREATE VIEW now passes the same definition guard as CREATE TABLE. Signed-off-by: Dennis Zhuang * feat: debug-log authorization exclusions; declared-edge TTL to 90d Sources the derivation contract silently excludes (per-table denial, whole-scan denial, the declared-edge table) are invisible from outside; a debug log at each names what was excluded and why. The declared-edge table's default TTL becomes 90d, overridable at creation time via GREPTIMEDB_DECLARED_RELATIONSHIPS_TTL (a proper configuration option is a TODO). Signed-off-by: Dennis Zhuang * fix: rank declared-edge revisions by the visible edge identity Ranking partitioned by the full primary key, but the projection drops scope and generation_id: two assertions of the same visible edge under different generations both ranked first and came out as duplicate, indistinguishable rows. Rank by the exposed identity (endpoints, rel_type, provenance) instead, with generation_id/scope as deterministic tie-breakers for same-timestamp assertions. Signed-off-by: Dennis Zhuang * test: drop redundant declared-edge tests The generations regression is already asserted by the revision and as-of tests; the DDL shape test restated the declarative builder against itself. Its one non-tautological check (attributes maps to the json type) moves into the schema-matcher test. Signed-off-by: Dennis Zhuang * fix: reject disjunctive graph windows and unmatchable future windows - OR/IN over observed_at collapse disjoint ranges into their convex hull; a declared edge's synthesized timestamp can land in a gap and be dropped by the re-applied filter even though the edge is valid at a requested instant. The strict extractor now rejects those shapes. - A lower bound in the future inverts against the implicit up-to-now upper bound; the declared branch then fabricated an edge observed at the future bound. Such windows now derive nothing. - The reserved-table gRPC create path classifies an instant-TTL table like every sibling path. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/auth/src/lib.rs | 2 +- src/auth/src/permission.rs | 5 + src/catalog/src/kvbackend/manager.rs | 8 +- .../src/system_schema/semantic_graph.rs | 57 +- src/common/catalog/src/consts.rs | 11 + src/frontend/src/instance/builder.rs | 2 + src/frontend/src/instance/entity_graph.rs | 269 ++++- src/operator/src/error.rs | 12 +- src/operator/src/insert.rs | 53 +- src/operator/src/statement/ddl.rs | 51 +- src/operator/src/statement/dml.rs | 35 +- src/operator/src/statement/semantic_graph.rs | 1004 +++++++++++++++-- src/table/src/predicate.rs | 171 +++ tests-integration/src/grpc.rs | 172 ++- tests-integration/src/instance.rs | 104 ++ .../common/system/semantic_graph.result | 124 ++ .../common/system/semantic_graph.sql | 69 ++ 17 files changed, 2004 insertions(+), 145 deletions(-) diff --git a/src/auth/src/lib.rs b/src/auth/src/lib.rs index 44382737ee8..d2a1c05b1cd 100644 --- a/src/auth/src/lib.rs +++ b/src/auth/src/lib.rs @@ -34,7 +34,7 @@ pub use permission::{ DefaultPermissionChecker, INFLUXDB_WRITE, JAEGER_QUERY, LOG_QUERY, LOG_WRITE, OPENTSDB_WRITE, OTLP_WRITE, PIPELINE_DELETE, PIPELINE_INSERT, PIPELINE_QUERY, PROM_STORE_READ, PROM_STORE_WRITE, PROMQL_QUERY, PermissionAction, PermissionChecker, PermissionReq, - PermissionResp, PermissionTableTarget, PermissionTableTargets, + PermissionResp, PermissionTableTarget, PermissionTableTargets, SEMANTIC_GRAPH_QUERY, }; pub use user_info::UserInfo; pub use user_provider::static_user_provider::StaticUserProvider; diff --git a/src/auth/src/permission.rs b/src/auth/src/permission.rs index 60a89225738..901b935d6a2 100644 --- a/src/auth/src/permission.rs +++ b/src/auth/src/permission.rs @@ -150,6 +150,10 @@ pub const PIPELINE_DELETE: PermissionAction = PermissionAction::write("pipeline. pub const DASHBOARD_QUERY: PermissionAction = PermissionAction::read("dashboard.query"); pub const DASHBOARD_SAVE: PermissionAction = PermissionAction::write("dashboard.save"); pub const DASHBOARD_DELETE: PermissionAction = PermissionAction::write("dashboard.delete"); +/// Reading a telemetry table through the entity-graph derivation +/// (`greptime_private.semantic_entities` / `semantic_relationships`); checked +/// once per candidate source table with that table as the target. +pub const SEMANTIC_GRAPH_QUERY: PermissionAction = PermissionAction::read("semantic_graph.query"); /// All permission actions built into GreptimeDB. /// @@ -171,6 +175,7 @@ pub const ALL_ACTIONS: &[PermissionAction] = &[ DASHBOARD_QUERY, DASHBOARD_SAVE, DASHBOARD_DELETE, + SEMANTIC_GRAPH_QUERY, ]; #[derive(Debug, Clone)] diff --git a/src/catalog/src/kvbackend/manager.rs b/src/catalog/src/kvbackend/manager.rs index db7361615d0..2afc77c0d23 100644 --- a/src/catalog/src/kvbackend/manager.rs +++ b/src/catalog/src/kvbackend/manager.rs @@ -686,8 +686,12 @@ impl SystemCatalog { // Constructed on demand (the provider is a name + a weak ref); the // system catalog is consulted before physical resolution, so the // computed tables shadow same-named physical tables by design. - SemanticGraphTableProvider::new(catalog.to_string(), self.catalog_manager.clone()) - .table(table_name) + SemanticGraphTableProvider::new( + catalog.to_string(), + self.catalog_manager.clone(), + query_ctx.cloned().map(Arc::new), + ) + .table(table_name) } else { None } diff --git a/src/catalog/src/system_schema/semantic_graph.rs b/src/catalog/src/system_schema/semantic_graph.rs index 049377f5a38..028ce79b906 100644 --- a/src/catalog/src/system_schema/semantic_graph.rs +++ b/src/catalog/src/system_schema/semantic_graph.rs @@ -18,10 +18,9 @@ //! They live in `greptime_private`, not `information_schema`: scanning them //! triggers read-time derivation over telemetry tables (trace self-joins, ...), //! which breaks the "cheap, metadata-only" expectation users have of -//! `information_schema`. `greptime_private` already signals "system-managed, -//! computed data objects", and it is also where the physical declared-edge -//! table will live (a follow-up), so derived and declared edges will share one -//! schema. +//! `information_schema`. `greptime_private` also hosts the physical +//! declared-edge table (`semantic_relationships_declared`), whose rows are +//! unioned into the computed `semantic_relationships`. //! //! These are thin forwarders: their rows are derived at read time by the injected //! [`EntityGraphProvider`], which enumerates the `table_semantics` declarations, @@ -49,6 +48,7 @@ use common_recordbatch::{ use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; use futures::StreamExt; +use session::context::QueryContextRef; use snafu::ResultExt; use store_api::storage::{ScanRequest, TableId}; use table::TableRef; @@ -59,17 +59,18 @@ use crate::system_schema::{SystemSchemaProviderInner, SystemTable, SystemTableRe pub type EntityGraphProviderRef = Arc; -/// Produces the rows of the computed entity-graph tables -/// (`semantic_entities` / `semantic_relationships`) at read time. +/// Produces the rows of the computed entity-graph tables at read time. /// -/// The computed tables are thin forwarders to this provider, which is -/// implemented above the query engine (in the frontend): it enumerates the entity -/// declarations from `table_semantics`, builds the typed derivation plans, and -/// executes them. It is injected into the catalog manager *after* construction — -/// the provider needs the engine, which needs the catalog manager — so this late -/// binding breaks the `catalog -> query` dependency cycle. Keeping derivation out -/// of `catalog` also respects the crate layering (the plan builders live in -/// `operator`). See `docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md`. +/// Implemented above the query engine (in the frontend) and injected into the +/// catalog manager *after* construction — the provider needs the engine, which +/// needs the catalog manager — so this late binding breaks the +/// `catalog -> query` dependency cycle. +/// +/// `query_ctx` is the outer query's context, captured when the computed table +/// is resolved: the derivation must read only sources that context may read +/// and execute under it, inheriting the caller's permissions, cancellation and +/// deadline (the RFC's derivation contract). `None` (context-less internal +/// resolution) keeps the provider's default behaviour. #[async_trait::async_trait] pub trait EntityGraphProvider: Send + Sync { /// Produces the entity registry (`semantic_entities`) rows for `catalog`. @@ -78,6 +79,7 @@ pub trait EntityGraphProvider: Send + Sync { &self, catalog: &str, request: ScanRequest, + query_ctx: Option, ) -> std::result::Result, BoxedError>; /// Produces the relationship set (`semantic_relationships`) rows for `catalog`. @@ -85,6 +87,7 @@ pub trait EntityGraphProvider: Send + Sync { &self, catalog: &str, request: ScanRequest, + query_ctx: Option, ) -> std::result::Result, BoxedError>; } @@ -94,13 +97,21 @@ pub trait EntityGraphProvider: Send + Sync { pub(crate) struct SemanticGraphTableProvider { catalog_name: String, catalog_manager: Weak, + /// The resolving query's context; captured per resolution (the provider is + /// built on demand, never cached), so it cannot leak across sessions. + query_ctx: Option, } impl SemanticGraphTableProvider { - pub(crate) fn new(catalog_name: String, catalog_manager: Weak) -> Self { + pub(crate) fn new( + catalog_name: String, + catalog_manager: Weak, + query_ctx: Option, + ) -> Self { Self { catalog_name, catalog_manager, + query_ctx, } } @@ -139,6 +150,7 @@ impl SystemSchemaProviderInner for SemanticGraphTableProvider { kind, self.catalog_name.clone(), self.catalog_manager.clone(), + self.query_ctx.clone(), )) as _) } } @@ -252,6 +264,7 @@ struct SemanticGraphTable { schema: SchemaRef, catalog_name: String, catalog_manager: Weak, + query_ctx: Option, } impl SemanticGraphTable { @@ -259,6 +272,7 @@ impl SemanticGraphTable { kind: GraphTableKind, catalog_name: String, catalog_manager: Weak, + query_ctx: Option, ) -> Self { let schema = match kind { GraphTableKind::Entities => ENTITIES_SCHEMA.clone(), @@ -269,6 +283,7 @@ impl SemanticGraphTable { schema, catalog_name, catalog_manager, + query_ctx, } } @@ -277,6 +292,7 @@ impl SemanticGraphTable { catalog: String, catalog_manager: Weak, request: ScanRequest, + query_ctx: Option, ) -> Result> { let provider = utils::entity_graph_provider(&catalog_manager)?; // No provider (engine not up / non-frontend node): stream empty. @@ -284,8 +300,12 @@ impl SemanticGraphTable { return Ok(None); }; match kind { - GraphTableKind::Entities => provider.scan_entities(&catalog, request).await, - GraphTableKind::Relationships => provider.scan_relationships(&catalog, request).await, + GraphTableKind::Entities => provider.scan_entities(&catalog, request, query_ctx).await, + GraphTableKind::Relationships => { + provider + .scan_relationships(&catalog, request, query_ctx) + .await + } } .context(InternalSnafu) } @@ -338,10 +358,11 @@ impl SystemTable for SemanticGraphTable { let kind = self.kind; let catalog = self.catalog_name.clone(); let catalog_manager = self.catalog_manager.clone(); + let query_ctx = self.query_ctx.clone(); let stream_schema = schema.clone(); let stream = async move { - let stream = Self::derive(kind, catalog, catalog_manager, request) + let stream = Self::derive(kind, catalog, catalog_manager, request, query_ctx) .await .map_err(BoxedError::new) .context(common_recordbatch::error::ExternalSnafu)?; diff --git a/src/common/catalog/src/consts.rs b/src/common/catalog/src/consts.rs index 64bbbcbde86..6053d5a67c8 100644 --- a/src/common/catalog/src/consts.rs +++ b/src/common/catalog/src/consts.rs @@ -167,6 +167,14 @@ pub fn is_readonly_table(schema: &str, table: &str) -> bool { )) } +/// Returns true for tables whose *definition* is system-owned while their rows +/// are user-written: created with a canonical schema on first write; user +/// CREATE/ALTER/RENAME-into rejected; DML and DROP/TRUNCATE allowed (the next +/// write recreates the table). Currently only the declared-edge table. +pub fn is_ddl_reserved_table(schema: &str, table: &str) -> bool { + schema == DEFAULT_PRIVATE_SCHEMA_NAME && table == SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME +} + // ---- special table and fields ---- pub const TRACE_ID_COLUMN: &str = "trace_id"; pub const SPAN_ID_COLUMN: &str = "span_id"; @@ -208,4 +216,7 @@ pub const SEMANTIC_ENTITIES_TABLE_NAME: &str = "semantic_entities"; /// Computed relationship set: the edge set of the observability graph, derived at /// read time (`calls`/`runs_on`/... ) and unioned with the declared-edge table. pub const SEMANTIC_RELATIONSHIPS_TABLE_NAME: &str = "semantic_relationships"; +/// Physical table of hand-declared edges, unioned into the computed +/// `semantic_relationships`; see [`is_ddl_reserved_table`] for its lifecycle. +pub const SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME: &str = "semantic_relationships_declared"; // ---- End of entity relationship graph tables ---- diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index 8e69c5951b8..85a2b8ff183 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -15,6 +15,7 @@ use std::sync::Arc; use std::sync::atomic::AtomicBool; +use auth::PermissionCheckerRef; use cache::{PARTITION_INFO_CACHE_NAME, TABLE_FLOWNODE_SET_CACHE_NAME, TABLE_ROUTE_CACHE_NAME}; use catalog::CatalogManagerRef; use catalog::kvbackend::KvBackendCatalogManager; @@ -283,6 +284,7 @@ impl FrontendBuilder { let provider: EntityGraphProviderRef = Arc::new(EntityGraphProviderImpl::new( query_engine.clone(), Arc::downgrade(&self.catalog_manager), + plugins.get::(), )); kv_catalog.set_entity_graph_provider(provider); } diff --git a/src/frontend/src/instance/entity_graph.rs b/src/frontend/src/instance/entity_graph.rs index 1d012b7fc60..e55e5f2b8d7 100644 --- a/src/frontend/src/instance/entity_graph.rs +++ b/src/frontend/src/instance/entity_graph.rs @@ -27,29 +27,37 @@ use std::collections::HashMap; use std::sync::Weak; use async_trait::async_trait; +use auth::{ + PermissionChecker, PermissionCheckerRef, PermissionReq, PermissionTableTarget, + PermissionTableTargets, SEMANTIC_GRAPH_QUERY, +}; use catalog::CatalogManager; use catalog::system_schema::semantic_graph::EntityGraphProvider; use common_catalog::consts::{ DEFAULT_PRIVATE_SCHEMA_NAME, DEFAULT_SCHEMA_NAME, INFORMATION_SCHEMA_NAME, PG_CATALOG_NAME, - SERVICE_NAME_COLUMN, + SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SERVICE_NAME_COLUMN, }; -use common_error::ext::BoxedError; +use common_error::ext::{BoxedError, ErrorExt}; +use common_error::status_code::StatusCode; use common_query::OutputData; use common_recordbatch::SendableRecordBatchStream; -use common_telemetry::warn; +use common_telemetry::{debug, warn}; +use common_time::timestamp::TimeUnit; use datafusion::dataframe::DataFrame; use datafusion_expr::LogicalPlan; use futures::TryStreamExt; use operator::statement::semantic_graph::{ - CallsSource, EntityDeclaration, GraphWindow, RegistrySource, build_calls_plan, - build_registry_plan, + CallsSource, DeclaredSource, EntityDeclaration, GraphQueryWindow, OBSERVED_AT_COLUMN, + RegistrySource, build_registry_plan, build_relationships_plan, + declared_relationships_schema_matches, }; use query::QueryEngineRef; -use session::context::{QueryContextBuilder, QueryContextRef}; +use session::context::{QueryContext, QueryContextBuilder, QueryContextRef}; use snafu::ResultExt; use store_api::storage::ScanRequest; use table::TableRef; use table::metadata::TableInfo; +use table::predicate::{TimeRangeExtraction, extract_time_range_strict}; use table::requests::{ EntityRole, is_trace_v1_table, parse_entity_columns, parse_entity_option_key, }; @@ -60,6 +68,7 @@ use crate::error; pub struct EntityGraphProviderImpl { query_engine: QueryEngineRef, catalog_manager: Weak, + permission_checker: Option, } struct EntitySource { @@ -68,10 +77,40 @@ struct EntitySource { } impl EntityGraphProviderImpl { - pub fn new(query_engine: QueryEngineRef, catalog_manager: Weak) -> Self { + pub fn new( + query_engine: QueryEngineRef, + catalog_manager: Weak, + permission_checker: Option, + ) -> Self { Self { query_engine, catalog_manager, + permission_checker, + } + } + + /// Whether the caller may read the derivation sources named by `targets`. + /// `Ok(false)` = denied: the source is silently excluded, per the + /// derivation contract. Errors other than a denial abort the scan. + fn authorize_sources( + &self, + query_ctx: Option<&QueryContext>, + targets: PermissionTableTargets, + ) -> Result { + let Some(ctx) = query_ctx else { + return Ok(true); + }; + match self + .permission_checker + .as_ref() + .check_permission_with_table_targets( + ctx.current_user(), + PermissionReq::Action(SEMANTIC_GRAPH_QUERY), + targets, + ) { + Ok(_) => Ok(true), + Err(err) if err.status_code() == StatusCode::PermissionDenied => Ok(false), + Err(err) => Err(BoxedError::new(err)), } } @@ -182,15 +221,37 @@ impl EntityGraphProviderImpl { async fn enumerate( &self, catalog: &str, + query_ctx: Option<&QueryContext>, ) -> Result<(Vec, Vec<(EntityDeclaration, TableRef)>), BoxedError> { let Some(catalog_manager) = self.catalog_manager.upgrade() else { return Ok((vec![], vec![])); }; + // A target-blind checker (e.g. the default mode-based one) answers the + // same for every table: ask once up front instead of per table. + let per_table_auth = match self.permission_checker.as_ref() { + Some(checker) if query_ctx.is_some() => { + if checker.uses_table_targets() { + true + } else if self + .authorize_sources(query_ctx, PermissionTableTargets::resolved(vec![]))? + { + false + } else { + debug!( + "Caller lacks the entity-graph read permission; deriving an empty graph \ + (catalog: {catalog})" + ); + return Ok((vec![], vec![])); + } + } + _ => false, + }; + let mut declarations = vec![]; let mut traces = vec![]; let schemas = catalog_manager - .schema_names(catalog, None) + .schema_names(catalog, query_ctx) .await .map_err(BoxedError::new)?; for schema in schemas { @@ -202,11 +263,31 @@ impl EntityGraphProviderImpl { { continue; } - let mut tables = catalog_manager.tables(catalog, &schema, None); + let mut tables = catalog_manager.tables(catalog, &schema, query_ctx); while let Some(table) = tables.try_next().await.map_err(BoxedError::new)? { let table_info = table.table_info(); let table_declarations = Self::declarations_for(&table_info); - if is_trace_v1_table(&table_info) { + let is_trace = is_trace_v1_table(&table_info); + // Authorize only tables that would contribute rows. + if per_table_auth + && (is_trace || !table_declarations.is_empty()) + && !self.authorize_sources( + query_ctx, + PermissionTableTargets::resolved(vec![PermissionTableTarget::new( + catalog, + &schema, + &table_info.name, + )]), + )? + { + debug!( + "Excluding `{schema}.{}` from the entity-graph derivation: caller lacks \ + read permission", + table_info.name + ); + continue; + } + if is_trace { // TODO(entity-graph): validate the complete fixed trace-v1 // schema before adding the table to calls derivation, and // skip malformed/stale tables instead of letting one poison @@ -235,37 +316,122 @@ impl EntityGraphProviderImpl { Ok((declarations, traces)) } - /// The time window to derive over, taken from the scan's time predicate. + /// The time window to derive over, taken from the scan's `observed_at` + /// predicate. /// - /// TODO(entity-graph): read the `observed_at` bounds out of `request.filters` - /// and build the window from them. Until then this always returns the default - /// (last hour), so a query whose time filter falls entirely outside the last - /// hour derives nothing and returns empty — a known limitation tracked for the - /// scan-time-predicate-pushdown follow-up. - fn query_window(_request: &ScanRequest) -> GraphWindow { - GraphWindow::default_last_hour() + /// The contract (RFC "The contract"): no `observed_at` predicate → the + /// product default (last hour); a missing upper bound means "up to now"; a + /// predicate without a lower bound or in a shape that cannot be safely + /// extracted is an error asking for an explicit range — never a silent + /// fallback into incomplete results. + /// `Ok(None)` = the window cannot match anything (e.g. a lower bound in + /// the future with the implicit "up to now" upper bound): the scan streams + /// empty instead of deriving. + fn query_window(request: &ScanRequest) -> Result, BoxedError> { + let invalid = + |err_msg: String| Err(BoxedError::new(error::InvalidSqlSnafu { err_msg }.build())); + match extract_time_range_strict(OBSERVED_AT_COLUMN, TimeUnit::Millisecond, &request.filters) + { + TimeRangeExtraction::Absent => Ok(Some(GraphQueryWindow::default_last_hour())), + TimeRangeExtraction::Extracted(range) => { + let Some(start) = range.start() else { + return invalid(format!( + "the {OBSERVED_AT_COLUMN} filter has no lower bound; the graph cannot \ + derive over unbounded history — add e.g. {OBSERVED_AT_COLUMN} >= \ + '2026-01-01 00:00:00'" + )); + }; + let end_ms = range + .end() + .map(|ts| ts.value()) + .unwrap_or_else(common_time::util::current_time_millis); + if start.value() >= end_ms { + return Ok(None); + } + Ok(Some(GraphQueryWindow::from_observed(start.value(), end_ms))) + } + TimeRangeExtraction::Unsupported => invalid(format!( + "cannot derive the graph window from the {OBSERVED_AT_COLUMN} filter; use plain \ + range predicates (>=, <, BETWEEN) with literal bounds" + )), + } } fn read_table(&self, table: TableRef) -> Result { self.query_engine.read_table(table).map_err(BoxedError::new) } - /// Executes a derivation plan and returns its live result stream. - /// - /// TODO(entity-graph): the QueryContext must come from the outer query - /// instead of being built here, so the derivation inherits the caller's - /// permissions, cancellation and deadline. Requires threading the context - /// through the computed-table scan path (planned next PR). + /// The declared-edge branch source, when the physical table exists, the + /// caller may read it, and it still matches the canonical definition. A + /// mismatch (upgrade skew) is an explicit error, not a silent skip: + /// dropping declared edges would misrepresent the graph. + async fn declared_source( + &self, + catalog: &str, + query_ctx: Option<&QueryContext>, + ) -> Result, BoxedError> { + let Some(catalog_manager) = self.catalog_manager.upgrade() else { + return Ok(None); + }; + let Some(table) = catalog_manager + .table( + catalog, + DEFAULT_PRIVATE_SCHEMA_NAME, + SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, + query_ctx, + ) + .await + .map_err(BoxedError::new)? + else { + return Ok(None); + }; + if !self.authorize_sources( + query_ctx, + PermissionTableTargets::resolved(vec![PermissionTableTarget::new( + catalog, + DEFAULT_PRIVATE_SCHEMA_NAME, + SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, + )]), + )? { + debug!( + "Excluding declared edges: caller lacks read permission on \ + `{DEFAULT_PRIVATE_SCHEMA_NAME}.{SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME}`" + ); + return Ok(None); + } + let table_info = table.table_info(); + if !declared_relationships_schema_matches(&table_info) { + return Err(BoxedError::new( + error::InvalidSqlSnafu { + err_msg: format!( + "{DEFAULT_PRIVATE_SCHEMA_NAME}.{SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME} \ + does not match its canonical schema; declared edges cannot be derived" + ), + } + .build(), + )); + } + Ok(Some(DeclaredSource { + scan: self.read_table(table)?, + })) + } + + /// Executes a derivation plan under the caller's context (inheriting its + /// permissions, cancellation and deadline); context-less internal scans + /// get a minimal default. async fn execute_plan( &self, catalog: &str, plan: LogicalPlan, + query_ctx: Option, ) -> Result, BoxedError> { - let query_ctx: QueryContextRef = QueryContextBuilder::default() - .current_catalog(catalog.to_string()) - .current_schema(DEFAULT_SCHEMA_NAME.to_string()) - .build() - .into(); + let query_ctx = query_ctx.unwrap_or_else(|| { + QueryContextBuilder::default() + .current_catalog(catalog.to_string()) + .current_schema(DEFAULT_SCHEMA_NAME.to_string()) + .build() + .into() + }); let output = self .query_engine .execute(plan, query_ctx) @@ -286,8 +452,9 @@ impl EntityGraphProvider for EntityGraphProviderImpl { &self, catalog: &str, request: ScanRequest, + query_ctx: Option, ) -> Result, BoxedError> { - let (sources, _) = self.enumerate(catalog).await?; + let (sources, _) = self.enumerate(catalog, query_ctx.as_deref()).await?; let mut plans = Vec::with_capacity(sources.len()); for source in sources { plans.push(RegistrySource { @@ -295,22 +462,25 @@ impl EntityGraphProvider for EntityGraphProviderImpl { scan: self.read_table(source.table)?, }); } - let window = Self::query_window(&request); + let Some(window) = Self::query_window(&request)? else { + return Ok(None); + }; let Some(plan) = build_registry_plan(plans, &window) .context(error::DataFusionSnafu) .map_err(BoxedError::new)? else { return Ok(None); }; - self.execute_plan(catalog, plan).await + self.execute_plan(catalog, plan, query_ctx).await } async fn scan_relationships( &self, catalog: &str, request: ScanRequest, + query_ctx: Option, ) -> Result, BoxedError> { - let (_, traces) = self.enumerate(catalog).await?; + let (_, traces) = self.enumerate(catalog, query_ctx.as_deref()).await?; let mut scans = Vec::with_capacity(traces.len()); for (service, trace) in traces { scans.push(CallsSource { @@ -318,14 +488,17 @@ impl EntityGraphProvider for EntityGraphProviderImpl { scan: self.read_table(trace)?, }); } - let window = Self::query_window(&request); - let Some(plan) = build_calls_plan(scans, &window) + let declared = self.declared_source(catalog, query_ctx.as_deref()).await?; + let Some(window) = Self::query_window(&request)? else { + return Ok(None); + }; + let Some(plan) = build_relationships_plan(scans, declared, &window) .context(error::DataFusionSnafu) .map_err(BoxedError::new)? else { return Ok(None); }; - self.execute_plan(catalog, plan).await + self.execute_plan(catalog, plan, query_ctx).await } } @@ -393,6 +566,30 @@ mod tests { .unwrap() } + #[test] + fn future_lower_only_window_matches_nothing() { + let future_ms = common_time::util::current_time_millis() + 86_400_000; + let request = ScanRequest { + filters: vec![ + datafusion_expr::col(OBSERVED_AT_COLUMN).gt_eq(datafusion_expr::lit( + datafusion::common::ScalarValue::TimestampMillisecond(Some(future_ms), None), + )), + ], + ..Default::default() + }; + assert!( + EntityGraphProviderImpl::query_window(&request) + .unwrap() + .is_none() + ); + + assert!( + EntityGraphProviderImpl::query_window(&ScanRequest::default()) + .unwrap() + .is_some() + ); + } + #[test] fn declaration_referencing_missing_column_is_skipped() { let info = table_info( diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs index 324825d4589..b5e76f0f2cc 100644 --- a/src/operator/src/error.rs +++ b/src/operator/src/error.rs @@ -356,6 +356,15 @@ pub enum Error { location: Location, }, + #[snafu(display( + "The definition of table `{name}` is managed by GreptimeDB; it cannot be created or altered (DROP recreates it on the next write)" + ))] + TableDdlReserved { + name: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Table occurs error"))] Table { #[snafu(implicit)] @@ -986,7 +995,8 @@ impl ErrorExt for Error { Error::NotSupported { .. } | Error::ShowCreateTableBaseOnly { .. } | Error::SchemaReadOnly { .. } - | Error::TableReadOnly { .. } => StatusCode::Unsupported, + | Error::TableReadOnly { .. } + | Error::TableDdlReserved { .. } => StatusCode::Unsupported, Error::TableMetadataManager { source, .. } => source.status_code(), Error::ParseSql { source, .. } => source.status_code(), Error::InvalidateTableCache { source, .. } => source.status_code(), diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 8205932708e..b5b6be093d9 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -29,8 +29,8 @@ use catalog::CatalogManagerRef; use client::{OutputData, OutputMeta}; use common_catalog::consts::{ PARENT_SPAN_ID_COLUMN, SERVICE_NAME_COLUMN, TRACE_ID_COLUMN, TRACE_TABLE_NAME, - TRACE_TABLE_NAME_SESSION_KEY, default_engine, trace_operations_table_name, - trace_services_table_name, + TRACE_TABLE_NAME_SESSION_KEY, default_engine, is_ddl_reserved_table, + trace_operations_table_name, trace_services_table_name, }; use common_grpc_expr::util::ColumnExpr; use common_meta::cache::TableFlownodeSetCacheRef; @@ -578,15 +578,25 @@ impl Inserter { if let Some(disabled_reason) = self.auto_create_disabled_reason(ctx)? { let mut instant_table_ids = HashSet::new(); for req in &requests.inserts { - let table = self - .get_table(catalog, &schema, &req.table_name) - .await? - .context(InvalidInsertRequestSnafu { - reason: format!( - "Table `{}` does not exist, and {}", - req.table_name, disabled_reason - ), - })?; + let table = match self.get_table(catalog, &schema, &req.table_name).await? { + Some(table) => table, + // System-defined table: created canonically by the system, + // so the auto-create config/hint does not apply. + None if is_ddl_reserved_table(&schema, &req.table_name) => { + statement_executor + .create_declared_relationships_table(catalog, ctx.clone()) + .await? + } + None => { + return InvalidInsertRequestSnafu { + reason: format!( + "Table `{}` does not exist, and {}", + req.table_name, disabled_reason + ), + } + .fail(); + } + }; let table_info = table.table_info(); if table_info.is_ttl_instant_table() { instant_table_ids.insert(table_info.table_id()); @@ -630,6 +640,19 @@ impl Inserter { table_infos.insert(table_info.table_id(), table.table_info()); } } + // A DDL-reserved table's definition never derives from the + // write request; the system creates it canonically, below the + // user-DDL guard that rejects the generic create path. + None if is_ddl_reserved_table(&schema, &req.table_name) => { + let table = statement_executor + .create_declared_relationships_table(catalog, ctx.clone()) + .await?; + let table_info = table.table_info(); + if table_info.is_ttl_instant_table() { + instant_table_ids.insert(table_info.table_id()); + } + table_infos.insert(table_info.table_id(), table_info); + } None => { let create_expr = self.get_create_table_expr_on_demand(req, &auto_create_table_type, ctx)?; @@ -911,6 +934,7 @@ impl Inserter { create_type: &AutoCreateTableType, ctx: &QueryContextRef, ) -> Result { + let schema = ctx.current_schema(); let mut table_options = std::collections::HashMap::with_capacity(4); fill_table_options_for_create(&mut table_options, create_type, ctx); apply_per_table_semantic_options(&mut table_options, ctx, &req.table_name); @@ -922,7 +946,6 @@ impl Inserter { default_engine() }; - let schema = ctx.current_schema(); let table_ref = TableReference::full(ctx.current_catalog(), &schema, &req.table_name); // SAFETY: `req.rows` is guaranteed to be `Some` by `handle_row_inserts_with_create_type()`. let request_schema = req.rows.as_ref().unwrap().schema.as_slice(); @@ -963,6 +986,12 @@ impl Inserter { let schema_name = ctx.current_schema(); let table_name = table.table_info().name.clone(); + // Never auto-alter a system-defined table to fit a write; a request + // with unknown columns fails instead. + if is_ddl_reserved_table(&schema_name, &table_name) { + return Ok(None); + } + let request_schema = req.rows.as_ref().unwrap().schema.as_slice(); let request_field_count = request_schema .iter() diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index ae55acd2fb8..3f637ebebcf 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -32,7 +32,8 @@ use catalog::CatalogManagerRef; use chrono::Utc; use common_base::regex_pattern::NAME_PATTERN_REG; use common_catalog::consts::{ - DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, is_readonly_schema, is_readonly_table, + DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, is_ddl_reserved_table, is_readonly_schema, + is_readonly_table, }; use common_catalog::{format_full_flow_name, format_full_table_name}; use common_error::ext::BoxedError; @@ -112,8 +113,8 @@ use crate::error::{ InvalidPartitionSnafu, InvalidSqlSnafu, InvalidTableNameSnafu, InvalidViewNameSnafu, InvalidViewStmtSnafu, NotSupportedSnafu, PartitionExprToPbSnafu, Result, SchemaInUseSnafu, SchemaNotFoundSnafu, SchemaReadOnlySnafu, SerializePartitionExprSnafu, SubstraitCodecSnafu, - TableAlreadyExistsSnafu, TableMetadataManagerSnafu, TableNotFoundSnafu, TableReadOnlySnafu, - UnrecognizedTableOptionSnafu, ViewAlreadyExistsSnafu, + TableAlreadyExistsSnafu, TableDdlReservedSnafu, TableMetadataManagerSnafu, TableNotFoundSnafu, + TableReadOnlySnafu, UnrecognizedTableOptionSnafu, ViewAlreadyExistsSnafu, }; use crate::expr_helper::{self, RepartitionRequest, RepartitionSource}; use crate::statement::StatementExecutor; @@ -379,6 +380,19 @@ impl StatementExecutor { .await } + /// Creates the declared-edge table with its canonical schema, entering + /// below the user-DDL guard ([`ensure_table_definition_writable`]). + /// `create_if_not_exists` makes concurrent first inserts race safely. + pub async fn create_declared_relationships_table( + &self, + catalog: &str, + query_ctx: QueryContextRef, + ) -> Result { + let mut expr = super::semantic_graph::build_declared_relationships_expr(catalog); + self.create_non_logic_table(&mut expr, None, query_ctx, TriggerReason::AutoCreate) + .await + } + #[tracing::instrument(skip_all)] pub async fn create_table_inner( &self, @@ -387,7 +401,7 @@ impl StatementExecutor { query_ctx: QueryContextRef, trigger_reason: TriggerReason, ) -> Result { - ensure_table_writable(&create_table.schema_name, &create_table.table_name)?; + ensure_table_definition_writable(&create_table.schema_name, &create_table.table_name)?; if create_table.engine == METRIC_ENGINE_NAME && create_table @@ -1002,6 +1016,9 @@ impl StatementExecutor { expr: CreateViewExpr, ctx: QueryContextRef, ) -> Result { + // A view could otherwise squat a reserved name and block the canonical + // table's first-write creation. + ensure_table_definition_writable(&expr.schema_name, &expr.view_name)?; ensure! { !(expr.create_if_not_exists & expr.or_replace), InvalidSqlSnafu { @@ -1540,7 +1557,7 @@ impl StatementExecutor { query_context: &QueryContextRef, ) -> Result { // Check if the schema is read-only. - ensure_table_writable(&request.schema_name, &request.table_name)?; + ensure_table_definition_writable(&request.schema_name, &request.table_name)?; let table_ref = TableReference::full( &request.catalog_name, @@ -1823,7 +1840,7 @@ impl StatementExecutor { query_context: QueryContextRef, trigger_reason: TriggerReason, ) -> Result { - ensure_table_writable(&expr.schema_name, &expr.table_name)?; + ensure_table_definition_writable(&expr.schema_name, &expr.table_name)?; let catalog_name = if expr.catalog_name.is_empty() { DEFAULT_CATALOG_NAME.to_string() @@ -2328,7 +2345,7 @@ pub fn verify_alter( ); // Renaming INTO a computed graph table's name would let the overlay // shadow the renamed physical table, orphaning its data. - ensure_table_writable(&table_info.schema_name, new_table_name)?; + ensure_table_definition_writable(&table_info.schema_name, new_table_name)?; } else if let AlterKind::AddColumns { columns } = alter_kind { // If all the columns are marked as add_if_not_exists and they already exist in the table, // there is no need to perform the alter. @@ -2548,6 +2565,22 @@ fn ensure_table_writable(schema: &str, table: &str) -> Result<()> { Ok(()) } +/// [`ensure_table_writable`] plus the definition guard for system-defined +/// tables: user CREATE, ALTER and RENAME-into are rejected so the canonical +/// schema cannot be squatted or mutated. DROP and TRUNCATE stay allowed — the +/// next INSERT recreates the table canonically, which is also the recovery +/// path if the canonical definition changes across an upgrade. +fn ensure_table_definition_writable(schema: &str, table: &str) -> Result<()> { + ensure_table_writable(schema, table)?; + ensure!( + !is_ddl_reserved_table(schema, table), + TableDdlReservedSnafu { + name: table.to_string() + } + ); + Ok(()) +} + /// Validates `greptime.semantic.entity..{id|descriptive|scope}` options /// against the table schema: every named column must exist, and `id` columns must /// additionally be tag (primary-key) columns so entity identity stays indexable @@ -2800,7 +2833,9 @@ async fn execute_undrop_table( table_name: TableName, query_context: QueryContextRef, ) -> Result { - ensure_table_writable(&table_name.schema_name, &table_name.table_name)?; + // Undropping restores a table definition and could resurrect a + // pre-canonical shape of a DDL-reserved table; rejected like CREATE. + ensure_table_definition_writable(&table_name.schema_name, &table_name.table_name)?; let dropped = table_metadata_manager .get_dropped_table(&table_name) diff --git a/src/operator/src/statement/dml.rs b/src/operator/src/statement/dml.rs index 827bfd8b663..d72daaaf1aa 100644 --- a/src/operator/src/statement/dml.rs +++ b/src/operator/src/statement/dml.rs @@ -12,19 +12,25 @@ // See the License for the specific language governing permissions and // limitations under the License. +use common_catalog::consts::is_ddl_reserved_table; +use common_error::ext::BoxedError; use common_query::Output; use common_telemetry::tracing; use query::parser::QueryStatement; use session::context::QueryContextRef; +use session::table_name::table_idents_to_full_name; +use snafu::ResultExt; use sql::statements::insert::Insert; use sql::statements::statement::Statement; -use crate::error::Result; +use crate::error::{ExternalSnafu, ParseSqlSnafu, Result}; use crate::statement::StatementExecutor; impl StatementExecutor { #[tracing::instrument(skip_all)] pub async fn insert(&self, insert: Box, query_ctx: QueryContextRef) -> Result { + self.create_ddl_reserved_target_on_demand(&insert, &query_ctx) + .await?; if insert.can_extract_values() { // Fast path: plain insert ("insert with literal values") is executed directly self.inserter @@ -36,4 +42,31 @@ impl StatementExecutor { self.plan_exec(statement, query_ctx).await } } + + /// A DDL-reserved table is defined by the system, not the user: its first + /// INSERT creates it here with the canonical schema. + async fn create_ddl_reserved_target_on_demand( + &self, + insert: &Insert, + query_ctx: &QueryContextRef, + ) -> Result<()> { + let table_name = insert.table_name().context(ParseSqlSnafu)?; + let (catalog, schema, table) = table_idents_to_full_name(table_name, query_ctx) + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + if !is_ddl_reserved_table(&schema, &table) { + return Ok(()); + } + let exists = self + .catalog_manager + .table_exists(&catalog, &schema, &table, Some(query_ctx)) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + if !exists { + self.create_declared_relationships_table(&catalog, query_ctx.clone()) + .await?; + } + Ok(()) + } } diff --git a/src/operator/src/statement/semantic_graph.rs b/src/operator/src/statement/semantic_graph.rs index b845dce3b75..15f76905555 100644 --- a/src/operator/src/statement/semantic_graph.rs +++ b/src/operator/src/statement/semantic_graph.rs @@ -12,9 +12,13 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Typed DataFusion plan builders for the read-time derivation behind the +//! The entity-relationship graph: the physical declared-edge table DDL and the +//! typed DataFusion plan builders for the read-time derivation behind the //! computed `semantic_entities` / `semantic_relationships` tables. //! +//! In OSS the graph is derived at read time, so the *only* stored part is the +//! declared-edge table: edges a user asserts by hand (`provenance = 'declared'`). +//! //! The derivation is built as typed [`Expr`]s over [`DataFrame`]s (never as SQL //! text), so user-controlled identifiers are plain values — no quoting or SQL //! injection surface — and the plans compose with DataFusion's optimizer, @@ -23,10 +27,16 @@ use std::sync::{Arc, LazyLock}; +use api::v1::column_data_type_extension::TypeExt; +use api::v1::{ + ColumnDataType, ColumnDataTypeExtension, ColumnDef, CreateTableExpr, JsonTypeExtension, + SemanticType, +}; use common_catalog::consts::{ - DURATION_NANO_COLUMN, PARENT_SPAN_ID_COLUMN, SPAN_ID_COLUMN, SPAN_KIND_CLIENT, - SPAN_KIND_COLUMN, SPAN_KIND_SERVER, SPAN_STATUS_CODE_COLUMN, SPAN_STATUS_ERROR, - TRACE_ID_COLUMN, TRACE_TIMESTAMP_COLUMN, + DEFAULT_PRIVATE_SCHEMA_NAME, DURATION_NANO_COLUMN, PARENT_SPAN_ID_COLUMN, + SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SPAN_ID_COLUMN, SPAN_KIND_CLIENT, SPAN_KIND_COLUMN, + SPAN_KIND_SERVER, SPAN_STATUS_CODE_COLUMN, SPAN_STATUS_ERROR, TRACE_ID_COLUMN, + TRACE_TIMESTAMP_COLUMN, }; use common_function::function::FunctionContext; use common_function::function_registry::FUNCTION_REGISTRY; @@ -35,13 +45,221 @@ use datafusion::dataframe::DataFrame; use datafusion::functions::{core as core_fns, datetime as datetime_fns, string as string_fns}; use datafusion::functions_aggregate::expr_fn::{count, sum}; use datafusion::functions_nested::expr_fn::make_array; +use datafusion::functions_window::expr_fn::row_number; use datafusion_common::{Column, Result as DfResult, ScalarValue}; use datafusion_expr::{Expr, ExprFunctionExt, JoinType, LogicalPlan, ScalarUDF, cast, ident, lit}; +use store_api::mito_engine_options::{APPEND_MODE_KEY, MERGE_MODE_KEY, TTL_KEY}; + +/// Whether an existing declared-edge table still matches the canonical +/// definition ([`build_declared_relationships_expr`]): columns, time index, +/// primary key, engine, and merge behaviour — the union branch's revision and +/// dedup semantics lean on all of these. A mismatch (upgrade skew) must be +/// surfaced, not silently derived wrong; DROP resets the table. +pub fn declared_relationships_schema_matches(table_info: &table::metadata::TableInfo) -> bool { + let meta = &table_info.meta; + if meta.engine != common_catalog::consts::MITO_ENGINE { + return false; + } + let schema = &meta.schema; + if schema + .timestamp_column() + .is_none_or(|column| column.name != OBSERVED_AT_COLUMN) + { + return false; + } + let primary_keys = meta + .primary_key_indices + .iter() + .filter_map(|idx| schema.column_schemas().get(*idx)) + .map(|column| column.name.as_str()); + if !primary_keys.eq(DECLARED_PRIMARY_KEY_COLUMNS) { + return false; + } + // LastRow merge is what makes a re-asserted edge a *revision*; append mode + // or another merge mode would change the read-side dedup semantics. + let options = &meta.options.extra_options; + if options.get(APPEND_MODE_KEY).map(String::as_str) == Some("true") + || options + .get(MERGE_MODE_KEY) + .is_some_and(|mode| mode != "last_row") + { + return false; + } + + let canonical = build_declared_relationships_expr("greptime"); + if schema.column_schemas().len() != canonical.column_defs.len() { + return false; + } + canonical + .column_defs + .iter() + .zip(schema.column_schemas()) + .all(|(def, column)| { + let Ok(wrapper) = api::helper::ColumnDataTypeWrapper::try_new( + def.data_type, + def.datatype_extension.clone(), + ) else { + return false; + }; + def.name == column.name + && datatypes::prelude::ConcreteDataType::from(wrapper) == column.data_type + }) +} /// Bin width for the temporal window of derived rows: 60s buckets, matching the /// service-graph convention. const BIN_NANOS: i64 = 60 * 1_000_000_000; +/// Time index of the graph tables: when an observation was recorded. +pub const OBSERVED_AT_COLUMN: &str = "observed_at"; + +/// Default retention for the declared-edge table; expiry slides the topology window. +const DEFAULT_DECLARED_RELATIONSHIPS_TTL: &str = "90d"; +/// Environment variable overriding the declared-edge table's TTL at creation +/// time (e.g. `180d`, `forever`). +// TODO(entity-graph): promote this to a real configuration option. +const DECLARED_RELATIONSHIPS_TTL_ENV: &str = "GREPTIMEDB_DECLARED_RELATIONSHIPS_TTL"; + +fn declared_relationships_ttl() -> String { + std::env::var(DECLARED_RELATIONSHIPS_TTL_ENV) + .ok() + .map(|ttl| ttl.trim().to_string()) + .filter(|ttl| !ttl.is_empty()) + .unwrap_or_else(|| DEFAULT_DECLARED_RELATIONSHIPS_TTL.to_string()) +} + +/// The primary-key (tag) columns, in key order. Starting with the source endpoint +/// makes out-edge lookup (`WHERE src_type=? AND src_id=?`) a key-prefix scan; +/// `provenance` and `generation_id` are in the key so a declared edge and a +/// (future) derived edge for the same pair coexist without clobbering. +pub const DECLARED_PRIMARY_KEY_COLUMNS: [&str; 8] = [ + "src_type", + "src_id", + "rel_type", + "dst_type", + "dst_id", + "provenance", + "scope", + "generation_id", +]; + +/// The externally visible edge identity: the primary key minus `scope` and +/// `generation_id`, which the computed table does not expose. Revision ranking +/// uses this identity, or assertions differing only in those two columns would +/// surface as indistinguishable duplicate rows. +const DECLARED_EDGE_IDENTITY_COLUMNS: [&str; 6] = [ + "src_type", + "src_id", + "rel_type", + "dst_type", + "dst_id", + "provenance", +]; + +fn column( + name: &str, + data_type: ColumnDataType, + semantic_type: SemanticType, + nullable: bool, +) -> ColumnDef { + ColumnDef { + name: name.to_string(), + data_type: data_type as i32, + is_nullable: nullable, + default_constraint: vec![], + semantic_type: semantic_type as i32, + comment: String::new(), + datatype_extension: None, + options: None, + } +} + +fn tag(name: &str) -> ColumnDef { + column(name, ColumnDataType::String, SemanticType::Tag, false) +} + +fn field(name: &str, data_type: ColumnDataType) -> ColumnDef { + column(name, data_type, SemanticType::Field, true) +} + +fn json_field(name: &str) -> ColumnDef { + let mut def = field(name, ColumnDataType::Binary); + def.datatype_extension = Some(ColumnDataTypeExtension { + type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())), + }); + def +} + +/// Builds the `CREATE TABLE` request for the declared-edge table. Columns mirror +/// the computed `semantic_relationships` shape (temporal window + endpoints + +/// provenance/confidence + RED metrics) plus the declared-only business validity +/// window (`valid_from` / `valid_until`), which — unlike TTL (physical retention) +/// — expresses whether a hand-declared edge is still in effect. +/// +/// The default `LastRow` merge dedups on primary key **plus** `observed_at` +/// (the time index is part of mito's dedup key): re-asserting an edge at a new +/// `observed_at` stores a new revision, and the read-side union keeps only the +/// latest revision per edge key. +pub fn build_declared_relationships_expr(catalog: &str) -> CreateTableExpr { + let column_defs = vec![ + // Temporal: when this revision of the edge was declared (time index, + // also the TTL clock). + column( + OBSERVED_AT_COLUMN, + ColumnDataType::TimestampMillisecond, + SemanticType::Timestamp, + false, + ), + field("window_start", ColumnDataType::TimestampMillisecond), + field("window_end", ColumnDataType::TimestampMillisecond), + field("fresh_until", ColumnDataType::TimestampMillisecond), + // Declared-only business validity. NULL valid_from = valid since the + // declaration; NULL valid_until = valid for as long as the row exists + // (TTL expiry retires the edge with the row). + field("valid_from", ColumnDataType::TimestampMillisecond), + field("valid_until", ColumnDataType::TimestampMillisecond), + // Endpoints + edge identity (all tags, in primary-key order). + tag("src_type"), + tag("src_id"), + tag("rel_type"), + tag("dst_type"), + tag("dst_id"), + tag("provenance"), + tag("scope"), + tag("generation_id"), + // Confidence + RED metrics (populated for derived edges; usually NULL here). + field("confidence", ColumnDataType::Float64), + field("request_count", ColumnDataType::Int64), + field("error_count", ColumnDataType::Int64), + field("duration_sum", ColumnDataType::Float64), + field("duration_count", ColumnDataType::Int64), + // JSONB, so the union matches the computed table's json column + // without a per-scan parse. + json_field("attributes"), + ]; + + let table_options = [(TTL_KEY.to_string(), declared_relationships_ttl())] + .into_iter() + .collect(); + + CreateTableExpr { + catalog_name: catalog.to_string(), + schema_name: DEFAULT_PRIVATE_SCHEMA_NAME.to_string(), + table_name: SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME.to_string(), + desc: "Hand-declared edges of the entity-relationship graph".to_string(), + column_defs, + time_index: OBSERVED_AT_COLUMN.to_string(), + primary_keys: DECLARED_PRIMARY_KEY_COLUMNS + .iter() + .map(|c| c.to_string()) + .collect(), + create_if_not_exists: true, + table_options, + table_id: None, + engine: common_catalog::consts::MITO_ENGINE.to_string(), + } +} + /// A single table's entity-identity declaration, projected from /// `information_schema.table_semantics` (`greptime.semantic.entity..*`). #[derive(Debug, Clone)] @@ -62,33 +280,85 @@ pub struct EntityDeclaration { pub scope_columns: Vec, } -/// The read-time query window as scalar [`Expr`]s. Kept as expressions so the -/// provider can fill them from a scan's time predicate or the product default. +/// The read-time query window, resolved from the scan's `observed_at` predicate +/// (or the product default). Two half-open `[start, end)` millisecond ranges: +/// +/// - *observed*: the queried `observed_at` range. The scan's filters are +/// re-applied above the computed table (`FilterPushDownType::Inexact`), so +/// every emitted row's `observed_at` must fall inside it. The declared-edge +/// branch also runs its validity overlap against this range. +/// - *source scan*: `observed` widened to whole 60s buckets. Derived rows are +/// `date_bin('60s')` aggregates keyed by the bucket start, so a bucket the +/// observed range selects must be scanned over its full extent — filtering +/// the source rows to the observed range instead would silently truncate the +/// RED numbers of the boundary buckets. +/// +/// The bounds are plain millisecond values turned into timestamp literals on +/// demand — wall-clock snapshot semantics like `now()`, but as constants they +/// prune the source table scans without depending on constant-folding. #[derive(Debug, Clone)] -pub struct GraphWindow { - pub start: Expr, - pub end: Expr, +pub struct GraphQueryWindow { + observed_start_ms: i64, + observed_end_ms: i64, + source_start_ms: i64, + source_end_ms: i64, } -impl GraphWindow { +impl GraphQueryWindow { + /// Builds the window for the queried `observed_at` range `[start_ms, end_ms)`. + pub fn from_observed(start_ms: i64, end_ms: i64) -> Self { + const BIN_MS: i64 = BIN_NANOS / 1_000_000; + // Ceiling division that stays correct for pre-epoch (negative) values. + let ceil_to_bin = |ms: i64| { + ms.div_euclid(BIN_MS) * BIN_MS + + if ms.rem_euclid(BIN_MS) == 0 { + 0 + } else { + BIN_MS + } + }; + // A bucket `b` (a 60s multiple, `date_bin`'s floor of its rows) is + // selected when `start <= b < end`: the first is ceil(start) and the + // last one's rows extend to ceil(end). + Self { + observed_start_ms: start_ms, + observed_end_ms: end_ms, + source_start_ms: ceil_to_bin(start_ms), + source_end_ms: ceil_to_bin(end_ms), + } + } + /// Conservative default when a query carries no explicit time predicate: the /// last hour, so a bare `SELECT * FROM semantic_entities` never scans every /// declaring table's full history. This is a product default, not a cap. - /// - /// The bounds are timestamp literals taken from the wall clock at plan-build - /// time — the same snapshot semantics `now()` would give, but as plain - /// constants they prune the source table scans without depending on - /// constant-folding. pub fn default_last_hour() -> Self { let end_ms = common_time::util::current_time_millis(); - Self { - start: lit(ScalarValue::TimestampMillisecond( - Some(end_ms - 60 * 60 * 1000), - None, - )), - end: lit(ScalarValue::TimestampMillisecond(Some(end_ms), None)), - } + Self::from_observed(end_ms - 60 * 60 * 1000, end_ms) } + + /// Start of the queried `observed_at` range (inclusive). + pub fn observed_start(&self) -> Expr { + ts_ms_lit(self.observed_start_ms) + } + + /// End of the queried `observed_at` range (exclusive). + pub fn observed_end(&self) -> Expr { + ts_ms_lit(self.observed_end_ms) + } + + /// Start of the source-table scan range (inclusive). + pub fn source_start(&self) -> Expr { + ts_ms_lit(self.source_start_ms) + } + + /// End of the source-table scan range (exclusive). + pub fn source_end(&self) -> Expr { + ts_ms_lit(self.source_end_ms) + } +} + +fn ts_ms_lit(ms: i64) -> Expr { + lit(ScalarValue::TimestampMillisecond(Some(ms), None)) } /// An `INTERVAL` literal of `nanos` nanoseconds. @@ -254,14 +524,14 @@ fn registry_source( first: &EntityDeclaration, rest: &[EntityDeclaration], df: DataFrame, - window: &GraphWindow, + window: &GraphQueryWindow, ) -> DfResult { let ts = ident(&first.time_index); let bin = bin_ms(ts.clone()); let window_predicate = ts .clone() - .gt_eq(window.start.clone()) - .and(ts.lt(window.end.clone())); + .gt_eq(window.source_start()) + .and(ts.lt(window.source_end())); let mut arrays = vec![Vec::new(); REGISTRY_COLUMNS.len() + 1]; for decl in std::iter::once(first).chain(rest) { @@ -350,19 +620,25 @@ pub struct RegistrySource { } /// A trace table's scan paired with the `service` declaration it derives -/// `calls` edges for — the unit `build_calls_plan` consumes. +/// `calls` edges for — a unit of `build_relationships_plan`. pub struct CallsSource { pub service: EntityDeclaration, pub scan: DataFrame, } +/// The declared-edge table's scan (`semantic_relationships_declared`), whose +/// rows `build_relationships_plan` unions into the edge set. +pub struct DeclaredSource { + pub scan: DataFrame, +} + /// Builds the `semantic_entities` registry plan: one branch and source scan per /// declaring table, filtered to `window`, then `UNION ALL` across source tables. /// Returns `None` when nothing declared an entity, so the computed table streams /// empty. pub fn build_registry_plan( sources: Vec, - window: &GraphWindow, + window: &GraphQueryWindow, ) -> DfResult> { let mut union_df: Option = None; for source in sources { @@ -375,11 +651,11 @@ pub fn build_registry_plan( } /// The projected columns of `semantic_relationships`, in order. Every derived -/// branch and the declared-edge scan must project exactly these so the top-level -/// `UNION ALL` type-aligns. (The physical declared table additionally stores -/// `valid_from`/`valid_until`, which are applied as a filter, not projected.) -/// Test-only until the declared-edge union branch lands and enforces it in code. -#[cfg(test)] +/// branch and the declared-edge branch must project exactly these so the +/// top-level `UNION ALL` type-aligns; `build_relationships_plan` re-selects +/// them over the union to enforce the contract. (The physical declared table +/// additionally stores `valid_from`/`valid_until`, which feed the validity +/// filter and the projected window columns.) const RELATIONSHIP_COLUMNS: [&str; 16] = [ "observed_at", "window_start", @@ -406,21 +682,117 @@ const RELATIONSHIP_COLUMNS: [&str; 16] = [ const CHILD_SPAN_EARLY_NANOS: i64 = 5 * 60 * 1_000_000_000; const CHILD_SPAN_LATE_NANOS: i64 = 60 * 60 * 1_000_000_000; -/// Builds the `calls` derivation (RFC §3a) over `traces`: pair each client span -/// with its child server span on `trace_id` + `parent_span_id`, project to -/// `service`, union the pairs of all trace tables, and aggregate to RED metrics -/// per 60s window in one pass — so an edge observed across several trace tables -/// yields one row, not per-table fragments. This is the plan form of the Tempo -/// servicegraph connector. Virtual-node edges (uninstrumented peers) are a -/// separate branch, added on top of this. Column names are the fixed -/// `greptime_trace_v1` schema (the reason `table_data_model = greptime_trace_v1` -/// is required); `span_status_code` is a string column (`STATUS_CODE_ERROR`), -/// verified against the trace ingest path. Returns `None` when there is no -/// trace table. -pub fn build_calls_plan( +/// Builds the `semantic_relationships` plan: the trace-derived `calls` branch +/// unioned with the declared-edge branch, re-projected to the 16-column +/// contract. Returns `None` when there is neither a trace table nor a +/// declared-edge table, so the computed table streams empty. +pub fn build_relationships_plan( traces: Vec, - window: &GraphWindow, + declared: Option, + window: &GraphQueryWindow, ) -> DfResult> { + let mut union_df = calls_branch(traces, window)?; + if let Some(declared) = declared { + union_df = union_all(union_df, declared_edges(declared.scan, window)?)?; + } + union_df + .map(|df| { + Ok(df + .select_columns(&RELATIONSHIP_COLUMNS)? + .into_unoptimized_plan()) + }) + .transpose() +} + +/// Ranking column used to pick the latest declared-edge revision per edge key. +const DECLARED_REVISION_COLUMN: &str = "__declared_revision"; + +/// The declared-edge branch: the latest revision per edge key *as of the +/// queried window* (mito dedups on primary key + `observed_at`, so re-asserting +/// an edge stores a new revision), filtered by business-validity overlap. +/// +/// `valid_from` defaults to the declaration time; a NULL `valid_until` means +/// the edge holds while its row exists (TTL retires it), so `window_end` / +/// `fresh_until` take the window's upper bound. The output `observed_at` is +/// synthesized inside the queried range: the scan's filters are re-applied +/// above the computed table, and the physical revision time would fail them. +fn declared_edges(scan: DataFrame, window: &GraphQueryWindow) -> DfResult { + let eff_valid_from = + core_fns::coalesce().call(vec![ident("valid_from"), ident(OBSERVED_AT_COLUMN)]); + let eff_valid_until = + core_fns::coalesce().call(vec![ident("valid_until"), window.observed_end()]); + + // As-of the queried window: a revision recorded after its end, or whose + // validity starts after it, was not the edge's state inside the window and + // must not outrank (and thereby hide) the revision that was. + let as_of = ident(OBSERVED_AT_COLUMN) + .lt(window.observed_end()) + .and(eff_valid_from.clone().lt(window.observed_end())); + + let revision = row_number() + .partition_by( + DECLARED_EDGE_IDENTITY_COLUMNS + .iter() + .map(|c| ident(*c)) + .collect(), + ) + // generation_id/scope break observed_at ties deterministically (rows + // differing only in them are not merged by the storage dedup). + .order_by(vec![ + ident(OBSERVED_AT_COLUMN).sort(false, false), + ident("generation_id").sort(false, false), + ident("scope").sort(false, false), + ]) + .build()? + .alias(DECLARED_REVISION_COLUMN); + + // The as-of filter already bounds eff_valid_from below the window's end; + // only the retirement side of the overlap remains to check. + let still_valid = ident("valid_until") + .is_null() + .or(ident("valid_until").gt(window.observed_start())); + // Tags come out of the storage engine dictionary-encoded; cast to plain + // strings so the union with the derived branches type-aligns. + let tag_utf8 = |name: &str| cast(ident(name), DataType::Utf8).alias(name); + + scan.filter(as_of)? + .window(vec![revision])? + .filter(ident(DECLARED_REVISION_COLUMN).eq(lit(1_u64)))? + .filter(still_valid)? + .select(vec![ + core_fns::greatest() + .call(vec![eff_valid_from.clone(), window.observed_start()]) + .alias(OBSERVED_AT_COLUMN), + eff_valid_from.alias("window_start"), + eff_valid_until.clone().alias("window_end"), + eff_valid_until.alias("fresh_until"), + tag_utf8("src_type"), + tag_utf8("src_id"), + tag_utf8("dst_type"), + tag_utf8("dst_id"), + tag_utf8("rel_type"), + tag_utf8("provenance"), + ident("confidence"), + ident("request_count"), + ident("error_count"), + ident("duration_sum"), + ident("duration_count"), + ident("attributes"), + ]) +} + +/// The `calls` derivation: pair each client span with its child server span on +/// `trace_id` + `parent_span_id`, project to `service`, union the pairs of all +/// trace tables, and aggregate RED metrics per 60s window in one pass — an edge +/// observed across several trace tables yields one row, not per-table +/// fragments (the plan form of the Tempo servicegraph connector). Column names +/// are the fixed `greptime_trace_v1` schema, the reason +/// `table_data_model = greptime_trace_v1` is required. Returns `None` when +/// there is no trace table. +fn calls_branch( + traces: Vec, + window: &GraphQueryWindow, +) -> DfResult> { let mut union_df: Option = None; for trace in traces { union_df = union_all(union_df, calls_pairs(&trace.service, trace.scan, window)?)?; @@ -429,7 +801,7 @@ pub fn build_calls_plan( return Ok(None); }; - let plan = pairs + let df = pairs .aggregate( vec![ident("observed_at"), ident("src_id"), ident("dst_id")], vec![ @@ -460,9 +832,8 @@ pub fn build_calls_plan( .alias("duration_sum"), ident("request_count").alias("duration_count"), null_json().alias("attributes"), - ])? - .into_unoptimized_plan(); - Ok(Some(plan)) + ])?; + Ok(Some(df)) } /// One trace table's client/server span pairs, projected to the aggregation @@ -473,12 +844,12 @@ pub fn build_calls_plan( fn calls_pairs( service: &EntityDeclaration, trace: DataFrame, - window: &GraphWindow, + window: &GraphQueryWindow, ) -> DfResult { let mut client_pred = ident(SPAN_KIND_COLUMN) .eq(lit(SPAN_KIND_CLIENT)) - .and(ident(TRACE_TIMESTAMP_COLUMN).gt_eq(window.start.clone())) - .and(ident(TRACE_TIMESTAMP_COLUMN).lt(window.end.clone())); + .and(ident(TRACE_TIMESTAMP_COLUMN).gt_eq(window.source_start())) + .and(ident(TRACE_TIMESTAMP_COLUMN).lt(window.source_end())); let mut server_pred = ident(SPAN_KIND_COLUMN) .eq(lit(SPAN_KIND_SERVER)) // Static bounds implied by the window and the join's time-proximity @@ -486,10 +857,10 @@ fn calls_pairs( // cannot prune the server-side scan. .and( ident(TRACE_TIMESTAMP_COLUMN) - .gt_eq(window.start.clone() - interval(CHILD_SPAN_EARLY_NANOS)), + .gt_eq(window.source_start() - interval(CHILD_SPAN_EARLY_NANOS)), ) .and( - ident(TRACE_TIMESTAMP_COLUMN).lt(window.end.clone() + interval(CHILD_SPAN_LATE_NANOS)), + ident(TRACE_TIMESTAMP_COLUMN).lt(window.source_end() + interval(CHILD_SPAN_LATE_NANOS)), ); // A NULL identity component identifies nothing, on either endpoint. for id in &service.id_columns { @@ -541,14 +912,34 @@ mod tests { use super::*; - fn test_window() -> GraphWindow { - GraphWindow { - start: lit(ScalarValue::TimestampMillisecond(Some(0), None)), - end: lit(ScalarValue::TimestampMillisecond( - Some(10 * 60 * 1000), - None, - )), - } + fn test_window() -> GraphQueryWindow { + GraphQueryWindow::from_observed(0, 10 * 60 * 1000) + } + + #[test] + fn window_source_scan_covers_whole_buckets() { + // Buckets selected by [10:30.5', 11:30') is exactly {11:00'}; its source + // rows span [11:00', 12:00'). + let window = GraphQueryWindow::from_observed(630_000, 690_000); + assert_eq!( + (window.source_start_ms, window.source_end_ms), + (660_000, 720_000) + ); + + // Aligned bounds stay put. + let window = GraphQueryWindow::from_observed(600_000, 720_000); + assert_eq!( + (window.source_start_ms, window.source_end_ms), + (600_000, 720_000) + ); + + // A sub-bucket range that selects no bucket start scans nothing. + let window = GraphQueryWindow::from_observed(610_000, 650_000); + assert!(window.source_start_ms >= window.source_end_ms); + + // Pre-epoch bounds round toward the correct buckets. + let window = GraphQueryWindow::from_observed(-90_000, -30_000); + assert_eq!((window.source_start_ms, window.source_end_ms), (-60_000, 0)); } async fn collect(ctx: &SessionContext, plan: LogicalPlan) -> Vec { @@ -954,7 +1345,7 @@ mod tests { async fn calls_plan_aggregates_red_metrics() { let ctx = trace_table_ctx(); let trace = ctx.table("opentelemetry_traces").await.unwrap(); - let plan = build_calls_plan( + let plan = build_relationships_plan_for_test( vec![CallsSource { service: trace_service_decl(&[SERVICE_NAME_COLUMN]), scan: trace, @@ -1016,7 +1407,7 @@ mod tests { let ctx = trace_table_ctx(); let t1 = ctx.table("opentelemetry_traces").await.unwrap(); let t2 = ctx.table("opentelemetry_traces_2").await.unwrap(); - let plan = build_calls_plan( + let plan = build_relationships_plan_for_test( vec![ CallsSource { service: trace_service_decl(&[SERVICE_NAME_COLUMN]), @@ -1062,7 +1453,7 @@ mod tests { async fn calls_endpoints_follow_service_declaration() { let ctx = trace_table_ctx(); let trace = ctx.table("opentelemetry_traces").await.unwrap(); - let plan = build_calls_plan( + let plan = build_relationships_plan_for_test( vec![CallsSource { service: trace_service_decl(&[SERVICE_NAME_COLUMN, "service_namespace"]), scan: trace, @@ -1087,4 +1478,487 @@ mod tests { vec!["service_name=cart,service_namespace=ns1"] ); } + + fn build_relationships_plan_for_test( + traces: Vec, + window: &GraphQueryWindow, + ) -> DfResult> { + build_relationships_plan(traces, None, window) + } + + /// A `TableInfo` derived from the canonical declared-edge definition, with + /// injection points for each aspect the schema matcher must check. + fn declared_table_info( + mutate_columns: impl FnOnce(&mut Vec), + engine: &str, + primary_key_indices: Vec, + extra_options: &[(&str, &str)], + ) -> table::metadata::TableInfo { + let canonical = build_declared_relationships_expr("greptime"); + let mut columns: Vec = canonical + .column_defs + .iter() + .map(|def| { + let wrapper = api::helper::ColumnDataTypeWrapper::try_new( + def.data_type, + def.datatype_extension.clone(), + ) + .unwrap(); + datatypes::schema::ColumnSchema::new( + &def.name, + datatypes::prelude::ConcreteDataType::from(wrapper), + def.is_nullable, + ) + .with_time_index(def.name == OBSERVED_AT_COLUMN) + }) + .collect(); + mutate_columns(&mut columns); + let schema = Arc::new( + datatypes::schema::SchemaBuilder::try_from_columns(columns) + .unwrap() + .build() + .unwrap(), + ); + let options = table::requests::TableOptions { + extra_options: extra_options + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(), + ..Default::default() + }; + let meta = table::metadata::TableMeta { + schema, + primary_key_indices, + value_indices: vec![], + engine: engine.to_string(), + next_column_id: 1, + options, + created_on: Default::default(), + updated_on: Default::default(), + partition_key_indices: vec![], + column_ids: vec![], + }; + table::metadata::TableInfoBuilder::default() + .table_id(1) + .name(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) + .catalog_name("greptime") + .schema_name(DEFAULT_PRIVATE_SCHEMA_NAME) + .table_version(0) + .table_type(table::metadata::TableType::Base) + .meta(meta) + .build() + .unwrap() + } + + #[test] + fn declared_schema_match_checks_more_than_columns() { + let mito = common_catalog::consts::MITO_ENGINE; + let canonical_pk = || (6..14).collect::>(); + + let ok = declared_table_info(|_| {}, mito, canonical_pk(), &[(TTL_KEY, "180d")]); + assert!(declared_relationships_schema_matches(&ok)); + // The canonical expr must declare attributes as json, matching the + // computed table's column. + assert_eq!( + ok.meta + .schema + .column_schema_by_name("attributes") + .map(|c| c.data_type.clone()), + Some(datatypes::prelude::ConcreteDataType::json_datatype()) + ); + let ok = declared_table_info( + |_| {}, + mito, + canonical_pk(), + &[(MERGE_MODE_KEY, "last_row")], + ); + assert!(declared_relationships_schema_matches(&ok)); + + let wrong_engine = declared_table_info(|_| {}, "metric", canonical_pk(), &[]); + assert!(!declared_relationships_schema_matches(&wrong_engine)); + + let pk_missing_generation_id = declared_table_info(|_| {}, mito, (6..13).collect(), &[]); + assert!(!declared_relationships_schema_matches( + &pk_missing_generation_id + )); + + let wrong_time_index = declared_table_info( + |columns| { + columns[0] = columns[0].clone().with_time_index(false); + columns[4] = datatypes::schema::ColumnSchema::new( + "valid_from", + datatypes::prelude::ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true); + }, + mito, + canonical_pk(), + &[], + ); + assert!(!declared_relationships_schema_matches(&wrong_time_index)); + + let append_mode = + declared_table_info(|_| {}, mito, canonical_pk(), &[(APPEND_MODE_KEY, "true")]); + assert!(!declared_relationships_schema_matches(&append_mode)); + + let wrong_column_type = declared_table_info( + |columns| { + columns[14] = datatypes::schema::ColumnSchema::new( + "confidence", + datatypes::prelude::ConcreteDataType::int64_datatype(), + true, + ); + }, + mito, + canonical_pk(), + &[], + ); + assert!(!declared_relationships_schema_matches(&wrong_column_type)); + } + + fn ts_values(batch: &RecordBatch, column: usize) -> Vec { + let array = batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + (0..array.len()).map(|i| array.value(i)).collect() + } + + /// Registers a declared-edge table holding: + /// - frontend->db: two revisions (0.5 then 1.0), the newer retiring the edge at 300s; + /// - api->cache: expired at 100s; + /// - agent-1->search: open-ended `provenance = 'agent'` with JSON attributes; + /// - webapp->auth: two revisions, the newer declared at 900s; + /// - batch->queue: declared at 100s but only valid from 700s; + /// - gateway->redis: asserted twice at the same instant as generations g1/g2. + fn declared_table(ctx: &SessionContext) { + let ts_field = |name: &str| { + Field::new( + name, + DataType::Timestamp(TimeUnit::Millisecond, None), + name != OBSERVED_AT_COLUMN, + ) + }; + let schema = Arc::new(Schema::new(vec![ + ts_field(OBSERVED_AT_COLUMN), + ts_field("window_start"), + ts_field("window_end"), + ts_field("fresh_until"), + ts_field("valid_from"), + ts_field("valid_until"), + Field::new("src_type", DataType::Utf8, false), + Field::new("src_id", DataType::Utf8, false), + Field::new("rel_type", DataType::Utf8, false), + Field::new("dst_type", DataType::Utf8, false), + Field::new("dst_id", DataType::Utf8, false), + Field::new("provenance", DataType::Utf8, false), + Field::new("scope", DataType::Utf8, false), + Field::new("generation_id", DataType::Utf8, false), + Field::new("confidence", DataType::Float64, true), + Field::new("request_count", DataType::Int64, true), + Field::new("error_count", DataType::Int64, true), + Field::new("duration_sum", DataType::Float64, true), + Field::new("duration_count", DataType::Int64, true), + Field::new("attributes", DataType::Binary, true), + ])); + let attrs = jsonb::parse_value(br#"{"connection_type":"virtual_node"}"#) + .unwrap() + .to_vec(); + let no_ts = || Arc::new(TimestampMillisecondArray::from(vec![None::; 9])) as ArrayRef; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![ + 100_000, 200_000, 50_000, 90_000, 100_000, 900_000, 100_000, 100_000, 100_000, + ])) as ArrayRef, + no_ts(), // window_start + no_ts(), // window_end + no_ts(), // fresh_until + Arc::new(TimestampMillisecondArray::from(vec![ + None, + None, + None, + None, + None, + None, + Some(700_000), + None, + None, + ])), + Arc::new(TimestampMillisecondArray::from(vec![ + None, + Some(300_000), + Some(100_000), + None, + None, + None, + None, + None, + None, + ])), + Arc::new(StringArray::from(vec![ + "service", "service", "service", "agent", "service", "service", "service", + "service", "service", + ])), + Arc::new(StringArray::from(vec![ + "frontend", "frontend", "api", "agent-1", "webapp", "webapp", "batch", + "gateway", "gateway", + ])), + Arc::new(StringArray::from(vec![ + "depends_on", + "depends_on", + "depends_on", + "uses", + "depends_on", + "depends_on", + "depends_on", + "depends_on", + "depends_on", + ])), + Arc::new(StringArray::from(vec![ + "service", "service", "service", "tool", "service", "service", "service", + "service", "service", + ])), + Arc::new(StringArray::from(vec![ + "db", "db", "cache", "search", "auth", "auth", "queue", "redis", "redis", + ])), + Arc::new(StringArray::from(vec![ + "declared", "declared", "declared", "agent", "declared", "declared", + "declared", "declared", "declared", + ])), + Arc::new(StringArray::from(vec![""; 9])), + Arc::new(StringArray::from(vec![ + "", "", "", "", "", "", "", "g1", "g2", + ])), + Arc::new(Float64Array::from(vec![ + Some(0.5), + Some(1.0), + Some(1.0), + Some(0.8), + Some(0.5), + Some(1.0), + Some(1.0), + Some(0.5), + Some(1.0), + ])), + Arc::new(Int64Array::from(vec![None::; 9])), + Arc::new(Int64Array::from(vec![None::; 9])), + Arc::new(Float64Array::from(vec![None::; 9])), + Arc::new(Int64Array::from(vec![None::; 9])), + Arc::new(BinaryArray::from(vec![ + None, + None, + None, + Some(attrs.as_slice()), + None, + None, + None, + None, + None, + ])), + ], + ) + .unwrap(); + ctx.register_table( + SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, + Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), + ) + .unwrap(); + } + + #[tokio::test] + async fn declared_edges_latest_revision_validity_and_synthetic_time() { + let ctx = SessionContext::new(); + declared_table(&ctx); + let scan = ctx + .table(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) + .await + .unwrap(); + + // Queried window [120s, 600s). + let window = GraphQueryWindow::from_observed(120_000, 600_000); + let plan = build_relationships_plan(vec![], Some(DeclaredSource { scan }), &window) + .unwrap() + .unwrap(); + let names = plan + .schema() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>(); + assert_eq!(names, RELATIONSHIP_COLUMNS); + + // (src_id, observed_at, window_start, window_end, fresh_until, + // provenance, attributes) + type DeclaredRow = (String, i64, i64, i64, i64, String, Option); + + let batches = collect(&ctx, plan).await; + let mut rows: Vec = vec![]; + for batch in &batches { + let observed = ts_values(batch, 0); + let window_start = ts_values(batch, 1); + let window_end = ts_values(batch, 2); + let fresh_until = ts_values(batch, 3); + let src = strings(batch, 5); + let provenance = strings(batch, 9); + let attributes = json_texts(batch, 15); + for i in 0..batch.num_rows() { + rows.push(( + src[i].clone(), + observed[i], + window_start[i], + window_end[i], + fresh_until[i], + provenance[i].clone(), + attributes[i].clone(), + )); + } + } + rows.sort(); + + // api->cache expired before the window; batch->queue is not yet valid; + // frontend->db keeps only the latest revision; webapp->auth keeps the + // first revision (the second postdates the window); gateway->redis + // collapses to one row across generations. + assert_eq!(rows.len(), 4); + + // agent->tool: open-ended validity declared at 90s. The synthesized + // observed_at is clamped into the queried window; window_end and + // fresh_until take the window's upper bound. + assert_eq!( + rows[0], + ( + "agent-1".to_string(), + 120_000, + 90_000, + 600_000, + 600_000, + "agent".to_string(), + Some(r#"{"connection_type":"virtual_node"}"#.to_string()), + ) + ); + + // frontend->db: the latest revision (declared at 200s, retired at 300s) + // supersedes the open-ended first revision. + assert_eq!( + rows[1], + ( + "frontend".to_string(), + 200_000, + 200_000, + 300_000, + 300_000, + "declared".to_string(), + None, + ) + ); + + assert_eq!( + rows[2], + ( + "gateway".to_string(), + 120_000, + 100_000, + 600_000, + 600_000, + "declared".to_string(), + None, + ) + ); + + assert_eq!( + rows[3], + ( + "webapp".to_string(), + 120_000, + 100_000, + 600_000, + 600_000, + "declared".to_string(), + None, + ) + ); + } + + #[tokio::test] + async fn declared_edges_pick_the_revision_as_of_the_window() { + let ctx = SessionContext::new(); + declared_table(&ctx); + let scan = ctx + .table(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) + .await + .unwrap(); + + // A historical window that predates frontend->db's second revision: + // the first revision was the edge's state then and must not be hidden + // by the newer one. + let window = GraphQueryWindow::from_observed(0, 150_000); + let plan = build_relationships_plan(vec![], Some(DeclaredSource { scan }), &window) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let mut rows: Vec<(String, f64)> = vec![]; + for batch in &batches { + let src = strings(batch, 5); + let confidence = batch + .column(10) + .as_any() + .downcast_ref::() + .unwrap(); + for (i, src) in src.into_iter().enumerate() { + rows.push((src, confidence.value(i))); + } + } + rows.sort_by(|a, b| a.0.cmp(&b.0)); + + // api->cache is valid inside this window; batch->queue is not yet. + assert_eq!( + rows, + vec![ + ("agent-1".to_string(), 0.8), + ("api".to_string(), 1.0), + ("frontend".to_string(), 0.5), + ("gateway".to_string(), 1.0), + ("webapp".to_string(), 0.5), + ] + ); + } + + #[tokio::test] + async fn declared_edges_union_calls_branch() { + let ctx = trace_table_ctx(); + declared_table(&ctx); + let trace = ctx.table("opentelemetry_traces").await.unwrap(); + let declared = ctx + .table(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) + .await + .unwrap(); + + let plan = build_relationships_plan( + vec![CallsSource { + service: trace_service_decl(&[SERVICE_NAME_COLUMN]), + scan: trace, + }], + Some(DeclaredSource { scan: declared }), + &test_window(), + ) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let mut provenance: Vec = batches.iter().flat_map(|b| strings(b, 9)).collect(); + provenance.sort(); + // One trace-derived calls edge plus the declared edges valid somewhere + // inside [0s, 600s): frontend, api, webapp, gateway, and the + // agent-asserted one. + assert_eq!( + provenance, + vec![ + "agent", "declared", "declared", "declared", "declared", "trace" + ] + ); + } } diff --git a/src/table/src/predicate.rs b/src/table/src/predicate.rs index 29b10430492..5fb341bc63b 100644 --- a/src/table/src/predicate.rs +++ b/src/table/src/predicate.rs @@ -23,6 +23,7 @@ use datafusion::common::ScalarValue; use datafusion::physical_optimizer::pruning::PruningPredicate; use datafusion_common::ToDFSchema; use datafusion_common::pruning::PruningStatistics; +use datafusion_common::tree_node::TreeNode; use datafusion_expr::expr::{Expr, InList}; use datafusion_expr::{Between, BinaryExpr, Operator}; use datafusion_physical_expr::execution_props::ExecutionProps; @@ -217,6 +218,83 @@ pub fn build_time_range_predicate( res } +/// The outcome of strictly extracting the time range of `ts_col_name` from a +/// scan's filters. Unlike [`build_time_range_predicate`], which quietly widens +/// to `min_to_max` on anything it cannot parse (fine for pruning), this +/// distinguishes "the column is not filtered at all" from "it is filtered in a +/// way that cannot be safely turned into a range" — for callers whose contract +/// forbids silently falling back to a default window. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum TimeRangeExtraction { + /// No filter references the column. + Absent, + /// Every filter referencing the column was folded into this range, which + /// over-approximates the filters' satisfying set. + Extracted(TimestampRange), + /// At least one filter references the column in a shape that cannot be + /// safely extracted (e.g. under `OR`/`NOT`, or compared to a non-literal). + Unsupported, +} + +/// Strictly extracts the time range of `ts_col_name` from the (implicitly +/// AND-ed) scan filters. See [`TimeRangeExtraction`]. +pub fn extract_time_range_strict( + ts_col_name: &str, + ts_col_unit: TimeUnit, + filters: &[Expr], +) -> TimeRangeExtraction { + let mut range: Option = None; + for expr in filters { + if !expr + .column_refs() + .iter() + .any(|column| column.name == ts_col_name) + { + continue; + } + // Disjunctive shapes (`OR`, `IN`) collapse disjoint ranges into their + // convex hull, which is not exactly representable as one contiguous + // range; callers deriving synthetic timestamps from the bounds would + // emit points inside the gaps. + if contains_disjunction_over_column(expr, ts_col_name) { + return TimeRangeExtraction::Unsupported; + } + // `Some` from the lenient extractor over-approximates the expression's + // satisfying set (an `AND` side it cannot parse is dropped, which only + // widens), so intersecting extracted conjuncts stays an + // over-approximation. + match extract_time_range_from_expr(ts_col_name, ts_col_unit, expr) { + Some(extracted) => { + range = Some(match range { + Some(acc) => acc.and(&extracted), + None => extracted, + }); + } + None => return TimeRangeExtraction::Unsupported, + } + } + match range { + Some(range) => TimeRangeExtraction::Extracted(range), + None => TimeRangeExtraction::Absent, + } +} + +fn contains_disjunction_over_column(expr: &Expr, ts_col_name: &str) -> bool { + let is_disjunction = |expr: &Expr| { + matches!( + expr, + Expr::BinaryExpr(BinaryExpr { + op: Operator::Or, + .. + }) | Expr::InList(_) + ) && expr + .column_refs() + .iter() + .any(|column| column.name == ts_col_name) + }; + expr.exists(|expr| Ok(is_disjunction(expr))).unwrap_or(true) +} + /// Extract time range filter from `WHERE`/`IN (...)`/`BETWEEN` clauses. /// Return None if no time range can be found in expr. pub fn extract_time_range_from_expr( @@ -593,6 +671,99 @@ mod tests { ); } + #[test] + fn test_extract_time_range_strict() { + fn ts_lit(ms: i64) -> Expr { + lit(ScalarValue::TimestampMillisecond(Some(ms), None)) + } + let extract = + |filters: &[Expr]| extract_time_range_strict("ts", TimeUnit::Millisecond, filters); + let range = |start: i64, end: i64| { + TimestampRange::new( + Timestamp::new_millisecond(start), + Timestamp::new_millisecond(end), + ) + .unwrap() + }; + + // No filter references the column. + assert_eq!(extract(&[]), TimeRangeExtraction::Absent); + assert_eq!( + extract(&[col("host").eq(lit("a"))]), + TimeRangeExtraction::Absent + ); + + // Both bounds across conjuncts; unrelated filters are ignored. + assert_eq!( + extract(&[ + col("ts").gt_eq(ts_lit(1000)), + col("ts").lt(ts_lit(2000)), + col("host").eq(lit("a")), + ]), + TimeRangeExtraction::Extracted(range(1000, 2000)) + ); + + // Lower bound only / upper bound only. + assert_eq!( + extract(&[col("ts").gt_eq(ts_lit(1000))]), + TimeRangeExtraction::Extracted(TimestampRange::from_start(Timestamp::new_millisecond( + 1000 + ))) + ); + assert_eq!( + extract(&[col("ts").lt(ts_lit(2000))]), + TimeRangeExtraction::Extracted(TimestampRange::until_end( + Timestamp::new_millisecond(2000), + false + )) + ); + + // BETWEEN is inclusive on both ends. + assert_eq!( + extract(&[col("ts").between(ts_lit(1000), ts_lit(2000))]), + TimeRangeExtraction::Extracted(range(1000, 2001)) + ); + + // Equality pins a single point. + assert_eq!( + extract(&[col("ts").eq(ts_lit(1500))]), + TimeRangeExtraction::Extracted(TimestampRange::single(Timestamp::new_millisecond( + 1500 + ))) + ); + + // An unparsable side under AND only widens; the extraction stays safe. + assert_eq!( + extract(&[col("ts").gt_eq(ts_lit(1000)).and(col("ts").lt(col("t2")))]), + TimeRangeExtraction::Extracted(TimestampRange::from_start(Timestamp::new_millisecond( + 1000 + ))) + ); + + // Contradictory bounds collapse to the empty range, not an error. + let TimeRangeExtraction::Extracted(empty) = + extract(&[col("ts").gt_eq(ts_lit(2000)), col("ts").lt(ts_lit(1000))]) + else { + panic!("expected an extraction"); + }; + assert!(empty.is_empty()); + + // Shapes that could widen the satisfying set beyond what is extractable + // must be refused: OR with an unparsable side, NOT, non-literal bounds. + assert_eq!( + extract(&[col("ts").gt(ts_lit(1000)).or(col("host").eq(lit("a")))]), + TimeRangeExtraction::Unsupported + ); + assert_eq!( + extract(&[!col("ts").gt(ts_lit(1000))]), + TimeRangeExtraction::Unsupported + ); + assert_eq!( + extract(&[col("ts").gt_eq(col("t2"))]), + TimeRangeExtraction::Unsupported + ); + } + async fn gen_test_parquet_file(dir: &TempDir, cnt: usize) -> (String, Arc) { let path = dir .path() diff --git a/tests-integration/src/grpc.rs b/tests-integration/src/grpc.rs index 7c2148e130e..904cd0f1116 100644 --- a/tests-integration/src/grpc.rs +++ b/tests-integration/src/grpc.rs @@ -63,7 +63,10 @@ mod test { }; use client::OutputData; use common_base::Plugins; - use common_catalog::consts::MITO_ENGINE; + use common_catalog::consts::{ + DEFAULT_CATALOG_NAME, DEFAULT_PRIVATE_SCHEMA_NAME, MITO_ENGINE, + SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, + }; use common_meta::rpc::router::region_distribution; use common_query::Output; use common_query::logical_plan::breakup_insert_plan; @@ -1342,4 +1345,171 @@ CREATE TABLE {table_name} ( +--------------------+---------------------------------------------------+"#; execute_sql_and_expect(&frontend, sql, expected).await; } + + fn declared_relationships_row_insert() -> Request { + use api::v1::value::ValueData; + use api::v1::{ColumnSchema, Row, RowInsertRequest, RowInsertRequests, Rows, Value}; + + let string_column = |name: &str, semantic: SemanticType| ColumnSchema { + column_name: name.to_string(), + datatype: ColumnDataType::String as i32, + semantic_type: semantic as i32, + ..Default::default() + }; + let string_value = |value: &str| Value { + value_data: Some(ValueData::StringValue(value.to_string())), + }; + let mut schema = vec![ColumnSchema { + column_name: "observed_at".to_string(), + datatype: ColumnDataType::TimestampMillisecond as i32, + semantic_type: SemanticType::Timestamp as i32, + ..Default::default() + }]; + let mut values = vec![Value { + value_data: Some(ValueData::TimestampMillisecondValue(1000)), + }]; + for (tag, value) in [ + ("src_type", "service"), + ("src_id", "frontend"), + ("rel_type", "depends_on"), + ("dst_type", "service"), + ("dst_id", "users-db"), + ("provenance", "declared"), + ("scope", ""), + ("generation_id", ""), + ] { + schema.push(string_column(tag, SemanticType::Tag)); + values.push(string_value(value)); + } + + Request::RowInserts(RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME.to_string(), + rows: Some(Rows { + schema, + rows: vec![Row { values }], + }), + }], + }) + } + + async fn assert_declared_relationships_table_is_canonical(instance: &Instance) { + let table = instance + .catalog_manager() + .table( + DEFAULT_CATALOG_NAME, + DEFAULT_PRIVATE_SCHEMA_NAME, + SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, + None, + ) + .await + .unwrap() + .expect("declared-edge table must exist"); + let table_info = table.table_info(); + assert_eq!(table_info.meta.schema.column_schemas().len(), 20); + assert_eq!( + table_info + .meta + .schema + .timestamp_column() + .map(|c| c.name.as_str()), + Some("observed_at") + ); + let primary_keys: Vec<&str> = table_info + .meta + .primary_key_indices + .iter() + .map(|idx| table_info.meta.schema.column_schemas()[*idx].name.as_str()) + .collect(); + assert_eq!( + primary_keys, + [ + "src_type", + "src_id", + "rel_type", + "dst_type", + "dst_id", + "provenance", + "scope", + "generation_id", + ] + ); + } + + /// The declared-edge table's first gRPC write must create it with the + /// canonical definition (never the request's shape), entering below the + /// user-DDL guard that rejects generic CREATE against the reserved name. + #[tokio::test(flavor = "multi_thread")] + async fn test_grpc_first_write_creates_declared_relationships_table() { + common_telemetry::init_default_ut_logging(); + let standalone = GreptimeDbStandaloneBuilder::new("test_grpc_declared_first_write") + .build() + .await; + let instance = standalone.fe_instance(); + + let ctx = Arc::new(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_PRIVATE_SCHEMA_NAME, + )); + let output = + GrpcQueryHandler::do_query(instance.as_ref(), declared_relationships_row_insert(), ctx) + .await + .unwrap(); + assert!(matches!(output.data, OutputData::AffectedRows(1))); + + assert_declared_relationships_table_is_canonical(instance).await; + } + + /// The system-defined table is created regardless of the auto-create hint: + /// its creation is a system action, not user auto-create. + #[tokio::test(flavor = "multi_thread")] + async fn test_grpc_declared_relationships_bypasses_auto_create_hint() { + use api::v1::value::ValueData; + use api::v1::{ColumnSchema, Row, RowInsertRequest, RowInsertRequests, Rows, Value}; + + common_telemetry::init_default_ut_logging(); + let standalone = GreptimeDbStandaloneBuilder::new("test_grpc_declared_no_auto_create") + .build() + .await; + let instance = standalone.fe_instance(); + + let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_PRIVATE_SCHEMA_NAME); + ctx.set_extension(table::requests::AUTO_CREATE_TABLE_KEY, "false"); + let ctx = Arc::new(ctx); + + let output = GrpcQueryHandler::do_query( + instance.as_ref(), + declared_relationships_row_insert(), + ctx.clone(), + ) + .await + .unwrap(); + assert!(matches!(output.data, OutputData::AffectedRows(1))); + assert_declared_relationships_table_is_canonical(instance).await; + + // An ordinary table stays subject to the hint. + let request = Request::RowInserts(RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "ordinary_table".to_string(), + rows: Some(Rows { + schema: vec![ColumnSchema { + column_name: "ts".to_string(), + datatype: ColumnDataType::TimestampMillisecond as i32, + semantic_type: SemanticType::Timestamp as i32, + ..Default::default() + }], + rows: vec![Row { + values: vec![Value { + value_data: Some(ValueData::TimestampMillisecondValue(1000)), + }], + }], + }), + }], + }); + let err = GrpcQueryHandler::do_query(instance.as_ref(), request, ctx) + .await + .unwrap_err(); + let msg = common_error::ext::ErrorExt::output_msg(&err); + assert!(msg.contains("auto_create_table"), "unexpected error: {msg}"); + } } diff --git a/tests-integration/src/instance.rs b/tests-integration/src/instance.rs index 17435a3b7a0..547f64436f6 100644 --- a/tests-integration/src/instance.rs +++ b/tests-integration/src/instance.rs @@ -536,4 +536,108 @@ mod tests { unreachable!(); } } + + /// The entity-graph derivation must run as the caller: a source table the + /// caller cannot read contributes neither entities nor edges. + #[tokio::test(flavor = "multi_thread")] + async fn test_entity_graph_derivation_runs_as_caller() { + common_telemetry::init_default_ut_logging(); + + /// Denies the entity-graph derivation on `secret_traces` only. + struct DenySecretTraces; + + impl auth::PermissionChecker for DenySecretTraces { + fn check_permission( + &self, + _user_info: auth::UserInfoRef, + _req: auth::PermissionReq, + ) -> auth::error::Result { + Ok(auth::PermissionResp::Allow) + } + + fn check_permission_with_table_targets( + &self, + _user_info: auth::UserInfoRef, + req: auth::PermissionReq, + targets: auth::PermissionTableTargets, + ) -> auth::error::Result { + if matches!(req, auth::PermissionReq::Action(auth::SEMANTIC_GRAPH_QUERY)) + && let auth::PermissionTableTargets::Resolved(targets) = targets + && targets.iter().any(|target| target.table == "secret_traces") + { + return Ok(auth::PermissionResp::Reject); + } + Ok(auth::PermissionResp::Allow) + } + } + + let plugins = Plugins::new(); + plugins.insert::(Arc::new(DenySecretTraces)); + + let standalone = GreptimeDbStandaloneBuilder::new("test_entity_graph_runs_as_caller") + .with_plugin(plugins) + .build() + .await; + let instance = standalone.fe_instance().clone(); + + for table in ["open_traces", "secret_traces"] { + let create = format!( + r#"create table {table} ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + primary key (service_name) + ) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true')"# + ); + create_table(&instance, &create).await; + let insert = format!( + "insert into {table} values \ + (now(), 't0', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'client-{table}', 0), \ + (now(), 't0', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'server-{table}', 100)" + ); + let output = query(&instance, &insert).await; + let OutputData::AffectedRows(x) = output.data else { + unreachable!() + }; + assert_eq!(x, 2); + } + + let sql = "select src_id, dst_id from greptime_private.semantic_relationships \ + order by src_id"; + let output = query(&instance, sql).await; + let OutputData::Stream(s) = output.data else { + unreachable!() + }; + let batches = common_recordbatch::util::collect_batches(s).await.unwrap(); + let pretty_print = batches.pretty_print().unwrap(); + assert!( + pretty_print.contains("client-open_traces"), + "allowed edge missing:\n{pretty_print}" + ); + assert!( + !pretty_print.contains("secret_traces"), + "denied source leaked into edges:\n{pretty_print}" + ); + + let sql = "select entity_id from greptime_private.semantic_entities order by entity_id"; + let output = query(&instance, sql).await; + let OutputData::Stream(s) = output.data else { + unreachable!() + }; + let batches = common_recordbatch::util::collect_batches(s).await.unwrap(); + let pretty_print = batches.pretty_print().unwrap(); + assert!( + pretty_print.contains("client-open_traces"), + "allowed entity missing:\n{pretty_print}" + ); + assert!( + !pretty_print.contains("secret_traces"), + "denied source leaked into entities:\n{pretty_print}" + ); + } } diff --git a/tests/cases/standalone/common/system/semantic_graph.result b/tests/cases/standalone/common/system/semantic_graph.result index bda28755f1e..73f00f57c6f 100644 --- a/tests/cases/standalone/common/system/semantic_graph.result +++ b/tests/cases/standalone/common/system/semantic_graph.result @@ -153,3 +153,127 @@ drop table graph_traces; Affected Rows: 0 +-- Declared edges: the physical table is created by the system on the first +-- INSERT with its canonical schema. Re-asserting an edge stores a new revision; +-- reads keep only the latest one per edge key. +insert into greptime_private.semantic_relationships_declared + (observed_at, src_type, src_id, rel_type, dst_type, dst_id, provenance, scope, generation_id, confidence) +values (now() - interval '10 minute', 'service', 'frontend', 'depends_on', 'service', 'users-db', 'declared', '', '', 0.5); + +Affected Rows: 1 + +insert into greptime_private.semantic_relationships_declared + (observed_at, src_type, src_id, rel_type, dst_type, dst_id, provenance, scope, generation_id, confidence) +values (now() - interval '5 minute', 'service', 'frontend', 'depends_on', 'service', 'users-db', 'declared', '', '', 1.0); + +Affected Rows: 1 + +-- An old open-ended declaration stays valid until its row expires; an edge +-- retired in the past must not appear. +insert into greptime_private.semantic_relationships_declared + (observed_at, src_type, src_id, rel_type, dst_type, dst_id, provenance, scope, generation_id, valid_until) +values + ('1970-01-01 00:00:01', 'service', 'legacy', 'depends_on', 'service', 'mainframe', 'declared', '', '', NULL), + ('1970-01-01 00:00:01', 'service', 'retired', 'depends_on', 'service', 'oldsys', 'declared', '', '', '2000-01-01 00:00:00'); + +Affected Rows: 2 + +select src_id, dst_id, rel_type, provenance, confidence +from greptime_private.semantic_relationships +order by src_id; + ++----------+-----------+------------+------------+------------+ +| src_id | dst_id | rel_type | provenance | confidence | ++----------+-----------+------------+------------+------------+ +| frontend | users-db | depends_on | declared | 1.0 | +| legacy | mainframe | depends_on | declared | | ++----------+-----------+------------+------------+------------+ + +-- An explicit observed_at window replaces the default last hour; the emitted +-- timestamps of declared edges are synthesized inside the queried window. +select observed_at, window_start, fresh_until, src_id, dst_id +from greptime_private.semantic_relationships +where observed_at >= '2001-01-01 00:00:00' and observed_at < '2001-01-02 00:00:00' +order by src_id; + ++---------------------+---------------------+---------------------+--------+-----------+ +| observed_at | window_start | fresh_until | src_id | dst_id | ++---------------------+---------------------+---------------------+--------+-----------+ +| 2001-01-01T00:00:00 | 1970-01-01T00:00:01 | 2001-01-02T00:00:00 | legacy | mainframe | ++---------------------+---------------------+---------------------+--------+-----------+ + +-- A lower bound alone is fine (the upper bound defaults to now)... +select src_id, dst_id, provenance +from greptime_private.semantic_relationships +where observed_at >= now() - interval '30 minute' +order by src_id; + ++----------+-----------+------------+ +| src_id | dst_id | provenance | ++----------+-----------+------------+ +| frontend | users-db | declared | +| legacy | mainframe | declared | ++----------+-----------+------------+ + +-- ...but an upper bound alone would scan unbounded history: explicit error. +select src_id from greptime_private.semantic_relationships +where observed_at < '2001-01-02 00:00:00'; + +Error: 3001(EngineExecuteQuery), Invalid SQL, error: the observed_at filter has no lower bound; the graph cannot derive over unbounded history — add e.g. observed_at >= '2026-01-01 00:00:00' + +-- The declared-edge table's definition is system-owned: user CREATE/ALTER are +-- rejected and a physical table cannot be renamed into the name, while plain +-- DML stays allowed. +create table greptime_private.semantic_relationships_declared (ts timestamp time index); + +Error: 1001(Unsupported), The definition of table `semantic_relationships_declared` is managed by GreptimeDB; it cannot be created or altered (DROP recreates it on the next write) + +alter table greptime_private.semantic_relationships_declared add column extra string; + +Error: 1001(Unsupported), The definition of table `semantic_relationships_declared` is managed by GreptimeDB; it cannot be created or altered (DROP recreates it on the next write) + +create table greptime_private.declared_rename_probe (ts timestamp time index); + +Affected Rows: 0 + +alter table greptime_private.declared_rename_probe rename semantic_relationships_declared; + +Error: 1001(Unsupported), The definition of table `semantic_relationships_declared` is managed by GreptimeDB; it cannot be created or altered (DROP recreates it on the next write) + +drop table greptime_private.declared_rename_probe; + +Affected Rows: 0 + +delete from greptime_private.semantic_relationships_declared; + +Affected Rows: 4 + +select src_id from greptime_private.semantic_relationships order by src_id; + +++ +++ + +-- DROP is allowed (nothing structural is lost: the next INSERT recreates the +-- canonical table) and cleans up after this test. +drop table greptime_private.semantic_relationships_declared; + +Affected Rows: 0 + +insert into greptime_private.semantic_relationships_declared + (observed_at, src_type, src_id, rel_type, dst_type, dst_id, provenance, scope, generation_id) +values (now(), 'service', 'reborn', 'depends_on', 'service', 'db', 'declared', '', ''); + +Affected Rows: 1 + +select src_id from greptime_private.semantic_relationships order by src_id; + ++--------+ +| src_id | ++--------+ +| reborn | ++--------+ + +drop table greptime_private.semantic_relationships_declared; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/system/semantic_graph.sql b/tests/cases/standalone/common/system/semantic_graph.sql index 000401e1666..c5a93d6c6b7 100644 --- a/tests/cases/standalone/common/system/semantic_graph.sql +++ b/tests/cases/standalone/common/system/semantic_graph.sql @@ -92,3 +92,72 @@ from greptime_private.semantic_entities order by entity_id; drop table graph_traces; + +-- Declared edges: the physical table is created by the system on the first +-- INSERT with its canonical schema. Re-asserting an edge stores a new revision; +-- reads keep only the latest one per edge key. +insert into greptime_private.semantic_relationships_declared + (observed_at, src_type, src_id, rel_type, dst_type, dst_id, provenance, scope, generation_id, confidence) +values (now() - interval '10 minute', 'service', 'frontend', 'depends_on', 'service', 'users-db', 'declared', '', '', 0.5); + +insert into greptime_private.semantic_relationships_declared + (observed_at, src_type, src_id, rel_type, dst_type, dst_id, provenance, scope, generation_id, confidence) +values (now() - interval '5 minute', 'service', 'frontend', 'depends_on', 'service', 'users-db', 'declared', '', '', 1.0); + +-- An old open-ended declaration stays valid until its row expires; an edge +-- retired in the past must not appear. +insert into greptime_private.semantic_relationships_declared + (observed_at, src_type, src_id, rel_type, dst_type, dst_id, provenance, scope, generation_id, valid_until) +values + ('1970-01-01 00:00:01', 'service', 'legacy', 'depends_on', 'service', 'mainframe', 'declared', '', '', NULL), + ('1970-01-01 00:00:01', 'service', 'retired', 'depends_on', 'service', 'oldsys', 'declared', '', '', '2000-01-01 00:00:00'); + +select src_id, dst_id, rel_type, provenance, confidence +from greptime_private.semantic_relationships +order by src_id; + +-- An explicit observed_at window replaces the default last hour; the emitted +-- timestamps of declared edges are synthesized inside the queried window. +select observed_at, window_start, fresh_until, src_id, dst_id +from greptime_private.semantic_relationships +where observed_at >= '2001-01-01 00:00:00' and observed_at < '2001-01-02 00:00:00' +order by src_id; + +-- A lower bound alone is fine (the upper bound defaults to now)... +select src_id, dst_id, provenance +from greptime_private.semantic_relationships +where observed_at >= now() - interval '30 minute' +order by src_id; + +-- ...but an upper bound alone would scan unbounded history: explicit error. +select src_id from greptime_private.semantic_relationships +where observed_at < '2001-01-02 00:00:00'; + +-- The declared-edge table's definition is system-owned: user CREATE/ALTER are +-- rejected and a physical table cannot be renamed into the name, while plain +-- DML stays allowed. +create table greptime_private.semantic_relationships_declared (ts timestamp time index); + +alter table greptime_private.semantic_relationships_declared add column extra string; + +create table greptime_private.declared_rename_probe (ts timestamp time index); + +alter table greptime_private.declared_rename_probe rename semantic_relationships_declared; + +drop table greptime_private.declared_rename_probe; + +delete from greptime_private.semantic_relationships_declared; + +select src_id from greptime_private.semantic_relationships order by src_id; + +-- DROP is allowed (nothing structural is lost: the next INSERT recreates the +-- canonical table) and cleans up after this test. +drop table greptime_private.semantic_relationships_declared; + +insert into greptime_private.semantic_relationships_declared + (observed_at, src_type, src_id, rel_type, dst_type, dst_id, provenance, scope, generation_id) +values (now(), 'service', 'reborn', 'depends_on', 'service', 'db', 'declared', '', ''); + +select src_id from greptime_private.semantic_relationships order by src_id; + +drop table greptime_private.semantic_relationships_declared; From 78084a9d4485ee61ff03e071f6428fc57c906332 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Mon, 10 Aug 2026 20:19:27 +0800 Subject: [PATCH 12/92] feat: add admin function to discard unflushed data (#8768) * feat: add admin function to discard unflushed data Signed-off-by: evenyag * test: cover discarding unflushed data by table Signed-off-by: evenyag * chore: fix license header Signed-off-by: evenyag * fix: reject discarding logical metric table data Signed-off-by: evenyag * refactor: defer table name formatting in error paths Signed-off-by: evenyag * chore(deps): update greptime-proto revision Signed-off-by: evenyag * refactor: rename discard unflushed admin function Signed-off-by: evenyag --------- Signed-off-by: evenyag --- AGENTS.md | 21 +- Cargo.lock | 2 +- Cargo.toml | 2 +- src/common/function/src/admin.rs | 4 +- .../src/admin/discard_unflushed_data.rs | 232 ++++++++++++++++++ src/common/function/src/handlers.rs | 14 ++ src/common/function/src/state.rs | 17 ++ src/operator/src/request.rs | 96 +++++++- src/operator/src/table.rs | 25 ++ src/store-api/src/region_request.rs | 39 +++ tests-integration/src/tests/instance_test.rs | 75 ++++++ .../admin/discard_unflushed_data.result | 133 ++++++++++ .../function/admin/discard_unflushed_data.sql | 62 +++++ 13 files changed, 709 insertions(+), 13 deletions(-) create mode 100644 src/common/function/src/admin/discard_unflushed_data.rs create mode 100644 tests/cases/standalone/common/function/admin/discard_unflushed_data.result create mode 100644 tests/cases/standalone/common/function/admin/discard_unflushed_data.sql diff --git a/AGENTS.md b/AGENTS.md index c99856c2941..70f2ae20b41 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -114,22 +114,25 @@ blast radius requires it. ## Before opening a PR -1. `make fmt` -2. `make clippy` -3. `make test` -4. `make check-udeps` (run `make fix-udeps` if it reports unused dependencies). -5. If you added or changed a public configuration option, update the applicable +1. If you added a `.rs`, `.py`, or `.ts` file, apply and verify its license + header with `hawkeye format` followed by `hawkeye check`. Use the inception + year from `licenserc.toml`, not the current year. +2. `make fmt` +3. `make clippy` +4. `make test` +5. `make check-udeps` (run `make fix-udeps` if it reports unused dependencies). +6. If you added or changed a public configuration option, update the applicable example TOMLs, configuration-loading and serialized-config snapshot tests, and related user-facing documentation. Run `make config-docs` (needs Docker) and commit the regenerated `config/config.md`. -6. If you changed a persisted or wire format, add a compatibility test case (see +7. If you changed a persisted or wire format, add a compatibility test case (see `.agents/architecture-invariants.md`). -7. If you added or gated an enterprise-only file, give it the enterprise license +8. If you added or gated an enterprise-only file, give it the enterprise license header, list it in `licenserc-enterprise.toml` (`includes`) and `licenserc.toml` (`excludes`), and run `make check-enterprise-license`. -8. Use a conventional-commit title, sign off commits (`git commit -s`), and sign +9. Use a conventional-commit title, sign off commits (`git commit -s`), and sign the CLA. -9. When creating or updating a pull request, follow +10. When creating or updating a pull request, follow [`.github/pull_request_template.md`](.github/pull_request_template.md): include the CLA statement, fill the change-intention section with enough detail, and update checklist items accurately. diff --git a/Cargo.lock b/Cargo.lock index e2210b80bfe..d6385dd7657 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6032,7 +6032,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=032510ded061277b7d1bbb29d4046abb5b1cbf4b#032510ded061277b7d1bbb29d4046abb5b1cbf4b" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=e7be20ff855b7522efe235383d82238d56e676c3#e7be20ff855b7522efe235383d82238d56e676c3" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", diff --git a/Cargo.toml b/Cargo.toml index 502e652ffc8..64e6dba2f7f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -159,7 +159,7 @@ fs2 = "0.4" fst = "0.4.7" futures = "0.3" futures-util = "0.3" -greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "032510ded061277b7d1bbb29d4046abb5b1cbf4b" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "e7be20ff855b7522efe235383d82238d56e676c3" } hex = "0.4" http = "1" humantime = "2.1" diff --git a/src/common/function/src/admin.rs b/src/common/function/src/admin.rs index 73bed86f7ad..ff8172ebb5a 100644 --- a/src/common/function/src/admin.rs +++ b/src/common/function/src/admin.rs @@ -13,6 +13,7 @@ // limitations under the License. mod build_index_table; +mod discard_unflushed_data; mod flush_compact_region; mod flush_compact_table; mod gc; @@ -34,6 +35,7 @@ use reconcile_database::ReconcileDatabaseFunction; use reconcile_table::ReconcileTableFunction; use crate::admin::build_index_table::BuildIndexFunction; +use crate::admin::discard_unflushed_data::DiscardUnflushedDataFunction; use crate::flush_flow::FlushFlowFunction; use crate::function_registry::FunctionRegistry; @@ -58,8 +60,8 @@ impl AdminFunction { } /// Register functions that must only be resolved by an ADMIN statement. - #[cfg_attr(not(feature = "enterprise"), allow(unused_variables))] pub fn register_admin_only(registry: &FunctionRegistry) { + registry.register(DiscardUnflushedDataFunction::factory()); #[cfg(feature = "enterprise")] registry.register(PurgeTableFunction::factory()); } diff --git a/src/common/function/src/admin/discard_unflushed_data.rs b/src/common/function/src/admin/discard_unflushed_data.rs new file mode 100644 index 00000000000..b9f46f62f96 --- /dev/null +++ b/src/common/function/src/admin/discard_unflushed_data.rs @@ -0,0 +1,232 @@ +// Copyright 2023 Greptime Team +// +// 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. + +use arrow::datatypes::DataType as ArrowDataType; +use common_error::ext::BoxedError; +use common_macro::admin_fn; +use common_query::error::{ + InvalidFuncArgsSnafu, MissingTableMutationHandlerSnafu, Result, TableMutationSnafu, + UnsupportedInputDataTypeSnafu, +}; +use datafusion_expr::{Signature, TypeSignature, Volatility}; +use datatypes::data_type::DataType; +use datatypes::prelude::*; +use session::context::QueryContextRef; +use session::table_name::table_name_to_full_name; +use snafu::{ResultExt, ensure}; +use store_api::storage::RegionId; +use table::table_name::TableName; + +use crate::handlers::TableMutationHandlerRef; +use crate::helper::cast_u64; + +/// Discards all unflushed data from a region. +#[admin_fn( + name = DiscardUnflushedDataFunction, + display_name = discard_unflushed, + sig_fn = signature, + ret = uint64, + single_row +)] +pub(crate) async fn discard_unflushed_data( + table_mutation_handler: &TableMutationHandlerRef, + query_ctx: &QueryContextRef, + params: &[ValueRef<'_>], +) -> Result { + ensure!( + params.len() == 1, + InvalidFuncArgsSnafu { + err_msg: format!( + "The length of the args is not correct, expect 1, have: {}", + params.len() + ), + } + ); + + let affected_rows = match params[0] { + ValueRef::String(table_name) => { + let (catalog_name, schema_name, table_name) = + table_name_to_full_name(table_name, query_ctx) + .map_err(BoxedError::new) + .context(TableMutationSnafu)?; + table_mutation_handler + .discard_unflushed_data_by_table( + TableName::new(catalog_name, schema_name, table_name), + query_ctx.clone(), + ) + .await? + } + _ => { + let Some(region_id) = cast_u64(¶ms[0])? else { + return UnsupportedInputDataTypeSnafu { + function: "discard_unflushed", + datatypes: params + .iter() + .map(|value| value.data_type()) + .collect::>(), + } + .fail(); + }; + table_mutation_handler + .discard_unflushed_data(RegionId::from_u64(region_id), query_ctx.clone()) + .await? + } + }; + + Ok(Value::from(affected_rows as u64)) +} + +fn signature() -> Signature { + Signature::one_of( + vec![ + TypeSignature::Uniform( + 1, + ConcreteDataType::numerics() + .into_iter() + .map(|data_type| data_type.as_arrow_type()) + .collect(), + ), + TypeSignature::Exact(vec![ArrowDataType::Utf8]), + ], + Volatility::Immutable, + ) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{StringArray, UInt64Array}; + use arrow::datatypes::{DataType, Field}; + use datafusion_expr::{ColumnarValue, TypeSignature}; + + use super::*; + use crate::function::FunctionContext; + use crate::function_factory::ScalarFunctionFactory; + use crate::function_registry::{FUNCTION_REGISTRY, get_admin_function}; + + #[test] + fn test_discard_unflushed_data_is_admin_only() { + assert!(get_admin_function("discard_unflushed").is_some()); + assert!( + FUNCTION_REGISTRY + .get_function("discard_unflushed") + .is_none() + ); + } + + #[test] + fn test_discard_unflushed_data_signature() { + let factory: ScalarFunctionFactory = DiscardUnflushedDataFunction::factory().into(); + let function = factory.provide(FunctionContext::mock()); + + assert_eq!("discard_unflushed", function.name()); + assert_eq!(DataType::UInt64, function.return_type(&[]).unwrap()); + assert!(matches!( + function.signature(), + Signature { + type_signature: TypeSignature::OneOf(valid_types), + volatility: Volatility::Immutable, + .. + } if valid_types == &vec![ + TypeSignature::Uniform( + 1, + ConcreteDataType::numerics() + .into_iter() + .map(|data_type| { + use datatypes::data_type::DataType; + data_type.as_arrow_type() + }) + .collect::>(), + ), + TypeSignature::Exact(vec![DataType::Utf8]), + ] + )); + } + + #[tokio::test] + async fn test_discard_unflushed_data() { + let factory: ScalarFunctionFactory = DiscardUnflushedDataFunction::factory().into(); + let function = factory.provide(FunctionContext::mock()); + let args = datafusion::logical_expr::ScalarFunctionArgs { + args: vec![ColumnarValue::Array(Arc::new(UInt64Array::from(vec![99])))], + arg_fields: vec![Arc::new(Field::new("arg_0", DataType::UInt64, false))], + return_field: Arc::new(Field::new("result", DataType::UInt64, false)), + number_rows: 1, + config_options: Arc::new(datafusion_common::config::ConfigOptions::default()), + }; + + let result = function + .as_async() + .unwrap() + .invoke_async_with_args(args) + .await + .unwrap(); + let ColumnarValue::Array(array) = result else { + panic!("expected array output"); + }; + let array = array.as_any().downcast_ref::().unwrap(); + assert_eq!(42, array.value(0)); + } + + #[tokio::test] + async fn test_discard_unflushed_data_by_table() { + let factory: ScalarFunctionFactory = DiscardUnflushedDataFunction::factory().into(); + let function = factory.provide(FunctionContext::mock()); + let args = datafusion::logical_expr::ScalarFunctionArgs { + args: vec![ColumnarValue::Array(Arc::new(StringArray::from(vec![ + "my_table", + ])))], + arg_fields: vec![Arc::new(Field::new("arg_0", DataType::Utf8, false))], + return_field: Arc::new(Field::new("result", DataType::UInt64, false)), + number_rows: 1, + config_options: Arc::new(datafusion_common::config::ConfigOptions::default()), + }; + + let result = function + .as_async() + .unwrap() + .invoke_async_with_args(args) + .await + .unwrap(); + let ColumnarValue::Array(array) = result else { + panic!("expected array output"); + }; + let array = array.as_any().downcast_ref::().unwrap(); + assert_eq!(42, array.value(0)); + } + + #[tokio::test] + async fn test_discard_unflushed_data_rejects_multiple_rows() { + let factory: ScalarFunctionFactory = DiscardUnflushedDataFunction::factory().into(); + let function = factory.provide(FunctionContext::mock()); + let args = datafusion::logical_expr::ScalarFunctionArgs { + args: vec![ColumnarValue::Array(Arc::new(UInt64Array::from(vec![ + 1, 2, + ])))], + arg_fields: vec![Arc::new(Field::new("arg_0", DataType::UInt64, false))], + return_field: Arc::new(Field::new("result", DataType::UInt64, false)), + number_rows: 2, + config_options: Arc::new(datafusion_common::config::ConfigOptions::default()), + }; + + let error = function + .as_async() + .unwrap() + .invoke_async_with_args(args) + .await + .unwrap_err(); + assert!(error.to_string().contains("received 2")); + } +} diff --git a/src/common/function/src/handlers.rs b/src/common/function/src/handlers.rs index 47de2824b8f..24ce115a0d6 100644 --- a/src/common/function/src/handlers.rs +++ b/src/common/function/src/handlers.rs @@ -69,6 +69,20 @@ pub trait TableMutationHandler: Send + Sync { region_id: RegionId, ctx: QueryContextRef, ) -> Result; + + /// Discard all unflushed data from a table region. + async fn discard_unflushed_data( + &self, + region_id: RegionId, + ctx: QueryContextRef, + ) -> Result; + + /// Discard all unflushed data from all regions of a table. + async fn discard_unflushed_data_by_table( + &self, + table_name: TableName, + ctx: QueryContextRef, + ) -> Result; } /// A trait for handling procedure service requests in `QueryEngine`. diff --git a/src/common/function/src/state.rs b/src/common/function/src/state.rs index 726785f0741..b1802b62578 100644 --- a/src/common/function/src/state.rs +++ b/src/common/function/src/state.rs @@ -48,6 +48,7 @@ impl FunctionState { BuildIndexTableRequest, CompactTableRequest, DeleteRequest, FlushTableRequest, InsertRequest, }; + use table::table_name::TableName; use crate::handlers::{FlowServiceHandler, ProcedureServiceHandler, TableMutationHandler}; struct MockProcedureServiceHandler; @@ -171,6 +172,22 @@ impl FunctionState { ) -> Result { Ok(ROWS) } + + async fn discard_unflushed_data( + &self, + _region_id: RegionId, + _ctx: QueryContextRef, + ) -> Result { + Ok(ROWS) + } + + async fn discard_unflushed_data_by_table( + &self, + _table_name: TableName, + _ctx: QueryContextRef, + ) -> Result { + Ok(ROWS) + } } #[async_trait] diff --git a/src/operator/src/request.rs b/src/operator/src/request.rs index aea395d8eed..b46a6b657a6 100644 --- a/src/operator/src/request.rs +++ b/src/operator/src/request.rs @@ -18,9 +18,11 @@ use api::helper::to_pb_time_unit; use api::v1::region::region_request::Body as RegionRequestBody; use api::v1::region::{ BuildIndexRequest, CompactRequest, CompactionTimeRange, FlushRequest, RegionRequestHeader, + TruncateRequest, Unflushed, truncate_request, }; use catalog::CatalogManagerRef; use common_catalog::build_db_string; +use common_catalog::consts::METRIC_ENGINE; use common_meta::node_manager::{AffectedRows, NodeManagerRef}; use common_meta::peer::Peer; use common_telemetry::tracing_context::TracingContext; @@ -34,10 +36,12 @@ use session::context::QueryContextRef; use snafu::prelude::*; use store_api::storage::RegionId; use table::requests::{BuildIndexTableRequest, CompactTableRequest, FlushTableRequest}; +use table::table_name::TableName; use crate::error::{ CatalogSnafu, FindRegionLeaderSnafu, FindTablePartitionRuleSnafu, JoinTaskSnafu, - RequestRegionSnafu, Result, TableNotFoundSnafu, UnsupportedRegionRequestSnafu, + NotSupportedSnafu, RequestRegionSnafu, Result, TableNotFoundSnafu, + UnsupportedRegionRequestSnafu, }; use crate::region_req_factory::RegionRequestFactory; @@ -206,6 +210,76 @@ impl Requester { info!("Handle region manual compaction request: {region_id}"); self.do_request(vec![request], None, &ctx).await } + + /// Discard all unflushed data from the region. + pub async fn handle_discard_unflushed_data( + &self, + region_id: RegionId, + ctx: QueryContextRef, + ) -> Result { + let request = RegionRequestBody::Truncate(TruncateRequest { + region_id: region_id.into(), + kind: Some(truncate_request::Kind::Unflushed(Unflushed {})), + }); + + info!("Handle region discard unflushed data request: {region_id}"); + self.do_request(vec![request], None, &ctx).await + } + + /// Discard all unflushed data from all regions of the table. + pub async fn handle_discard_unflushed_data_by_table( + &self, + table_name: TableName, + ctx: QueryContextRef, + ) -> Result { + let table = self + .catalog_manager + .table( + &table_name.catalog_name, + &table_name.schema_name, + &table_name.table_name, + None, + ) + .await + .context(CatalogSnafu)?; + let table = table.with_context(|| TableNotFoundSnafu { + table_name: table_name.to_string(), + })?; + let table_info = table.table_info(); + ensure_discard_unflushed_supported( + &table_info.meta.engine, + table_info.is_physical_table(), + )?; + + let partitions = &self + .partition_manager + .find_physical_partition_info(table_info.ident.table_id) + .await + .with_context(|_| FindTablePartitionRuleSnafu { + table_name: table_name.to_string(), + })? + .partitions; + let requests = partitions + .iter() + .map(|partition| { + RegionRequestBody::Truncate(TruncateRequest { + region_id: partition.id.into(), + kind: Some(truncate_request::Kind::Unflushed(Unflushed {})), + }) + }) + .collect(); + + info!("Handle table discard unflushed data request: {table_name}"); + self.do_request( + requests, + Some(build_db_string( + &table_name.catalog_name, + &table_name.schema_name, + )), + &ctx, + ) + .await + } } fn to_pb_compaction_time_range(range: TimestampRange) -> Result { @@ -291,6 +365,7 @@ impl Requester { RegionRequestBody::Flush(req) => req.region_id, RegionRequestBody::Compact(req) => req.region_id, RegionRequestBody::BuildIndex(req) => req.region_id, + RegionRequestBody::Truncate(req) => req.region_id, _ => { error!("Unsupported region request: {:?}", req); return UnsupportedRegionRequestSnafu {}.fail(); @@ -329,6 +404,16 @@ impl Requester { } } +fn ensure_discard_unflushed_supported(engine: &str, is_physical_table: bool) -> Result<()> { + ensure!( + engine != METRIC_ENGINE || is_physical_table, + NotSupportedSnafu { + feat: "discarding unflushed data from a Metric Engine logical table" + } + ); + Ok(()) +} + #[cfg(test)] mod tests { use api::v1::TimeUnit; @@ -351,6 +436,15 @@ mod tests { assert_eq!(TimeUnit::Second as i32, pb_range.time_unit); } + #[test] + fn test_discard_unflushed_rejects_metric_logical_table() { + let error = ensure_discard_unflushed_supported(METRIC_ENGINE, false).unwrap_err(); + assert!(matches!(error, crate::error::Error::NotSupported { .. })); + + ensure_discard_unflushed_supported(METRIC_ENGINE, true).unwrap(); + ensure_discard_unflushed_supported("mito", false).unwrap(); + } + #[test] fn test_to_pb_compaction_time_range() { let range = TimestampRange::new( diff --git a/src/operator/src/table.rs b/src/operator/src/table.rs index 13ed57200c2..117351280b9 100644 --- a/src/operator/src/table.rs +++ b/src/operator/src/table.rs @@ -26,6 +26,7 @@ use table::requests::{ BuildIndexTableRequest, CompactTableRequest, DeleteRequest as TableDeleteRequest, FlushTableRequest, InsertRequest as TableInsertRequest, }; +use table::table_name::TableName; use crate::delete::DeleterRef; use crate::insert::InserterRef; @@ -132,4 +133,28 @@ impl TableMutationHandler for TableMutationOperator { .map_err(BoxedError::new) .context(query_error::TableMutationSnafu) } + + async fn discard_unflushed_data( + &self, + region_id: RegionId, + ctx: QueryContextRef, + ) -> QueryResult { + self.requester + .handle_discard_unflushed_data(region_id, ctx) + .await + .map_err(BoxedError::new) + .context(query_error::TableMutationSnafu) + } + + async fn discard_unflushed_data_by_table( + &self, + table_name: TableName, + ctx: QueryContextRef, + ) -> QueryResult { + self.requester + .handle_discard_unflushed_data_by_table(table_name, ctx) + .await + .map_err(BoxedError::new) + .context(query_error::TableMutationSnafu) + } } diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index 490e6c4d6e6..3bb2bb916c1 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -458,6 +458,10 @@ fn make_region_truncate(truncate: TruncateRequest) -> Result Ok(vec![( + region_id, + RegionRequest::Truncate(RegionTruncateRequest::Unflushed), + )]), } } @@ -1811,6 +1815,41 @@ mod tests { ); } + #[test] + fn test_make_region_truncate_unflushed() { + let region_id = RegionId::new(42, 3); + let requests = + RegionRequest::try_from_request_body(region_request::Body::Truncate(TruncateRequest { + region_id: region_id.as_u64(), + kind: Some(truncate_request::Kind::Unflushed( + api::v1::region::Unflushed {}, + )), + })) + .unwrap(); + + assert_eq!(region_id, requests[0].0); + assert!(matches!( + requests[0].1, + RegionRequest::Truncate(RegionTruncateRequest::Unflushed) + )); + } + + #[test] + fn test_make_region_truncate_requires_kind() { + let error = + RegionRequest::try_from_request_body(region_request::Body::Truncate(TruncateRequest { + region_id: RegionId::new(42, 3).as_u64(), + kind: None, + })) + .unwrap_err(); + + assert!( + error + .to_string() + .contains("missing kind in TruncateRequest") + ); + } + #[test] fn test_from_proto_location() { let proto_location = v1::AddColumnLocation { diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index 9c64e841254..f21b459c2a7 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -90,6 +90,81 @@ async fn test_create_database_and_insert_query(instance: Arc) } } +#[apply(both_instances_cases)] +async fn test_admin_discard_unflushed_data(instance: Arc) { + let instance = instance.frontend(); + + execute_sql( + &instance, + r#"CREATE TABLE discard_unflushed_data_test ( + host STRING PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX + ) ENGINE = mito"#, + ) + .await; + execute_sql( + &instance, + "INSERT INTO discard_unflushed_data_test VALUES ('persisted', 1, 1)", + ) + .await; + execute_sql( + &instance, + "ADMIN FLUSH_TABLE('discard_unflushed_data_test')", + ) + .await; + execute_sql( + &instance, + "INSERT INTO discard_unflushed_data_test VALUES ('unflushed', 2, 2)", + ) + .await; + + let region_output = execute_sql( + &instance, + "SELECT greptime_partition_id FROM information_schema.partitions \ + WHERE table_name = 'discard_unflushed_data_test' LIMIT 1", + ) + .await; + let OutputData::Stream(stream) = region_output.data else { + panic!("expected region id stream"); + }; + let batches = util::collect(stream).await.unwrap(); + let region_ids = batches[0].column(0); + let region_ids = region_ids.as_any().downcast_ref::().unwrap(); + let region_id = region_ids.value(0); + + let admin_sql = format!("ADMIN discard_unflushed({region_id})"); + for _ in 0..2 { + let output = execute_sql(&instance, &admin_sql).await; + let OutputData::RecordBatches(batches) = output.data else { + panic!("expected ADMIN result batches"); + }; + let result = batches.iter().next().unwrap().column(0); + let result = result.as_any().downcast_ref::().unwrap(); + assert_eq!(0, result.value(0)); + } + + let output = execute_sql( + &instance, + "SELECT host, val FROM discard_unflushed_data_test ORDER BY host", + ) + .await; + assert_eq!( + "+-----------+-----+\n\ + | host | val |\n\ + +-----------+-----+\n\ + | persisted | 1.0 |\n\ + +-----------+-----+", + output.data.pretty_print().await + ); + + assert!( + try_execute_sql(&instance, &format!("SELECT discard_unflushed({region_id})"),) + .await + .is_err() + ); +} + #[tokio::test(flavor = "multi_thread")] async fn test_distributed_scalar_latest_with_query_parallelism_below_regions() { common_telemetry::init_default_ut_logging(); diff --git a/tests/cases/standalone/common/function/admin/discard_unflushed_data.result b/tests/cases/standalone/common/function/admin/discard_unflushed_data.result new file mode 100644 index 00000000000..9723baa8638 --- /dev/null +++ b/tests/cases/standalone/common/function/admin/discard_unflushed_data.result @@ -0,0 +1,133 @@ +CREATE TABLE discard_unflushed_data_test ( + host STRING PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX +) ENGINE = mito; + +Affected Rows: 0 + +INSERT INTO discard_unflushed_data_test VALUES ('persisted', 1, 1); + +Affected Rows: 1 + +ADMIN FLUSH_TABLE('discard_unflushed_data_test'); + ++--------------------------------------------------+ +| ADMIN FLUSH_TABLE('discard_unflushed_data_test') | ++--------------------------------------------------+ +| 0 | ++--------------------------------------------------+ + +INSERT INTO discard_unflushed_data_test VALUES ('unflushed', 2, 2); + +Affected Rows: 1 + +ADMIN discard_unflushed('discard_unflushed_data_test'); + ++--------------------------------------------------------+ +| ADMIN discard_unflushed('discard_unflushed_data_test') | ++--------------------------------------------------------+ +| 0 | ++--------------------------------------------------------+ + +-- Repeating the operation is idempotent. +ADMIN discard_unflushed('discard_unflushed_data_test'); + ++--------------------------------------------------------+ +| ADMIN discard_unflushed('discard_unflushed_data_test') | ++--------------------------------------------------------+ +| 0 | ++--------------------------------------------------------+ + +SELECT host, val FROM discard_unflushed_data_test ORDER BY host; + ++-----------+-----+ +| host | val | ++-----------+-----+ +| persisted | 1.0 | ++-----------+-----+ + +-- The function must only be available through ADMIN. +-- SQLNESS REPLACE \nDid\syou\smean.* +SELECT discard_unflushed(0); + +Error: 3000(PlanQuery), Failed to plan SQL: Error during planning: Invalid function 'discard_unflushed'. + +DROP TABLE discard_unflushed_data_test; + +Affected Rows: 0 + +CREATE TABLE discard_unflushed_data_test ( + host STRING PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX +) ENGINE = mito; + +Affected Rows: 0 + +INSERT INTO discard_unflushed_data_test VALUES ('unflushed', 1, 1); + +Affected Rows: 1 + +-- The table has no persisted data, so discarding leaves it empty. +ADMIN discard_unflushed('discard_unflushed_data_test'); + ++--------------------------------------------------------+ +| ADMIN discard_unflushed('discard_unflushed_data_test') | ++--------------------------------------------------------+ +| 0 | ++--------------------------------------------------------+ + +SELECT COUNT(*) FROM discard_unflushed_data_test; + ++----------+ +| count(*) | ++----------+ +| 0 | ++----------+ + +DROP TABLE discard_unflushed_data_test; + +Affected Rows: 0 + +CREATE TABLE discard_unflushed_data_physical ( + ts TIMESTAMP TIME INDEX, + val DOUBLE +) ENGINE = metric WITH ("physical_metric_table" = ""); + +Affected Rows: 0 + +CREATE TABLE discard_unflushed_data_logical ( + host STRING PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX +) ENGINE = metric WITH ("on_physical_table" = "discard_unflushed_data_physical"); + +Affected Rows: 0 + +INSERT INTO discard_unflushed_data_logical VALUES ('unflushed', 1, 1); + +Affected Rows: 1 + +-- Logical Metric Engine tables share their physical regions with other logical tables. +-- Discarding by logical table name must not truncate those shared regions. +ADMIN discard_unflushed('discard_unflushed_data_logical'); + +Error: 1002(Unexpected), Failed to execute admin function discard_unflushed: Execution error: Not supported: discarding unflushed data from a Metric Engine logical table + +SELECT host, val FROM discard_unflushed_data_logical; + ++-----------+-----+ +| host | val | ++-----------+-----+ +| unflushed | 1.0 | ++-----------+-----+ + +DROP TABLE discard_unflushed_data_logical; + +Affected Rows: 0 + +DROP TABLE discard_unflushed_data_physical; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/function/admin/discard_unflushed_data.sql b/tests/cases/standalone/common/function/admin/discard_unflushed_data.sql new file mode 100644 index 00000000000..9183489495f --- /dev/null +++ b/tests/cases/standalone/common/function/admin/discard_unflushed_data.sql @@ -0,0 +1,62 @@ +CREATE TABLE discard_unflushed_data_test ( + host STRING PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX +) ENGINE = mito; + +INSERT INTO discard_unflushed_data_test VALUES ('persisted', 1, 1); + +ADMIN FLUSH_TABLE('discard_unflushed_data_test'); + +INSERT INTO discard_unflushed_data_test VALUES ('unflushed', 2, 2); + +ADMIN discard_unflushed('discard_unflushed_data_test'); + +-- Repeating the operation is idempotent. +ADMIN discard_unflushed('discard_unflushed_data_test'); + +SELECT host, val FROM discard_unflushed_data_test ORDER BY host; + +-- The function must only be available through ADMIN. +-- SQLNESS REPLACE \nDid\syou\smean.* +SELECT discard_unflushed(0); + +DROP TABLE discard_unflushed_data_test; + +CREATE TABLE discard_unflushed_data_test ( + host STRING PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX +) ENGINE = mito; + +INSERT INTO discard_unflushed_data_test VALUES ('unflushed', 1, 1); + +-- The table has no persisted data, so discarding leaves it empty. +ADMIN discard_unflushed('discard_unflushed_data_test'); + +SELECT COUNT(*) FROM discard_unflushed_data_test; + +DROP TABLE discard_unflushed_data_test; + +CREATE TABLE discard_unflushed_data_physical ( + ts TIMESTAMP TIME INDEX, + val DOUBLE +) ENGINE = metric WITH ("physical_metric_table" = ""); + +CREATE TABLE discard_unflushed_data_logical ( + host STRING PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX +) ENGINE = metric WITH ("on_physical_table" = "discard_unflushed_data_physical"); + +INSERT INTO discard_unflushed_data_logical VALUES ('unflushed', 1, 1); + +-- Logical Metric Engine tables share their physical regions with other logical tables. +-- Discarding by logical table name must not truncate those shared regions. +ADMIN discard_unflushed('discard_unflushed_data_logical'); + +SELECT host, val FROM discard_unflushed_data_logical; + +DROP TABLE discard_unflushed_data_logical; + +DROP TABLE discard_unflushed_data_physical; From f88226c3ff64f3cf410950b483f43f170fb07d0e Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 20:27:47 +0800 Subject: [PATCH 13/92] ci: update dev-builder image tag (#8831) Signed-off-by: greptimedb-ci Co-authored-by: greptimedb-ci --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 5bdb63d3520..0f280b7b9be 100644 --- a/Makefile +++ b/Makefile @@ -8,7 +8,7 @@ CARGO_BUILD_OPTS := --locked IMAGE_REGISTRY ?= docker.io IMAGE_NAMESPACE ?= greptime IMAGE_TAG ?= latest -DEV_BUILDER_IMAGE_TAG ?= 2026-03-21-9c9d9e9e-20260331090344 +DEV_BUILDER_IMAGE_TAG ?= 2026-03-21-7fa6f5f9-20260810091725 DEV_BUILDER_RISCV64_IMAGE_TAG ?= $(DEV_BUILDER_IMAGE_TAG) DEV_BUILDER_RISCV64_BASE_IMAGE ?= ubuntu:22.04 RISCV64_TARGET ?= riscv64gc-unknown-linux-gnu From e2ab11763a945504988d4e246181b68a1db07197 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 10 Aug 2026 21:14:56 +0800 Subject: [PATCH 14/92] fix(ci): identify team members by repository permission (#8822) * fix(ci): identify team members by repository permission `author_association` is computed from what the caller can see, so GITHUB_TOKEN reports a private organization member as CONTRIBUTOR. Only 5 of GreptimeTeam's members have public membership, so the open-pull-request check skipped almost everyone it was written for. Use the repository permission of the author instead, which is viewer-independent. On error, apply the limit rather than skipping, so a token that cannot read permissions cannot silently disable the check again. Signed-off-by: Dennis Zhuang * fix(ci): do not log repository permission levels Job logs are public. Resolving the author's permission is fine; printing the level is not. Signed-off-by: Dennis Zhuang * Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --------- Signed-off-by: Dennis Zhuang Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .github/scripts/pr-open-limit.js | 37 ++++++++++++++++++++++------- .github/workflows/pr-open-limit.yml | 1 - 2 files changed, 29 insertions(+), 9 deletions(-) diff --git a/.github/scripts/pr-open-limit.js b/.github/scripts/pr-open-limit.js index 1ba908b7450..dd1650b0d64 100644 --- a/.github/scripts/pr-open-limit.js +++ b/.github/scripts/pr-open-limit.js @@ -9,9 +9,8 @@ import { appendFileSync } from "node:fs"; const MARKER = ""; // GITHUB_TOKEN can only edit comments authored by the Actions bot itself. const COMMENT_AUTHOR = "github-actions[bot]"; -// Associations that identify an organization member. Outside collaborators and -// community contributors are out of scope. -const MEMBER_ASSOCIATIONS = ["OWNER", "MEMBER"]; +// Repository permissions that mark someone as part of the team. +const TEAM_PERMISSIONS = ["admin", "maintain", "write"]; // Keep the comment readable for authors who are far over the limit. const MAX_LISTED_PRS = 10; @@ -22,6 +21,29 @@ function summary(text) { console.log(text); } +// `author_association` cannot be used here: it is computed from what the caller +// is allowed to see, so GITHUB_TOKEN reports a private organization member as +// CONTRIBUTOR. Repository permission is viewer-independent. +// +// On error this returns true rather than false. A token that cannot read +// permissions would otherwise make the whole check silently pass everyone. +async function isTeamMember(octokit, owner, repo, username) { + try { + const { data } = await octokit.repos.getCollaboratorPermissionLevel({ + owner, + repo, + username, + }); + // Never log the level itself: job logs are public. + return TEAM_PERMISSIONS.includes(data.permission); + } catch (error) { + console.log( + `Cannot read repository permission for \`${username}\` (HTTP ${error.status}); applying the limit anyway.` + ); + return true; + } +} + function buildComment(author, total, limit, openPrs) { const listed = openPrs .slice(0, MAX_LISTED_PRS) @@ -69,7 +91,6 @@ async function upsertComment(octokit, params, body) { const [owner, repo] = process.env.GITHUB_REPOSITORY.split("/"); const prNumber = Number(process.env.PR_NUMBER); const author = process.env.PR_AUTHOR; - const association = process.env.PR_AUTHOR_ASSOCIATION; const rawLimit = Number(process.env.MAX_OPEN_PRS); const limit = Number.isInteger(rawLimit) && rawLimit >= 0 ? rawLimit : 5; @@ -78,13 +99,13 @@ async function upsertComment(octokit, params, body) { return; } - if (!MEMBER_ASSOCIATIONS.includes(association)) { - summary(`Skipping \`${author}\`: association is \`${association}\`, not an org member.`); + const octokit = new Octokit({ auth: process.env.GITHUB_TOKEN }); + + if (!(await isTeamMember(octokit, owner, repo, author))) { + summary(`Skipping \`${author}\`: not a team member.`); return; } - const octokit = new Octokit({ auth: process.env.GITHUB_TOKEN }); - const allOpen = await octokit.paginate(octokit.pulls.list, { owner, repo, diff --git a/.github/workflows/pr-open-limit.yml b/.github/workflows/pr-open-limit.yml index d22929911a8..de5b723cb25 100644 --- a/.github/workflows/pr-open-limit.yml +++ b/.github/workflows/pr-open-limit.yml @@ -44,5 +44,4 @@ jobs: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} PR_NUMBER: ${{ github.event.pull_request.number }} PR_AUTHOR: ${{ github.event.pull_request.user.login }} - PR_AUTHOR_ASSOCIATION: ${{ github.event.pull_request.author_association }} MAX_OPEN_PRS: ${{ vars.MAX_OPEN_PRS_PER_AUTHOR || 5 }} From e4175475bb7e630cefefa189b17282b063868a49 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 10 Aug 2026 21:30:34 +0800 Subject: [PATCH 15/92] fix(mito): re-encode bulk WAL entry after filling missing columns (#8808) When a bulk insert request carries a stale schema, the worker fills the missing columns via BulkPart::fill_missing_columns before writing. The method replaced the batch but kept raw_data (the original Arrow Flight IPC bytes), while BulkWalEntry::try_from(&BulkPart) prefers raw_data, so the memtable received the filled batch but the WAL recorded the pre-fill bytes. Replaying such an entry restores a batch that misses the filled columns: - Bulk memtable (flat format): convert_bulk_part fails with ColumnNotFound; the error is swallowed by the no-op write notifier and the rows are silently lost after restart. - Time series memtable: BulkPart::to_mutation builds rows shorter than the declared schema and the region worker panics with index out of bounds during replay, hanging the region open. Fixes: - fill_missing_columns clears raw_data so the WAL entry is re-encoded from the filled batch. - replay_memtable fills missing columns for replayed bulk parts of dense regions, so entries already written by affected versions replay correctly. Signed-off-by: jeremyhi --- src/mito2/src/engine/alter_test.rs | 214 +++++++++++++++++++++++++++- src/mito2/src/memtable/bulk/part.rs | 64 +++++++++ src/mito2/src/region/opener.rs | 13 +- 3 files changed, 284 insertions(+), 7 deletions(-) diff --git a/src/mito2/src/engine/alter_test.rs b/src/mito2/src/engine/alter_test.rs index be4daa1d88e..2461f2ca0cd 100644 --- a/src/mito2/src/engine/alter_test.rs +++ b/src/mito2/src/engine/alter_test.rs @@ -17,21 +17,26 @@ use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; +use api::v1::bulk_wal_entry::Body; use api::v1::helper::{row, tag_column_schema}; use api::v1::value::ValueData; -use api::v1::{ColumnDataType, Row, Rows, SemanticType, Value}; +use api::v1::{ArrowIpc, BulkWalEntry, ColumnDataType, Row, Rows, SemanticType, Value, WalEntry}; use common_error::ext::ErrorExt; use common_meta::ddl::utils::{parse_column_metadatas, parse_manifest_infos_from_extensions}; -use common_recordbatch::RecordBatches; +use common_recordbatch::{DfRecordBatch, RecordBatches}; +use common_test_util::flight::encode_to_flight_data; use datafusion_expr::col; +use datatypes::arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray}; +use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, FulltextAnalyzer, FulltextBackend, FulltextOptions}; +use store_api::logstore::provider::Provider; use store_api::metadata::ColumnMetadata; use store_api::metric_engine_consts::TABLE_COLUMN_METADATA_EXTENSION_KEY; use store_api::region_engine::{RegionEngine, RegionManifestInfo, RegionRole}; use store_api::region_request::{ - AddColumn, AddColumnLocation, AlterKind, PathType, RegionAlterRequest, RegionOpenRequest, - RegionRequest, SetIndexOption, SetRegionOption, + AddColumn, AddColumnLocation, AlterKind, PathType, RegionAlterRequest, + RegionBulkInsertsRequest, RegionOpenRequest, RegionRequest, SetIndexOption, SetRegionOption, }; use store_api::storage::{ColumnId, RegionId, ScanRequest}; @@ -42,9 +47,10 @@ use crate::error; use crate::sst::FormatType; use crate::test_util::batch_util::sort_batches_and_print; use crate::test_util::{ - CreateRequestBuilder, TestEnv, build_rows, build_rows_for_key, - column_metadata_to_column_schema, flush_region, put_rows, rows_schema, + CreateRequestBuilder, LogStoreImpl, TestEnv, build_rows, build_rows_for_key, + column_metadata_to_column_schema, flush_region, put_rows, reopen_region, rows_schema, }; +use crate::wal::Wal; async fn scan_check_after_alter(engine: &MitoEngine, region_id: RegionId, expected: &str) { let request = ScanRequest::default(); @@ -1978,3 +1984,199 @@ async fn test_alter_region_append_mode_invalid() { // append_mode should still be true check_append_mode(&engine, true); } + +/// Builds a batch against the schema before [add_nullable_field1] +/// (schema version 0): `(tag_0, field_0, ts)`. +fn build_schema_v0_batch() -> DfRecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("tag_0", DataType::Utf8, true), + Field::new("field_0", DataType::Float64, true), + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])); + let tag = Arc::new(StringArray::from_iter_values( + (0..3).map(|v| format!("a{v}")), + )) as ArrayRef; + let field = Arc::new(Float64Array::from_iter_values((0..3).map(|v| v as f64))) as ArrayRef; + let ts = Arc::new(TimestampMillisecondArray::from_iter_values( + (0..3).map(|v| v as i64 * 1000), + )) as ArrayRef; + DfRecordBatch::try_new(schema, vec![tag, field, ts]).unwrap() +} + +fn encode_arrow_ipc(batch: &DfRecordBatch) -> ArrowIpc { + let (schema, record_batch) = encode_to_flight_data(batch.clone()); + ArrowIpc { + schema: schema.data_header, + data_header: record_batch.data_header, + payload: record_batch.data_body, + } +} + +/// Builds a bulk insert request whose batch was built against the schema +/// before [add_nullable_field1] (schema version 0): `(tag_0, field_0, ts)`. +fn build_schema_v0_bulk_request(region_id: RegionId) -> RegionBulkInsertsRequest { + let payload = build_schema_v0_batch(); + let raw_data = encode_arrow_ipc(&payload); + + RegionBulkInsertsRequest { + region_id, + payload, + raw_data, + partition_expr_version: None, + // The writer only knows schema version 0 while the region is already + // at version 1, so the worker has to fill the missing columns. + aligned_schema_version: Some(0), + } +} + +async fn scan_all_sorted(engine: &MitoEngine, region_id: RegionId) -> String { + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + sort_batches_and_print(&batches, &["tag_0", "ts"]) +} + +const SCHEMA_V0_BULK_ROWS: &str = "\ ++-------+---------+---------------------+---------+ +| tag_0 | field_0 | ts | field_1 | ++-------+---------+---------------------+---------+ +| a0 | 0.0 | 1970-01-01T00:00:00 | | +| a1 | 1.0 | 1970-01-01T00:00:01 | | +| a2 | 2.0 | 1970-01-01T00:00:02 | | ++-------+---------+---------------------+---------+"; + +#[tokio::test] +async fn test_bulk_insert_stale_schema_wal_replay() { + test_bulk_insert_stale_schema_wal_replay_with_format(false).await; + test_bulk_insert_stale_schema_wal_replay_with_format(true).await; +} + +async fn test_bulk_insert_stale_schema_wal_replay_with_format(flat_format: bool) { + common_telemetry::init_default_ut_logging(); + + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let table_dir = request.table_dir.clone(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + // Bumps the region schema version to 1. + engine + .handle_request(region_id, RegionRequest::Alter(add_nullable_field1())) + .await + .unwrap(); + + // Bulk insert whose batch misses `field_1`. The worker fills the missing + // column before writing to the memtable and the WAL. + let response = engine + .handle_request( + region_id, + RegionRequest::BulkInserts(build_schema_v0_bulk_request(region_id)), + ) + .await + .unwrap(); + assert_eq!(3, response.affected_rows); + + assert_eq!( + SCHEMA_V0_BULK_ROWS, + scan_all_sorted(&engine, region_id).await, + "flat_format: {flat_format}" + ); + + // Reopens the region to replay the WAL. The rows must survive the replay. + reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + + assert_eq!( + SCHEMA_V0_BULK_ROWS, + scan_all_sorted(&engine, region_id).await, + "flat_format: {flat_format}" + ); +} + +/// Replays a bulk WAL entry that misses columns added by an alter. This +/// simulates entries written by old versions that keep the stale raw data +/// when filling missing columns. +#[tokio::test] +async fn test_bulk_insert_stale_wal_entry_replay() { + test_bulk_insert_stale_wal_entry_replay_with_format(false).await; + test_bulk_insert_stale_wal_entry_replay_with_format(true).await; +} + +async fn test_bulk_insert_stale_wal_entry_replay_with_format(flat_format: bool) { + common_telemetry::init_default_ut_logging(); + + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let table_dir = request.table_dir.clone(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + // Bumps the region schema version to 1. + engine + .handle_request(region_id, RegionRequest::Alter(add_nullable_field1())) + .await + .unwrap(); + + // Writes a bulk entry whose batch misses `field_1` to the WAL directly. + let batch = build_schema_v0_batch(); + let entry = WalEntry { + mutations: vec![], + bulk_entries: vec![BulkWalEntry { + sequence: 1, + max_ts: 2000, + min_ts: 0, + timestamp_index: 2, + body: Some(Body::ArrowIpc(encode_arrow_ipc(&batch))), + }], + }; + let LogStoreImpl::RaftEngine(log_store) = env.get_log_store().unwrap() else { + unreachable!() + }; + let wal = Wal::new(log_store); + let mut writer = wal.writer(); + writer + .add_entry( + region_id, + 1, + &entry, + &Provider::raft_engine_provider(region_id.as_u64()), + ) + .unwrap(); + writer.write_to_wal().await.unwrap(); + + // Reopens the region to replay the WAL. The replay must fill the missing + // column instead of losing the rows. + reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + + assert_eq!( + SCHEMA_V0_BULK_ROWS, + scan_all_sorted(&engine, region_id).await, + "flat_format: {flat_format}" + ); +} diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index fcba4519478..20482794079 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -262,6 +262,12 @@ impl BulkPart { // Update the batch self.batch = new_batch; + // The raw data no longer matches the batch after filling missing + // columns. Clears it so the WAL entry is re-encoded from the filled + // batch, otherwise replaying the entry restores a batch without the + // filled columns and the memtable rejects it. + self.raw_data = None; + Ok(()) } @@ -2896,4 +2902,62 @@ mod tests { let total_rows: usize = reader.map(|r| r.unwrap().num_rows()).sum(); assert_eq!(total_rows, 6); } + + #[test] + fn test_fill_missing_columns_resets_raw_data() { + let metadata = metadata_for_test(); + + // A batch missing the `v1` field column, e.g. built against an older + // schema before `v1` was added. + let k0_array = Arc::new(arrow::array::StringArray::from(vec!["key1", "key2"])); + let k1_array = Arc::new(arrow::array::UInt32Array::from(vec![1u32, 2])); + let v0_array = Arc::new(arrow::array::Int64Array::from(vec![100i64, 200])); + let ts_array = Arc::new(TimestampMillisecondArray::from(vec![1000i64, 2000])); + let input_schema = Arc::new(Schema::new(vec![ + Field::new("k0", ArrowDataType::Utf8, false), + Field::new("k1", ArrowDataType::UInt32, false), + Field::new("v0", ArrowDataType::Int64, true), + Field::new( + "ts", + ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + ])); + let batch = + RecordBatch::try_new(input_schema, vec![k0_array, k1_array, v0_array, ts_array]) + .unwrap(); + + // Encodes the batch as raw data like the bulk insert request does. + let mut encoder = FlightEncoder::default(); + let schema_bytes = encoder.encode_schema(batch.schema().as_ref()).data_header; + let [flight_data] = encoder + .encode(FlightMessage::RecordBatch(batch.clone())) + .try_into() + .unwrap(); + let mut part = BulkPart { + batch, + max_timestamp: 2000, + min_timestamp: 1000, + sequence: 5, + timestamp_index: 3, + raw_data: Some(ArrowIpc { + schema: schema_bytes, + data_header: flight_data.data_header, + payload: flight_data.data_body, + }), + }; + + part.fill_missing_columns(&metadata).unwrap(); + assert!(part.batch.column_by_name("v1").is_some()); + // The raw data no longer matches the batch after filling missing + // columns and must be cleared, otherwise the WAL entry built from + // this part loses the filled columns. + assert!(part.raw_data.is_none()); + + // The WAL entry round trip keeps the filled column. + let entry = BulkWalEntry::try_from(&part).unwrap(); + let replayed = BulkPart::try_from(entry).unwrap(); + assert_eq!(2, replayed.num_rows()); + assert!(replayed.batch.column_by_name("v1").is_some()); + } } diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 9388f943ceb..5d2d36df107 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -32,6 +32,7 @@ use object_store::manager::ObjectStoreManagerRef; use object_store::util::{is_object_storage, normalize_dir}; use parquet::file::metadata::PageIndexPolicy; use snafu::{OptionExt, ResultExt, ensure}; +use store_api::codec::PrimaryKeyEncoding; use store_api::logstore::LogStore; use store_api::logstore::provider::Provider; use store_api::metadata::{ @@ -832,6 +833,7 @@ where // data in the WAL. let mut last_entry_id = flushed_entry_id; let replay_from_entry_id = flushed_entry_id + 1; + let region_metadata = version_control.current().version.metadata.clone(); let mut wal_stream = wal_entry_reader.read(provider, replay_from_entry_id)?; while let Some(res) = wal_stream.next().await { @@ -876,7 +878,16 @@ where } for bulk_entry in entry.bulk_entries { - let part = BulkPart::try_from(bulk_entry)?; + let mut part = BulkPart::try_from(bulk_entry)?; + // The entry may miss columns added by a concurrent alter if it was + // written by a writer with a stale schema (older versions kept the + // stale raw data when filling missing columns). Fills missing + // columns like the write path, otherwise the memtable rejects the + // batch. Skips sparse batches as they don't carry tag columns by + // design. + if region_metadata.primary_key_encoding != PrimaryKeyEncoding::Sparse { + part.fill_missing_columns(®ion_metadata)?; + } rows_replayed += part.num_rows(); // During replay, we should adopt the sequence from WAL. let bulk_sequence_from_wal = part.sequence; From 00e45c860222195367943942b8ec21fef46b7cf5 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 10 Aug 2026 22:57:08 +0800 Subject: [PATCH 16/92] test: fix unknown channel event context expectation (#8832) Signed-off-by: WenyXu --- tests-integration/tests/table_ddl_event.rs | 31 ++++++++++++---------- 1 file changed, 17 insertions(+), 14 deletions(-) diff --git a/tests-integration/tests/table_ddl_event.rs b/tests-integration/tests/table_ddl_event.rs index 45d87f30971..c53842d66d6 100644 --- a/tests-integration/tests/table_ddl_event.rs +++ b/tests-integration/tests/table_ddl_event.rs @@ -90,7 +90,7 @@ async fn test_table_ddl_procedure_events() { "create_table", &auto_create_procedure_id, "auto_create", - "prometheus", + Some("prometheus"), ) .await; @@ -102,7 +102,7 @@ async fn test_table_ddl_procedure_events() { "create_table", &auto_influx_create_procedure_id, "auto_create", - "influx", + Some("influx"), ) .await; @@ -151,7 +151,7 @@ async fn test_table_ddl_procedure_events() { "create_table", &grpc_create_procedure_id, "manual", - "grpc", + Some("grpc"), ) .await; GrpcQueryHandler::do_query( @@ -190,7 +190,7 @@ async fn test_table_ddl_procedure_events() { "alter_table", &grpc_alter_procedure_id, "manual", - "grpc", + Some("grpc"), ) .await; @@ -215,7 +215,7 @@ async fn test_table_ddl_procedure_events() { "create_table", &create_table_procedure_id, "manual", - "httpsql", + Some("httpsql"), ) .await; assert_named_submitted_event( @@ -263,7 +263,7 @@ async fn test_table_ddl_procedure_events() { "alter_table", &alter_table_procedure_id, "manual", - "httpsql", + Some("httpsql"), ) .await; assert_named_submitted_event( @@ -307,7 +307,7 @@ async fn test_table_ddl_procedure_events() { "truncate_table", &truncate_table_procedure_id, "manual", - "httpsql", + Some("httpsql"), ) .await; assert_named_submitted_event( @@ -348,7 +348,7 @@ async fn test_table_ddl_procedure_events() { "drop_table", &drop_table_procedure_id, "manual", - "httpsql", + Some("httpsql"), ) .await; assert_named_submitted_event( @@ -389,7 +389,7 @@ async fn test_table_ddl_procedure_events() { "undrop_table", &undrop_table_procedure_id, "manual", - "unknown", + None, ) .await; assert_id_submitted_event( @@ -429,7 +429,7 @@ async fn test_table_ddl_procedure_events() { "purge_dropped_table", &purge_table_procedure_id, "manual", - "unknown", + None, ) .await; assert_id_submitted_event( @@ -483,7 +483,7 @@ async fn test_table_ddl_procedure_events() { "create_logical_tables", &create_logical_tables_procedure_id, "manual", - "httpsql", + Some("httpsql"), ) .await; @@ -536,7 +536,7 @@ async fn test_table_ddl_procedure_events() { "alter_logical_tables", &alter_logical_tables_procedure_id, "manual", - "httpsql", + Some("httpsql"), ) .await; assert_named_submitted_event( @@ -629,7 +629,7 @@ async fn assert_event_context( event_type: &str, procedure_id: &str, reason: &str, - protocol: &str, + protocol: Option<&str>, ) { let event_filter = format!( "FROM {EVENTS_TABLE} WHERE {} = '{event_type}' AND {} = '{procedure_id}' AND json_get_string({}, 'type') = 'Submitted'", @@ -649,7 +649,10 @@ async fn assert_event_context( "event_context", ) .await; - let expected = format!(r#"{{"protocol":"{protocol}","reason":"{reason}"}}"#); + let expected = match protocol { + Some(protocol) => format!(r#"{{"protocol":"{protocol}","reason":"{reason}"}}"#), + None => format!(r#"{{"reason":"{reason}"}}"#), + }; assert_eq!(expected, actual); } From cd6efa0abf6ec99737ad8bce36376c4946023edc Mon Sep 17 00:00:00 2001 From: discord9 Date: Tue, 11 Aug 2026 11:43:17 +0800 Subject: [PATCH 17/92] fix(query): ignore field metadata in MergeScan remote schema validation (#8818) * fix(query): ignore field metadata in MergeScan remote schema validation validate_remote_schema compared Arrow Fields with the default Field equality, which includes field metadata. A column whose SST region metadata carries greptime:skipping_index (or greptime:inverted_index) while the frontend table schema does not declare it was misreported as a schema mismatch (HTTP 500 'advertised remote stream schema field mismatch'), even though name, data_type, and nullability were identical. Field metadata is auxiliary (index/encoding info) and is not part of field semantics. Compare name + data_type + nullability only; JSON fields keep their existing semantic comparison (wire Binary vs decoded Struct) and timestamp timezone-only differences remain accepted. * refactor(query): drop all field metadata comparison from MergeScan schema validation fields_semantically_equal now compares name + data_type + nullability only. The previous version still compared JSON identity metadata keys (TYPE_KEY, EXTENSION_TYPE_METADATA_KEY) and is_json_field; those are auxiliary and must not participate in field semantics. JSON fields keep their separate physical-type exemption (wire Binary vs decoded Struct) via json_fields_compatible, which never applies to non-JSON fields. --- src/query/src/dist_plan/merge_scan.rs | 127 ++++++++++++++++++++------ 1 file changed, 101 insertions(+), 26 deletions(-) diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index 0245f98ddc4..aafc17c469f 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -113,6 +113,33 @@ fn merge_scan_schema_error_count_for_test() -> u64 { TEST_MERGE_SCAN_SCHEMA_ERRORS.with(Cell::get) } +/// Returns true when the field is a JSON column, identified either by its +/// Arrow extension type (`greptime.json` / `greptime.json2`) or by its +/// `greptime:type=Json` marker. +fn is_json_field(field: &Field) -> bool { + is_any_json_extension_type(field) + || field + .metadata() + .get(datatypes::schema::TYPE_KEY) + .map(String::as_ref) + == Some("Json") +} + +/// Returns true when two fields describe the same column semantics: same +/// name, same data type, and same nullability. +/// +/// Arrow `Field` metadata is auxiliary information (indexing, encoding, +/// compression hints, ...) and does not participate in the comparison. This +/// applies to every field, JSON columns included: JSON identity keys +/// (`greptime:type`, `ARROW:extension:metadata`) are never compared here. +/// The wire/decoded physical-type difference of JSON columns is exempted +/// separately by [`json_fields_compatible`]. +fn fields_semantically_equal(expected_field: &Field, actual_field: &Field) -> bool { + expected_field.name() == actual_field.name() + && expected_field.data_type() == actual_field.data_type() + && expected_field.is_nullable() == actual_field.is_nullable() +} + /// Returns true when two fields are semantically equivalent JSON columns. /// /// A JSON column is carried on the wire in its binary-encoded form (e.g. @@ -124,17 +151,8 @@ fn merge_scan_schema_error_count_for_test() -> u64 { /// settings (`ARROW:extension:metadata`, e.g. type hints) are the semantic /// parts of the field and must match. fn json_fields_compatible(expected_field: &Field, actual_field: &Field) -> bool { - let is_json = |field: &Field| { - is_any_json_extension_type(field) - || field - .metadata() - .get(datatypes::schema::TYPE_KEY) - .map(String::as_ref) - == Some("Json") - }; - - is_json(expected_field) - && is_json(actual_field) + is_json_field(expected_field) + && is_json_field(actual_field) && expected_field.name() == actual_field.name() && expected_field.is_nullable() == actual_field.is_nullable() // Both must carry the same JSON marker. @@ -151,12 +169,15 @@ fn json_fields_compatible(expected_field: &Field, actual_field: &Field) -> bool /// Validates the remote schema before positional column handling. /// -/// A timestamp timezone difference is the only intentional exception. It is -/// accepted by directly comparing the same timestamp unit, distinct timezones, -/// and equal name, nullability, and field metadata. Top-level Arrow schema -/// metadata is non-semantic at this boundary. JSON columns are additionally -/// compared semantically (see [`json_fields_compatible`]) because their wire -/// and decoded representations use different physical arrow types. +/// Field metadata (indexing/encoding/compression hints) is non-semantic at +/// this boundary and is ignored: two fields are equal when their name, data +/// type, and nullability match (see [`fields_semantically_equal`]). A +/// timestamp timezone difference is the only intentional exception, accepted +/// when the fields share the same timestamp unit, distinct timezones, and +/// equal name and nullability. Top-level Arrow schema metadata is +/// non-semantic at this boundary as well. JSON columns are compared +/// semantically (see [`json_fields_compatible`]) because their wire and +/// decoded representations use different physical arrow types. fn validate_remote_schema( expected: &ArrowSchema, actual: &ArrowSchema, @@ -176,7 +197,8 @@ fn validate_remote_schema( .zip(actual.fields().iter()) .enumerate() { - if expected_field == actual_field { + // Field metadata does not participate in the semantic comparison. + if fields_semantically_equal(expected_field, actual_field) { continue; } @@ -197,7 +219,6 @@ fn validate_remote_schema( && expected_timezone != actual_timezone && expected_field.name() == actual_field.name() && expected_field.is_nullable() == actual_field.is_nullable() - && expected_field.metadata() == actual_field.metadata() ); if !timezone_only_difference { return Err(remote_schema_mismatch(format!( @@ -2693,9 +2714,12 @@ mod tests { } #[test] - fn merge_scan_remote_schema_identity_rejects_json_vs_plain_binary_field() { - // A JSON column must not be accepted as compatible with a plain Binary - // column lacking the JSON extension metadata. + fn merge_scan_remote_schema_identity_ignores_json_extension_metadata_when_physical_type_matches() + { + // Field metadata (including the JSON extension identity) is not part + // of the semantic field identity: a JSON column in its binary wire + // form and a plain Binary column share the same name, physical type, + // and nullability, so they validate as equal. let expected = ArrowSchema::new(vec![ Field::new("j", TestArrowDataType::Binary, true).with_metadata(json_field_metadata( true, @@ -2703,11 +2727,26 @@ mod tests { )), ]); let actual = ArrowSchema::new(vec![Field::new("j", TestArrowDataType::Binary, true)]); - assert!(validate_remote_schema(&expected, &actual, "test").is_err()); + assert!(validate_remote_schema(&expected, &actual, "test").is_ok()); + + // The JSON compatibility exemption is scoped to JSON columns only: a + // pair of non-JSON fields with different physical types is still + // rejected even though one side is Binary. + let plain_binary = ArrowSchema::new(vec![Field::new("v", TestArrowDataType::Binary, true)]); + let plain_struct = ArrowSchema::new(vec![Field::new( + "v", + TestArrowDataType::Struct(arrow_schema::Fields::from(vec![Arc::new(Field::new( + "a", + TestArrowDataType::Utf8View, + true, + ))])), + true, + )]); + assert!(validate_remote_schema(&plain_binary, &plain_struct, "test").is_err()); } #[test] - fn merge_scan_remote_schema_identity_rejects_field_metadata_and_nullability_mismatch() { + fn merge_scan_remote_schema_identity_ignores_field_metadata_but_rejects_nullability_mismatch() { let expected = expected_int64_schema(); let metadata_mismatch = ArrowSchema::new_with_metadata( vec![ @@ -2730,10 +2769,42 @@ mod tests { ], expected.metadata().clone(), ); - assert!(validate_remote_schema(&expected, &metadata_mismatch, "test").is_err()); + // Field metadata is auxiliary and does not participate in the + // semantic field comparison. + assert!(validate_remote_schema(&expected, &metadata_mismatch, "test").is_ok()); assert!(validate_remote_schema(&expected, &nullability_mismatch, "test").is_err()); } + #[test] + fn merge_scan_remote_schema_identity_ignores_skipping_index_field_metadata() { + // Regression: a remote stream may advertise auxiliary field metadata + // (e.g. `greptime:skipping_index`) that the expected schema does not + // carry. Field metadata is not part of the semantic field identity + // (name + data type + nullability) and must not fail validation. + let skipping_index_metadata = StdHashMap::from([( + "greptime:skipping_index".to_string(), + r#"{"granularity":1,"false-positive-rate-in-10000":100,"index-type":"BloomFilter"}"# + .to_string(), + )]); + let expected = + ArrowSchema::new(vec![Field::new("value", TestArrowDataType::Float64, true)]); + let actual_with_metadata = ArrowSchema::new(vec![ + Field::new("value", TestArrowDataType::Float64, true) + .with_metadata(skipping_index_metadata), + ]); + assert!(validate_remote_schema(&expected, &actual_with_metadata, "test").is_ok()); + + // A real data type mismatch is still rejected. + let wrong_type = + ArrowSchema::new(vec![Field::new("value", TestArrowDataType::Int64, true)]); + assert!(validate_remote_schema(&expected, &wrong_type, "test").is_err()); + + // A name mismatch is still rejected. + let wrong_name = + ArrowSchema::new(vec![Field::new("other", TestArrowDataType::Float64, true)]); + assert!(validate_remote_schema(&expected, &wrong_name, "test").is_err()); + } + #[test] fn merge_scan_remote_schema_identity_rejects_timestamp_timezone_plus_field_mismatches() { let expected = ArrowSchema::new(vec![Field::new( @@ -2767,9 +2838,13 @@ mod tests { "different".to_string(), )])), ]); - for actual in [&different_unit, &different_name, &nullability, &metadata] { + for actual in [&different_unit, &different_name, &nullability] { assert!(validate_remote_schema(&expected, actual, "test").is_err()); } + // Field metadata is non-semantic and is ignored even for timezone + // pairs: a metadata difference on an otherwise timezone-only + // difference is accepted. + assert!(validate_remote_schema(&expected, &metadata, "test").is_ok()); } #[test] From b3faf22290e75650f1bc356ccf6b38355f1cbb29 Mon Sep 17 00:00:00 2001 From: discord9 Date: Tue, 11 Aug 2026 13:52:31 +0800 Subject: [PATCH 18/92] fix(tests): make two Windows CI failures deterministic (Nightly CI #8837) (#8840) * fix(tests): reject overlay directories before opening on all platforms DatanodeOverlay::load() opened the target before checking is_file(). On Unix, File::open on a directory succeeds and the loader rejects it with "must be a regular file". On Windows, File::open on a directory fails up front with "Access is denied", so the type check was never reached and the rejects_directories_and_parse_errors test failed 4/4 in Nightly CI (issue #8837). Check std::fs::metadata before File::open: metadata succeeds on directories on both platforms, so the error message is now identical everywhere and the test assertion holds on Windows too. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(datanode): make test_region_error deterministic across platforms The second phase raced a 100ms mock handle delay against a 200ms replay_timeout; on busy Windows CI runners the error could land after the timeout fired, flaking reply.error.is_some() (Nightly CI, issue #8837). Use a mock handle that returns the error on its first poll with no delay: the catchup future completes before replay_timeout can ever fire, so the test no longer depends on wall-clock scheduling. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/heartbeat/handler/upgrade_region.rs | 31 ++++++++++++++++--- tests/runner/src/cmd/datanode_overlay.rs | 24 +++++++++----- 2 files changed, 43 insertions(+), 12 deletions(-) diff --git a/src/datanode/src/heartbeat/handler/upgrade_region.rs b/src/datanode/src/heartbeat/handler/upgrade_region.rs index 05292415be8..e2a5677fcad 100644 --- a/src/datanode/src/heartbeat/handler/upgrade_region.rs +++ b/src/datanode/src/heartbeat/handler/upgrade_region.rs @@ -408,9 +408,12 @@ mod tests { #[tokio::test] async fn test_region_error() { common_telemetry::init_default_ut_logging(); - let mock_region_server = mock_region_server(); let region_id = RegionId::new(1024, 1); + // With the default (zero) `replay_timeout`, the handler times out before the + // mock handle's 100ms delay elapses: it didn't wait for handle returns; it + // had no idea about the error. + let region_server = mock_region_server(); let (mock_engine, _) = MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| { // Region is not ready. @@ -424,9 +427,9 @@ mod tests { // Note: Don't change. region_engine.handle_request_delay = Some(Duration::from_millis(100)); }); - mock_region_server.register_test_region(region_id, mock_engine); + region_server.register_test_region(region_id, mock_engine); let kv_backend = Arc::new(MemoryKvBackend::new()); - let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend); + let handler_context = HandlerContext::new_for_test(region_server, kv_backend); let reply = UpgradeRegionsHandler::new_test() .handle( &handler_context, @@ -437,12 +440,32 @@ mod tests { ) .await; - // It didn't wait for handle returns; it had no idea about the error. let reply = &reply.unwrap().expect_upgrade_regions_reply()[0]; assert!(!reply.ready); assert!(reply.exists); assert!(reply.error.is_none()); + // With a 200ms `replay_timeout` and a mock handle that returns the error + // immediately (no delay), the handler waits for the handle and propagates + // its error deterministically: the catchup future completes on its first + // poll, so the `replay_timeout` can never fire first. This avoids racing + // real wall-clock time (a 100ms delay vs 200ms timeout is flaky on + // Windows CI). + let mock_region_server = mock_region_server(); + let (mock_engine, _) = + MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |region_engine| { + // Region is not ready. + region_engine.mock_role = Some(Some(RegionRole::Follower)); + region_engine.handle_request_mock_fn = Some(Box::new(|_, _| { + error::UnexpectedSnafu { + violated: "mock_error".to_string(), + } + .fail() + })); + }); + mock_region_server.register_test_region(region_id, mock_engine); + let kv_backend = Arc::new(MemoryKvBackend::new()); + let handler_context = HandlerContext::new_for_test(mock_region_server, kv_backend); let reply = UpgradeRegionsHandler::new_test() .handle( &handler_context, diff --git a/tests/runner/src/cmd/datanode_overlay.rs b/tests/runner/src/cmd/datanode_overlay.rs index 21892b9a091..5feefd99632 100644 --- a/tests/runner/src/cmd/datanode_overlay.rs +++ b/tests/runner/src/cmd/datanode_overlay.rs @@ -106,13 +106,11 @@ impl DatanodeOverlay { )); } - let mut file = File::open(&canonical_target).map_err(|error| { - format!( - "Failed to open datanode overlay {}: {error}", - canonical_target.display() - ) - })?; - let metadata = file.metadata().map_err(|error| { + // Check the type before opening so the error is identical on every + // platform: `std::fs::metadata` succeeds on directories on both Unix + // and Windows, whereas `File::open` on a directory fails up front on + // Windows (with a platform-specific message). + let metadata = std::fs::metadata(&canonical_target).map_err(|error| { format!( "Failed to inspect datanode overlay {}: {error}", canonical_target.display() @@ -125,6 +123,13 @@ impl DatanodeOverlay { )); } + let mut file = File::open(&canonical_target).map_err(|error| { + format!( + "Failed to open datanode overlay {}: {error}", + canonical_target.display() + ) + })?; + let mut content = String::new(); file.read_to_string(&mut content).map_err(|error| { format!( @@ -613,7 +618,10 @@ mod tests { let directory_error = DatanodeOverlay::load(temp_dir.path(), Path::new("directory.toml")).unwrap_err(); - assert!(directory_error.contains("regular file")); + assert!( + directory_error.contains("regular file"), + "unexpected error for directory.toml: {directory_error}" + ); let parse_error = DatanodeOverlay::load(temp_dir.path(), Path::new("broken.toml")).unwrap_err(); assert!(parse_error.contains("Failed to parse datanode overlay")); From 4eeb4052d6e4ed486266f6c599e60280671a4aac Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 11 Aug 2026 14:20:08 +0800 Subject: [PATCH 19/92] feat(servers): stamp prometheus remote write v2 metadata as semantic table options (#8797) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(servers): stamp prometheus remote write v2 metadata as semantic table options Remote write 2.0 carries per-series metadata (type, unit, help) that the v2 ingest decoded and dropped; tables kept the name-based 'inferred' quality. Wire it into the semantic layer: - generalize the OTLP per-table semantic index into a shared, schema- aware servers::semantic module: v2 lets each series override its target schema, so the index is keyed {schema -> table -> options} and the same metric name in two schemas no longer collapses; - into_write_requests records metric type and unit per written table; an explicit type upgrades the table's metadata quality to declared, UNSPECIFIED series keep the request-level inferred stamp, and units are canonicalised from OpenMetrics words to the UCUM codes the vocabulary is defined in (unknown units are dropped, help text is not persisted); - the consumer folds the index in on both auto-create paths: the operator row-insert path and the pending-rows batched create, which bypasses the former. Table options are stamped at auto-create only; updating existing tables from later metadata is future work (a metadata registry, see the native histograms RFC). Signed-off-by: Dennis Zhuang * chore: trim over-commenting in the remote write metadata path Signed-off-by: Dennis Zhuang * fix: parse the per-table semantic index once per create round The index was re-parsed from JSON for every table being created — a first write creating N tables (a fleet's first scrape) paid O(N x index size). Parsing now happens lazily once per create-planning round, on both consumers: the operator row-insert auto-create (also serving OTLP metrics) and the pending-rows batched create. Also validate every non-zero metadata symbol reference up front, as the remote write 2.0 spec requires: help_ref was never checked, and unit_ref escaped checking when the metric type was UNSPECIFIED. Signed-off-by: Dennis Zhuang * fix(servers): stamp remote-write v2 unit independently of metric type OpenMetrics models TYPE and UNIT as independent MetricFamily metadata, and the Prometheus v2 sender emits UNSPECIFIED-type series that still carry a unit. The early return on UNSPECIFIED dropped that unit, which is unrecoverable after table auto-create (units are not stored in rows). Stamp the mapped unit whenever present; the type and the metadata_quality=declared upgrade still require an explicit type. Signed-off-by: Dennis Zhuang * chore: trim restating comments in the v2 metadata path Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/frontend/src/instance/otlp.rs | 5 +- src/frontend/src/instance/prom_store.rs | 15 +- src/operator/src/insert.rs | 142 +++++++--- src/servers/src/http/prom_store.rs | 23 +- src/servers/src/lib.rs | 1 + src/servers/src/otlp/metrics.rs | 20 +- src/servers/src/otlp/metrics/semantic.rs | 177 ------------ src/servers/src/prom_remote_write/v2.rs | 300 ++++++++++++++++++++- src/servers/src/semantic.rs | 325 +++++++++++++++++++++++ tests-integration/tests/http.rs | 97 ++++++- 10 files changed, 865 insertions(+), 240 deletions(-) delete mode 100644 src/servers/src/otlp/metrics/semantic.rs create mode 100644 src/servers/src/semantic.rs diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index c4616ae578d..459a05321cb 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -132,8 +132,9 @@ impl OpenTelemetryProtocolHandler for Instance { c.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC); c.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY); // Per-table metric specifics + resource/scope lineage ride this - // internal channel; the auto-create path folds them per table name. - if let Some(index) = semantic_index.encode() { + // internal channel; the auto-create path folds them per schema and + // table name. + if let Some(index) = semantic_index.encode(&c.current_schema()) { c.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index); } if !is_legacy { diff --git a/src/frontend/src/instance/prom_store.rs b/src/frontend/src/instance/prom_store.rs index 1bfb5f56437..428628df123 100644 --- a/src/frontend/src/instance/prom_store.rs +++ b/src/frontend/src/instance/prom_store.rs @@ -36,7 +36,8 @@ use common_query::prelude::{GREPTIME_PHYSICAL_TABLE, greptime_value}; use common_recordbatch::RecordBatches; use common_telemetry::{debug, tracing}; use operator::insert::{ - AutoCreateTableType, InserterRef, build_create_table_expr, fill_table_options_for_create, + AutoCreateTableType, InserterRef, PerTableSemanticIndex, apply_per_table_semantic_options, + build_create_table_expr, fill_table_options_for_create, parse_per_table_semantic_index, }; use operator::statement::StatementExecutor; use prost::Message; @@ -313,6 +314,7 @@ impl PendingRowsSchemaAlterer for Instance { // Check which tables actually still need to be created (may have been // concurrently created by another request). let mut create_exprs: Vec = Vec::with_capacity(tables.len()); + let mut per_table_semantics: Option> = None; for &(table_name, request_schema) in tables { let existing = self .catalog_manager() @@ -331,6 +333,17 @@ impl PendingRowsSchemaAlterer for Instance { let mut table_options = std::collections::HashMap::with_capacity(4); fill_table_options_for_create(&mut table_options, &create_type, &ctx); + // The batched create path bypasses the operator's auto-create, so + // fold the per-table semantic index in here too. + let semantic_index = per_table_semantics + .get_or_insert_with(|| parse_per_table_semantic_index(&ctx)) + .as_ref(); + apply_per_table_semantic_options( + &mut table_options, + semantic_index, + schema, + table_name, + ); create_table_expr.table_options.extend(table_options); create_exprs.push(create_table_expr); } diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index b5b6be093d9..fabc1b3778b 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::BTreeMap; use std::sync::Arc; use ahash::{HashMap, HashMapExt, HashSet, HashSetExt}; @@ -614,6 +615,7 @@ impl Inserter { let mut alter_tables = vec![]; let mut need_refresh_table_infos = HashSet::new(); let mut instant_table_ids = HashSet::new(); + let mut per_table_semantics: Option> = None; for req in &mut requests.inserts { match self.get_table(catalog, &schema, &req.table_name).await? { @@ -654,8 +656,15 @@ impl Inserter { table_infos.insert(table_info.table_id(), table_info); } None => { - let create_expr = - self.get_create_table_expr_on_demand(req, &auto_create_table_type, ctx)?; + let semantic_index = per_table_semantics + .get_or_insert_with(|| parse_per_table_semantic_index(ctx)) + .as_ref(); + let create_expr = self.get_create_table_expr_on_demand( + req, + &auto_create_table_type, + ctx, + semantic_index, + )?; create_tables.push(create_expr); } } @@ -933,11 +942,17 @@ impl Inserter { req: &RowInsertRequest, create_type: &AutoCreateTableType, ctx: &QueryContextRef, + semantic_index: Option<&PerTableSemanticIndex>, ) -> Result { let schema = ctx.current_schema(); let mut table_options = std::collections::HashMap::with_capacity(4); fill_table_options_for_create(&mut table_options, create_type, ctx); - apply_per_table_semantic_options(&mut table_options, ctx, &req.table_name); + apply_per_table_semantic_options( + &mut table_options, + semantic_index, + ctx.current_schema().as_str(), + &req.table_name, + ); let engine_name = if let AutoCreateTableType::Logical(_) = create_type { // engine should be metric engine when creating logical tables. @@ -1284,32 +1299,45 @@ pub fn fill_table_options_for_create( } } -/// Folds the semantic keys for `table_name` carried on the internal per-table -/// index extension into `table_options`. +/// The parsed per-table semantic index: `{schema -> {table -> {key -> value}}}`, +/// produced by the OTLP metrics encode path (where one metric can fan out into +/// several tables with distinct keys) and the Prometheus remote write v2 path +/// (where per-series metadata declares type/unit, and a series may override its +/// target schema). +pub type PerTableSemanticIndex = BTreeMap>>; + +/// Parses the per-table semantic index off the context extension. Call once per +/// create-planning round: a first write creating N tables would otherwise +/// re-parse the whole index N times. `None` when the request carries no index +/// (logs, traces, Prom RW v1) or it fails to parse. +pub fn parse_per_table_semantic_index(ctx: &QueryContextRef) -> Option { + let raw = ctx.extension(SEMANTIC_PER_TABLE_INDEX_KEY)?; + match serde_json::from_str(raw) { + Ok(index) => Some(index), + Err(_) => { + warn!("failed to parse semantic per-table index, skipping per-table options"); + None + } + } +} + +/// Folds the semantic keys of the table being created into `table_options`. /// -/// The index is a `{table_name -> {semantic_key: value}}` JSON blob produced by -/// the OTLP metrics encode path (where one metric can fan out into several -/// tables with distinct keys). Common keys shared by every table in a request -/// travel as plain semantic extensions and are handled by -/// [`fill_table_options_for_create`]; this carries only the per-table tail. -/// Keys are re-checked against the vocabulary defensively. Ingestion paths -/// without a per-table index (logs, traces, Prom RW) carry no extension, so this -/// is a no-op for them. -fn apply_per_table_semantic_options( +/// Common keys shared by every table in a request travel as plain semantic +/// extensions and are handled by [`fill_table_options_for_create`]; this +/// carries only the per-table tail and is applied after it, so a per-table +/// value (e.g. `declared` quality) wins. Keys are re-checked against the +/// vocabulary defensively. +pub fn apply_per_table_semantic_options( table_options: &mut std::collections::HashMap, - ctx: &QueryContextRef, + index: Option<&PerTableSemanticIndex>, + schema: &str, table_name: &str, ) { - let Some(raw) = ctx.extension(SEMANTIC_PER_TABLE_INDEX_KEY) else { - return; - }; - let Ok(index) = serde_json::from_str::< - std::collections::BTreeMap>, - >(raw) else { - warn!("failed to parse semantic per-table index, skipping per-table options"); - return; - }; - let Some(entry) = index.get(table_name) else { + let Some(entry) = index + .and_then(|index| index.get(schema)) + .and_then(|tables| tables.get(table_name)) + else { return; }; for (key, value) in entry { @@ -1734,22 +1762,41 @@ mod tests { SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, SEMANTIC_PER_TABLE_INDEX_KEY, }; - let index = r#"{ - "http_requests_total": { - "greptime.semantic.metric.type": "counter", - "greptime.semantic.metric.unit": "By", - "greptime.semantic.metric.type_BOGUS": "x" - }, - "other_table": { - "greptime.semantic.metric.type": "gauge" - } - }"#; + let index = format!( + r#"{{ + "{DEFAULT_SCHEMA_NAME}": {{ + "http_requests_total": {{ + "greptime.semantic.metric.type": "counter", + "greptime.semantic.metric.unit": "By", + "greptime.semantic.metric.type_BOGUS": "x" + }}, + "other_table": {{ + "greptime.semantic.metric.type": "gauge" + }} + }}, + "other_schema": {{ + "http_requests_total": {{ + "greptime.semantic.metric.type": "gauge" + }} + }} + }}"# + ); let mut ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME); ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index); let ctx = Arc::new(ctx); + let index = parse_per_table_semantic_index(&ctx); + assert!(index.is_some()); + let index = index.as_ref(); + let mut table_options = std::collections::HashMap::new(); - apply_per_table_semantic_options(&mut table_options, &ctx, "http_requests_total"); + apply_per_table_semantic_options( + &mut table_options, + index, + DEFAULT_SCHEMA_NAME, + "http_requests_total", + ); + // The write schema's entry applies — not other_schema's `gauge`. assert_eq!( table_options.get(SEMANTIC_METRIC_TYPE).map(String::as_str), Some("counter") @@ -1764,16 +1811,33 @@ mod tests { assert_eq!(table_options.len(), 2); let mut empty = std::collections::HashMap::new(); - apply_per_table_semantic_options(&mut empty, &ctx, "not_in_index"); + apply_per_table_semantic_options(&mut empty, index, DEFAULT_SCHEMA_NAME, "not_in_index"); assert!(empty.is_empty()); - // No extension at all is a no-op (e.g. logs / Prom RW). + // A schema with no entry is a no-op even when the table name matches + // elsewhere. + let mut opts = std::collections::HashMap::new(); + apply_per_table_semantic_options( + &mut opts, + index, + "schema_without_entry", + "http_requests_total", + ); + assert!(opts.is_empty()); + + // No extension at all parses to no index (e.g. logs / Prom RW v1). let bare = Arc::new(QueryContext::with( DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, )); + assert!(parse_per_table_semantic_index(&bare).is_none()); let mut opts = std::collections::HashMap::new(); - apply_per_table_semantic_options(&mut opts, &bare, "http_requests_total"); + apply_per_table_semantic_options( + &mut opts, + None, + DEFAULT_SCHEMA_NAME, + "http_requests_total", + ); assert!(opts.is_empty()); } diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index f659656e03e..ac4f39628d1 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -35,8 +35,9 @@ use serde::{Deserialize, Serialize}; use session::context::{Channel, QueryContext, QueryContextRef}; use snafu::prelude::*; use table::requests::{ - METADATA_QUALITY_INFERRED, SEMANTIC_METRIC_METADATA_QUALITY, SEMANTIC_SIGNAL_TYPE, - SEMANTIC_SOURCE, SEMANTIC_SOURCE_VERSION, SIGNAL_TYPE_METRIC, SOURCE_PROMETHEUS, + METADATA_QUALITY_INFERRED, SEMANTIC_METRIC_METADATA_QUALITY, SEMANTIC_PER_TABLE_INDEX_KEY, + SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SEMANTIC_SOURCE_VERSION, SIGNAL_TYPE_METRIC, + SOURCE_PROMETHEUS, }; use crate::error::{self, InternalSnafu, PipelineSnafu, Result}; @@ -155,6 +156,7 @@ async fn remote_write_v1( let (db, query_ctx, _timer) = prepare_remote_write_context(¶ms, query_ctx, REMOTE_WRITE_V1_VERSION); + let query_ctx = Arc::new(query_ctx); let mut processor = PromSeriesProcessor::default_processor(); @@ -230,7 +232,7 @@ async fn remote_write_v2( // optional pipeline parameter and ingest samples directly. let _ = pipeline_info; - let (db, query_ctx, _timer) = + let (db, mut query_ctx, _timer) = prepare_remote_write_context(¶ms, query_ctx, REMOTE_WRITE_V2_VERSION); let request = match decode_remote_write_v2_request(is_zstd, body) { @@ -253,6 +255,12 @@ async fn remote_write_v2( Ok(req) => req, Err(error) => return Ok(remote_write_v2_error_response(error, 0, 0, 0)), }; + // The v2 per-series metadata upgrades the written tables' semantic options + // (metric type/unit, declared quality) at auto-create time. + if let Some(index) = req.semantic_index.encode(&query_ctx.current_schema()) { + query_ctx.set_extension(SEMANTIC_PER_TABLE_INDEX_KEY, index); + } + let query_ctx = Arc::new(query_ctx); let sample_count = req.sample_count; let histogram_count = req.histogram_count; let sample_batches = into_prom_write_batches(req.samples, query_ctx.clone()); @@ -310,11 +318,14 @@ fn vm_proto_version_response(params: &RemoteWriteQuery) -> Option (String, Arc, HistogramTimer) { +) -> (String, QueryContext, HistogramTimer) { let db = params.db.clone().unwrap_or_default(); query_ctx.set_channel(Channel::Prometheus); let physical_table = params @@ -325,12 +336,12 @@ fn prepare_remote_write_context( // Stamp the Prometheus metric identity here, before `as_req_iter` splits into the // batched and direct write paths, so both inherit it (the batched path bypasses // `PromStoreProtocolHandler::write`). Prometheus remote-write metadata is weak - // here, so the type is inferred from naming. + // here, so the type is inferred from naming; v2 upgrades tables whose series + // carry inline metadata via the per-table index. query_ctx.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC); query_ctx.set_extension(SEMANTIC_SOURCE, SOURCE_PROMETHEUS); query_ctx.set_extension(SEMANTIC_SOURCE_VERSION, remote_write_version); query_ctx.set_extension(SEMANTIC_METRIC_METADATA_QUALITY, METADATA_QUALITY_INFERRED); - let query_ctx = Arc::new(query_ctx); let timer = crate::metrics::METRIC_HTTP_PROM_STORE_WRITE_ELAPSED .with_label_values(&[db.as_str()]) .start_timer(); diff --git a/src/servers/src/lib.rs b/src/servers/src/lib.rs index 55d3b65fed5..2a3796aab21 100644 --- a/src/servers/src/lib.rs +++ b/src/servers/src/lib.rs @@ -50,6 +50,7 @@ pub mod repeated_field; pub mod request_memory_limiter; pub mod request_memory_metrics; mod row_writer; +pub mod semantic; pub mod server; pub mod tls; diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index 56ab0530338..b43cc73338f 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -29,11 +29,14 @@ use table::requests::{ use crate::error::Result; use crate::otlp::trace::{KEY_SERVICE_INSTANCE_ID, KEY_SERVICE_NAME}; use crate::row_writer::{self, MultiTableData, TableData}; +pub use crate::semantic::SemanticIndex; +use crate::semantic::{ + METRIC_TYPE_COUNTER, METRIC_TYPE_GAUGE, METRIC_TYPE_HISTOGRAM, METRIC_TYPE_SUMMARY, + METRIC_TYPE_UPDOWN_COUNTER, +}; -mod semantic; mod translator; -pub use semantic::SemanticIndex; pub use translator::legacy_normalize_otlp_name; use translator::{translate_label_name, translate_metric_name}; @@ -44,15 +47,6 @@ const COUNT_TABLE_SUFFIX: &str = "_count"; const SUM_TABLE_SUFFIX: &str = "_sum"; const BUCKET_TABLE_SUFFIX: &str = "_bucket"; -// `greptime.semantic.metric.type` values stamped per emitted table. Must stay -// within the domain accepted by `validate_semantic_option`; the drift-guard test -// asserts this. -const METRIC_TYPE_COUNTER: &str = "counter"; -const METRIC_TYPE_UPDOWN_COUNTER: &str = "updown_counter"; -const METRIC_TYPE_GAUGE: &str = "gauge"; -const METRIC_TYPE_HISTOGRAM: &str = "histogram"; -const METRIC_TYPE_SUMMARY: &str = "summary"; - const JOB_KEY: &str = "job"; const INSTANCE_KEY: &str = "instance"; @@ -1183,7 +1177,9 @@ mod tests { use table::requests::validate_semantic_option; fn decode(index: &SemanticIndex) -> BTreeMap> { - serde_json::from_str(&index.encode().expect("non-empty index")).unwrap() + let nested: BTreeMap>> = + serde_json::from_str(&index.encode("public").expect("non-empty index")).unwrap(); + nested.into_values().next().unwrap() } fn record(metric: &Metric, metric_type: MetricType, name: &str) -> SemanticIndex { diff --git a/src/servers/src/otlp/metrics/semantic.rs b/src/servers/src/otlp/metrics/semantic.rs deleted file mode 100644 index 0726dc15c12..00000000000 --- a/src/servers/src/otlp/metrics/semantic.rs +++ /dev/null @@ -1,177 +0,0 @@ -// Copyright 2023 Greptime Team -// -// 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. - -//! Per-table semantic metadata accumulated during one OTLP metrics encode pass. -//! -//! A metric emits one or more tables (histogram and summary fan out into -//! `_bucket`/`_sum`/`_count` companions). Each emitted table collects the -//! metric's scalar semantic keys. The resulting index is serialized onto the -//! `greptime.internal.semantic.per_table_index` context extension and folded -//! into each table's options at auto-create time. -//! -//! Conflict handling follows the RFC: when two sources disagree on a -//! single-valued key the value collapses to `mixed` (or `unknown` for keys whose -//! domain has no `mixed`). - -use std::collections::{BTreeMap, HashMap}; - -use table::requests::{SEMANTIC_VALUE_MIXED, SEMANTIC_VALUE_UNKNOWN, validate_semantic_option}; - -/// Index of `{table_name -> {semantic_key -> value}}` built while encoding. -#[derive(Debug, Default)] -pub struct SemanticIndex { - /// Per-table scalar keys; conflicting values collapse to `mixed`/`unknown`. - tables: HashMap>, -} - -impl SemanticIndex { - pub fn is_empty(&self) -> bool { - self.tables.is_empty() - } - - /// Records a scalar semantic key for `table`. A value conflicting with one - /// already recorded collapses the key to `mixed`/`unknown`; once collapsed - /// it stays collapsed. - pub fn record_scalar(&mut self, table: &str, key: &'static str, value: &str) { - // Avoid allocating the table name (and an empty map) on the common path - // where the table is already present. - if let Some(scalars) = self.tables.get_mut(table) { - match scalars.get(key).map(String::as_str) { - Some(existing) if existing == value => {} - Some(SEMANTIC_VALUE_MIXED) | Some(SEMANTIC_VALUE_UNKNOWN) => {} - Some(_) => { - scalars.insert(key, collapse_value(key)); - } - None => { - scalars.insert(key, value.to_string()); - } - } - } else { - self.tables.insert( - table.to_string(), - BTreeMap::from([(key, value.to_string())]), - ); - } - } - - /// Serializes to the JSON `{table -> {key -> value}}` carried on the context - /// extension, or `None` when nothing was recorded. - pub fn encode(&self) -> Option { - if self.tables.is_empty() { - return None; - } - serde_json::to_string(&self.tables).ok() - } - - #[cfg(test)] - fn options_of(&self, table: &str) -> Option<&BTreeMap<&'static str, String>> { - self.tables.get(table) - } -} - -/// The collapsed value for a conflicting scalar key: `mixed` when the key's -/// domain accepts it, else `unknown`. Uses the vocabulary validator as the -/// single source of truth for which keys allow `mixed`. -fn collapse_value(key: &str) -> String { - if validate_semantic_option(key, SEMANTIC_VALUE_MIXED) { - SEMANTIC_VALUE_MIXED.to_string() - } else { - SEMANTIC_VALUE_UNKNOWN.to_string() - } -} - -#[cfg(test)] -mod tests { - use table::requests::{ - SEMANTIC_METRIC_METADATA_QUALITY, SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, - }; - - use super::*; - - #[test] - fn test_scalar_recording_keeps_first_then_collapses_on_conflict() { - let mut index = SemanticIndex::default(); - index.record_scalar("t", SEMANTIC_METRIC_TYPE, "counter"); - index.record_scalar("t", SEMANTIC_METRIC_TYPE, "counter"); - assert_eq!( - index - .options_of("t") - .unwrap() - .get(SEMANTIC_METRIC_TYPE) - .map(String::as_str), - Some("counter") - ); - - // Conflict on a key whose domain has `mixed` collapses to `mixed`. - index.record_scalar("t", SEMANTIC_METRIC_TYPE, "gauge"); - assert_eq!( - index - .options_of("t") - .unwrap() - .get(SEMANTIC_METRIC_TYPE) - .map(String::as_str), - Some("mixed") - ); - // Further writes stay collapsed. - index.record_scalar("t", SEMANTIC_METRIC_TYPE, "histogram"); - assert_eq!( - index - .options_of("t") - .unwrap() - .get(SEMANTIC_METRIC_TYPE) - .map(String::as_str), - Some("mixed") - ); - } - - #[test] - fn test_scalar_conflict_without_mixed_domain_collapses_to_unknown() { - let mut index = SemanticIndex::default(); - index.record_scalar("t", SEMANTIC_METRIC_METADATA_QUALITY, "declared"); - index.record_scalar("t", SEMANTIC_METRIC_METADATA_QUALITY, "inferred"); - // metadata_quality accepts only declared/inferred/unknown, so a conflict - // is `unknown`. - assert_eq!( - index - .options_of("t") - .unwrap() - .get(SEMANTIC_METRIC_METADATA_QUALITY) - .map(String::as_str), - Some("unknown") - ); - } - - #[test] - fn test_encode_is_none_when_empty_and_round_trips() { - let index = SemanticIndex::default(); - assert!(index.is_empty()); - assert_eq!(index.encode(), None); - - let mut index = SemanticIndex::default(); - index.record_scalar("metric_a", SEMANTIC_METRIC_TYPE, "counter"); - index.record_scalar("metric_a", SEMANTIC_METRIC_UNIT, "By"); - let json = index.encode().unwrap(); - let parsed: BTreeMap> = - serde_json::from_str(&json).unwrap(); - let table = parsed.get("metric_a").unwrap(); - assert_eq!( - table.get(SEMANTIC_METRIC_TYPE).map(String::as_str), - Some("counter") - ); - assert_eq!( - table.get(SEMANTIC_METRIC_UNIT).map(String::as_str), - Some("By") - ); - } -} diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index cfa772f6919..9940ad2e7c2 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -17,10 +17,10 @@ use std::collections::hash_map::Entry; use ahash::{HashMap, HashMapExt, HashSet, HashSetExt}; use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; use api::greptime_proto::io::prometheus::write::v2::{ - BucketSpan, Histogram, Request, Sample, TimeSeries, + BucketSpan, Histogram, Request, Sample, TimeSeries, metadata, }; #[cfg(test)] -use api::greptime_proto::io::prometheus::write::v2::{Exemplar, Metadata, metadata}; +use api::greptime_proto::io::prometheus::write::v2::{Exemplar, Metadata}; use api::helper::ColumnDataTypeWrapper; use api::v1::value::ValueData; use api::v1::{ColumnSchema, ListValue, RowInsertRequest, Rows, SemanticType, Value}; @@ -31,6 +31,10 @@ use common_query::prelude::{greptime_native_histogram, greptime_timestamp, grept use pipeline::{ContextOpt, ContextReq}; use prost::Message; use snafu::{OptionExt, ResultExt, ensure}; +use table::requests::{ + METADATA_QUALITY_DECLARED, SEMANTIC_METRIC_METADATA_QUALITY, SEMANTIC_METRIC_TYPE, + SEMANTIC_METRIC_UNIT, +}; use crate::error::{self, Result}; use crate::prom_remote_write::row_builder::PromCtx; @@ -42,6 +46,11 @@ use crate::prom_store::{ PHYSICAL_TABLE_LABEL_ALT, SCHEMA_LABEL, }; use crate::row_writer::{self, TableData}; +use crate::semantic::{ + METRIC_TYPE_COUNTER, METRIC_TYPE_GAUGE, METRIC_TYPE_GAUGE_HISTOGRAM, METRIC_TYPE_HISTOGRAM, + METRIC_TYPE_INFO, METRIC_TYPE_STATESET, METRIC_TYPE_SUMMARY, SemanticIndexes, + openmetrics_unit_to_ucum, +}; type PromTags = Vec<(String, String)>; type ResolvedSeriesLabels = (PromCtx, String, PromTags); @@ -67,6 +76,9 @@ pub(crate) struct RemoteWriteV2WriteRequests { pub histograms: ContextReq, pub sample_count: u64, pub histogram_count: u64, + /// Per-table semantic metadata from the series' inline `Metadata`, folded + /// into table options at auto-create time. + pub semantic_index: SemanticIndexes, } /// Converts a PRW v2 request into normal sample writes and native histogram writes. @@ -92,6 +104,7 @@ pub(crate) fn into_write_requests(request: Request) -> Result Result 0 { histogram_count > 0 || histogram_tables @@ -166,9 +186,78 @@ pub(crate) fn into_write_requests(request: Request) -> Result Result<()> { + let Some(series_metadata) = &series.metadata else { + return Ok(()); + }; + // Symbol 0 is the mandatory empty string: no help / no unit. + if series_metadata.help_ref != 0 { + symbol_ref(symbols, series_metadata.help_ref, "metadata help")?; + } + let unit = if series_metadata.unit_ref != 0 { + Some(symbol_ref( + symbols, + series_metadata.unit_ref, + "metadata unit", + )?) + } else { + None + }; + let metric_type = metric_type_value(series_metadata.r#type); + let ucum = unit.and_then(|unit| openmetrics_unit_to_ucum(unit.trim())); + if metric_type.is_none() && ucum.is_none() { + return Ok(()); + } + + let index = index.index_for(prom_ctx.schema.as_deref()); + if let Some(metric_type) = metric_type { + index.record_scalar(table_name, SEMANTIC_METRIC_TYPE, metric_type); + index.record_scalar( + table_name, + SEMANTIC_METRIC_METADATA_QUALITY, + METADATA_QUALITY_DECLARED, + ); + } + if let Some(ucum) = ucum { + index.record_scalar(table_name, SEMANTIC_METRIC_UNIT, ucum); + } + Ok(()) +} + +/// The `greptime.semantic.metric.type` value for a wire metric type; `None` +/// for `UNSPECIFIED` (nothing was declared) and out-of-range values. +fn metric_type_value(wire_type: i32) -> Option<&'static str> { + match metadata::MetricType::try_from(wire_type).ok()? { + metadata::MetricType::Unspecified => None, + metadata::MetricType::Counter => Some(METRIC_TYPE_COUNTER), + metadata::MetricType::Gauge => Some(METRIC_TYPE_GAUGE), + metadata::MetricType::Histogram => Some(METRIC_TYPE_HISTOGRAM), + metadata::MetricType::Gaugehistogram => Some(METRIC_TYPE_GAUGE_HISTOGRAM), + metadata::MetricType::Summary => Some(METRIC_TYPE_SUMMARY), + metadata::MetricType::Info => Some(METRIC_TYPE_INFO), + metadata::MetricType::Stateset => Some(METRIC_TYPE_STATESET), + } +} + fn get_or_create_table_data( tables: &mut HashMap>, prom_ctx: PromCtx, @@ -1823,4 +1912,211 @@ mod tests { fn is_empty_list(value: Option) -> bool { matches!(value, Some(ValueData::ListValue(list)) if list.items.is_empty()) } + + fn push_test_symbol(symbols: &mut Vec, symbol: &str) -> u32 { + if let Some(idx) = symbols.iter().position(|s| s == symbol) { + return idx as u32; + } + let idx = symbols.len(); + symbols.push(symbol.to_string()); + idx as u32 + } + + /// One sample-carrying test series: `(labels, metric_type, unit)`. + type MetadataSeries<'a> = (Vec<(&'a str, &'a str)>, i32, Option<&'a str>); + + fn metadata_request(series: Vec>) -> Request { + let mut symbols = vec!["".to_string()]; + let timeseries = series + .into_iter() + .map(|(labels, metric_type, unit)| { + let mut labels_refs = Vec::with_capacity(labels.len() * 2); + for (name, value) in labels { + labels_refs.push(push_test_symbol(&mut symbols, name)); + labels_refs.push(push_test_symbol(&mut symbols, value)); + } + let unit_ref = unit.map_or(0, |unit| push_test_symbol(&mut symbols, unit)); + TimeSeries { + labels_refs, + samples: vec![Sample { + value: 1.0, + timestamp: 1000, + start_timestamp: 0, + }], + histograms: Vec::new(), + exemplars: Vec::new(), + metadata: Some(Metadata { + r#type: metric_type, + help_ref: 0, + unit_ref, + }), + } + }) + .collect(); + Request { + symbols, + timeseries, + } + } + + type DecodedIndex = std::collections::BTreeMap< + String, + std::collections::BTreeMap>, + >; + + fn decoded_index(request: Request) -> DecodedIndex { + let encoded = into_write_requests(request) + .unwrap() + .semantic_index + .encode("public") + .expect("non-empty semantic index"); + serde_json::from_str(&encoded).unwrap() + } + + #[test] + fn test_metadata_stamps_semantic_index() { + use table::requests::METADATA_QUALITY_DECLARED; + + let index = decoded_index(metadata_request(vec![ + ( + vec![(METRIC_NAME_LABEL, "http_requests_total")], + metadata::MetricType::Counter as i32, + Some("seconds"), + ), + ( + vec![(METRIC_NAME_LABEL, "queue_depth")], + metadata::MetricType::Gauge as i32, + // Outside the OpenMetrics base set: dropped, not passed through. + Some("requests"), + ), + ])); + + let tables = &index["public"]; + let typed = &tables["http_requests_total"]; + assert_eq!(typed[SEMANTIC_METRIC_TYPE], "counter"); + assert_eq!( + typed[SEMANTIC_METRIC_METADATA_QUALITY], + METADATA_QUALITY_DECLARED + ); + assert_eq!(typed[SEMANTIC_METRIC_UNIT], "s"); + + let unitless = &tables["queue_depth"]; + assert_eq!(unitless[SEMANTIC_METRIC_TYPE], "gauge"); + assert!(!unitless.contains_key(SEMANTIC_METRIC_UNIT)); + } + + #[test] + fn test_metadata_unspecified_stamps_unit_but_not_type() { + let index = decoded_index(metadata_request(vec![( + vec![(METRIC_NAME_LABEL, "untyped_total")], + metadata::MetricType::Unspecified as i32, + Some("seconds"), + )])); + let untyped = &index["public"]["untyped_total"]; + assert_eq!(untyped[SEMANTIC_METRIC_UNIT], "s"); + assert!(!untyped.contains_key(SEMANTIC_METRIC_TYPE)); + assert!(!untyped.contains_key(SEMANTIC_METRIC_METADATA_QUALITY)); + + let requests = into_write_requests(metadata_request(vec![( + vec![(METRIC_NAME_LABEL, "untyped_unitless_total")], + metadata::MetricType::Unspecified as i32, + None, + )])) + .unwrap(); + assert!(requests.semantic_index.is_empty()); + + // No metadata at all behaves the same. + let requests = into_write_requests(test_util::request_with_labels_and_samples( + vec![(METRIC_NAME_LABEL, "bare_total")], + vec![Sample { + value: 1.0, + timestamp: 1000, + start_timestamp: 0, + }], + )) + .unwrap(); + assert!(requests.semantic_index.is_empty()); + } + + #[test] + fn test_metadata_type_conflict_collapses_to_mixed() { + let index = decoded_index(metadata_request(vec![ + ( + vec![(METRIC_NAME_LABEL, "flappy_metric"), ("job", "a")], + metadata::MetricType::Counter as i32, + None, + ), + ( + vec![(METRIC_NAME_LABEL, "flappy_metric"), ("job", "b")], + metadata::MetricType::Gauge as i32, + None, + ), + ])); + assert_eq!( + index["public"]["flappy_metric"][SEMANTIC_METRIC_TYPE], + "mixed" + ); + } + + #[test] + fn test_metadata_schema_overrides_stay_apart() { + // The same metric name written into two schemas by one request must not + // collapse each other's metadata. + let index = decoded_index(metadata_request(vec![ + ( + vec![(METRIC_NAME_LABEL, "cpu_usage")], + metadata::MetricType::Counter as i32, + None, + ), + ( + vec![ + (METRIC_NAME_LABEL, "cpu_usage"), + (DATABASE_LABEL, "tenant_b"), + ], + metadata::MetricType::Gauge as i32, + None, + ), + ])); + assert_eq!( + index["public"]["cpu_usage"][SEMANTIC_METRIC_TYPE], + "counter" + ); + assert_eq!( + index["tenant_b"]["cpu_usage"][SEMANTIC_METRIC_TYPE], + "gauge" + ); + } + + #[test] + fn test_metadata_out_of_range_refs_are_rejected() { + let mut request = metadata_request(vec![( + vec![(METRIC_NAME_LABEL, "broken_total")], + metadata::MetricType::Counter as i32, + None, + )]); + request.timeseries[0].metadata.as_mut().unwrap().unit_ref = 999; + let err = into_write_requests(request).err().unwrap(); + assert!(err.to_string().contains("out of range"), "{err}"); + + // unit_ref must be validated even when the type is UNSPECIFIED + // (which persists nothing). + let mut request = metadata_request(vec![( + vec![(METRIC_NAME_LABEL, "broken_total")], + metadata::MetricType::Unspecified as i32, + None, + )]); + request.timeseries[0].metadata.as_mut().unwrap().unit_ref = 999; + let err = into_write_requests(request).err().unwrap(); + assert!(err.to_string().contains("out of range"), "{err}"); + + // help_ref is validated although help is never persisted. + let mut request = metadata_request(vec![( + vec![(METRIC_NAME_LABEL, "broken_total")], + metadata::MetricType::Counter as i32, + None, + )]); + request.timeseries[0].metadata.as_mut().unwrap().help_ref = 999; + let err = into_write_requests(request).err().unwrap(); + assert!(err.to_string().contains("out of range"), "{err}"); + } } diff --git a/src/servers/src/semantic.rs b/src/servers/src/semantic.rs new file mode 100644 index 00000000000..63501dffc31 --- /dev/null +++ b/src/servers/src/semantic.rs @@ -0,0 +1,325 @@ +// Copyright 2023 Greptime Team +// +// 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. + +//! Per-table semantic metadata accumulated during one ingest pass (OTLP +//! metrics, Prometheus remote write v2). +//! +//! Each written table collects the scalar semantic keys its wire metadata +//! declares (an OTLP histogram fans out into `_bucket`/`_sum`/`_count` +//! companions; a remote write series names one metric family). The resulting +//! index is serialized onto the `greptime.internal.semantic.per_table_index` +//! context extension as `{schema -> {table -> {key -> value}}}` and folded into +//! each table's options at auto-create time. +//! +//! Conflict handling follows the RFC: when two sources disagree on a +//! single-valued key the value collapses to `mixed` (or `unknown` for keys whose +//! domain has no `mixed`). + +use std::collections::{BTreeMap, HashMap}; + +use table::requests::{SEMANTIC_VALUE_MIXED, SEMANTIC_VALUE_UNKNOWN, validate_semantic_option}; + +// `greptime.semantic.metric.type` values stamped per emitted table. Must stay +// within the domain accepted by `validate_semantic_option`; the drift-guard test +// asserts this. +pub const METRIC_TYPE_COUNTER: &str = "counter"; +pub const METRIC_TYPE_UPDOWN_COUNTER: &str = "updown_counter"; +pub const METRIC_TYPE_GAUGE: &str = "gauge"; +pub const METRIC_TYPE_HISTOGRAM: &str = "histogram"; +pub const METRIC_TYPE_GAUGE_HISTOGRAM: &str = "gauge_histogram"; +pub const METRIC_TYPE_SUMMARY: &str = "summary"; +pub const METRIC_TYPE_INFO: &str = "info"; +pub const METRIC_TYPE_STATESET: &str = "stateset"; + +/// Maps an OpenMetrics unit name (open vocabulary full words: `seconds`, +/// `bytes`) to the UCUM code the `greptime.semantic.metric.unit` option is +/// defined in. Only the OpenMetrics base units are mapped; anything else is +/// dropped — a missing unit beats a corrupted cross-protocol vocabulary. +pub fn openmetrics_unit_to_ucum(unit: &str) -> Option<&'static str> { + Some(match unit { + "seconds" => "s", + "celsius" => "Cel", + "meters" => "m", + "bytes" => "By", + "ratios" => "1", + "volts" => "V", + "amperes" => "A", + "joules" => "J", + "grams" => "g", + _ => return None, + }) +} + +/// Index of `{table_name -> {semantic_key -> value}}` for one target schema. +#[derive(Debug, Default)] +pub struct SemanticIndex { + /// Per-table scalar keys; conflicting values collapse to `mixed`/`unknown`. + tables: HashMap>, +} + +impl SemanticIndex { + pub fn is_empty(&self) -> bool { + self.tables.is_empty() + } + + /// Records a scalar semantic key for `table`. A value conflicting with one + /// already recorded collapses the key to `mixed`/`unknown`; once collapsed + /// it stays collapsed. + pub fn record_scalar(&mut self, table: &str, key: &'static str, value: &str) { + // Avoid allocating the table name (and an empty map) on the common path + // where the table is already present. + if let Some(scalars) = self.tables.get_mut(table) { + match scalars.get(key).map(String::as_str) { + Some(existing) if existing == value => {} + Some(SEMANTIC_VALUE_MIXED) | Some(SEMANTIC_VALUE_UNKNOWN) => {} + Some(_) => { + scalars.insert(key, collapse_value(key)); + } + None => { + scalars.insert(key, value.to_string()); + } + } + } else { + self.tables.insert( + table.to_string(), + BTreeMap::from([(key, value.to_string())]), + ); + } + } + + /// Serializes to the JSON `{schema -> {table -> {key -> value}}}` carried + /// on the context extension, with every table under `schema`. `None` when + /// nothing was recorded. + pub fn encode(&self, schema: &str) -> Option { + if self.tables.is_empty() { + return None; + } + serde_json::to_string(&BTreeMap::from([(schema, &self.tables)])).ok() + } + + fn merge_from(&mut self, other: &SemanticIndex) { + for (table, scalars) in &other.tables { + for (key, value) in scalars { + self.record_scalar(table, key, value); + } + } + } + + #[cfg(test)] + fn options_of(&self, table: &str) -> Option<&BTreeMap<&'static str, String>> { + self.tables.get(table) + } +} + +/// Schema-aware collection of [`SemanticIndex`]es: Prometheus remote write lets +/// each series override the target schema with a special label, so one request +/// may write the same metric name into several schemas — their metadata must +/// not collapse into each other. +#[derive(Debug, Default)] +pub struct SemanticIndexes { + /// Tables written into the request's default schema (no override). + default: SemanticIndex, + /// Tables written under a per-series schema override. + overrides: HashMap, +} + +impl SemanticIndexes { + pub fn is_empty(&self) -> bool { + self.default.is_empty() && self.overrides.values().all(SemanticIndex::is_empty) + } + + /// The index for `schema` (`None` = the request's default schema). + pub fn index_for(&mut self, schema: Option<&str>) -> &mut SemanticIndex { + match schema { + None => &mut self.default, + Some(schema) => { + if !self.overrides.contains_key(schema) { + self.overrides + .insert(schema.to_string(), SemanticIndex::default()); + } + self.overrides.get_mut(schema).expect("just inserted") + } + } + } + + /// Serializes to the JSON `{schema -> {table -> {key -> value}}}` carried on + /// the context extension, resolving the default index to `default_schema`. + /// An override explicitly naming `default_schema` merges into the default + /// with the usual conflict collapse. + pub fn encode(&self, default_schema: &str) -> Option { + if self.is_empty() { + return None; + } + let mut by_schema: BTreeMap<&str, &HashMap>> = + BTreeMap::new(); + let mut merged_default; + if let Some(aliased) = self.overrides.get(default_schema) { + merged_default = SemanticIndex::default(); + merged_default.merge_from(&self.default); + merged_default.merge_from(aliased); + by_schema.insert(default_schema, &merged_default.tables); + } else if !self.default.is_empty() { + by_schema.insert(default_schema, &self.default.tables); + } + for (schema, index) in &self.overrides { + if schema != default_schema && !index.is_empty() { + by_schema.insert(schema, &index.tables); + } + } + serde_json::to_string(&by_schema).ok() + } +} + +/// The collapsed value for a conflicting scalar key: `mixed` when the key's +/// domain accepts it, else `unknown`. Uses the vocabulary validator as the +/// single source of truth for which keys allow `mixed`. +fn collapse_value(key: &str) -> String { + if validate_semantic_option(key, SEMANTIC_VALUE_MIXED) { + SEMANTIC_VALUE_MIXED.to_string() + } else { + SEMANTIC_VALUE_UNKNOWN.to_string() + } +} + +#[cfg(test)] +mod tests { + use table::requests::{ + SEMANTIC_METRIC_METADATA_QUALITY, SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, + }; + + use super::*; + + type Decoded = BTreeMap>>; + + #[test] + fn test_scalar_recording_keeps_first_then_collapses_on_conflict() { + let mut index = SemanticIndex::default(); + index.record_scalar("t", SEMANTIC_METRIC_TYPE, "counter"); + index.record_scalar("t", SEMANTIC_METRIC_TYPE, "counter"); + assert_eq!( + index + .options_of("t") + .unwrap() + .get(SEMANTIC_METRIC_TYPE) + .map(String::as_str), + Some("counter") + ); + + // Conflict on a key whose domain has `mixed` collapses to `mixed`. + index.record_scalar("t", SEMANTIC_METRIC_TYPE, "gauge"); + assert_eq!( + index + .options_of("t") + .unwrap() + .get(SEMANTIC_METRIC_TYPE) + .map(String::as_str), + Some("mixed") + ); + // Further writes stay collapsed. + index.record_scalar("t", SEMANTIC_METRIC_TYPE, "histogram"); + assert_eq!( + index + .options_of("t") + .unwrap() + .get(SEMANTIC_METRIC_TYPE) + .map(String::as_str), + Some("mixed") + ); + } + + #[test] + fn test_scalar_conflict_without_mixed_domain_collapses_to_unknown() { + let mut index = SemanticIndex::default(); + index.record_scalar("t", SEMANTIC_METRIC_METADATA_QUALITY, "declared"); + index.record_scalar("t", SEMANTIC_METRIC_METADATA_QUALITY, "inferred"); + // metadata_quality accepts only declared/inferred/unknown, so a conflict + // is `unknown`. + assert_eq!( + index + .options_of("t") + .unwrap() + .get(SEMANTIC_METRIC_METADATA_QUALITY) + .map(String::as_str), + Some("unknown") + ); + } + + #[test] + fn test_encode_is_none_when_empty_and_round_trips() { + let index = SemanticIndex::default(); + assert!(index.is_empty()); + assert_eq!(index.encode("public"), None); + + let mut index = SemanticIndex::default(); + index.record_scalar("metric_a", SEMANTIC_METRIC_TYPE, "counter"); + index.record_scalar("metric_a", SEMANTIC_METRIC_UNIT, "By"); + let json = index.encode("public").unwrap(); + let parsed: Decoded = serde_json::from_str(&json).unwrap(); + let table = parsed.get("public").unwrap().get("metric_a").unwrap(); + assert_eq!( + table.get(SEMANTIC_METRIC_TYPE).map(String::as_str), + Some("counter") + ); + assert_eq!( + table.get(SEMANTIC_METRIC_UNIT).map(String::as_str), + Some("By") + ); + } + + #[test] + fn test_indexes_keep_schemas_apart_and_merge_default_alias() { + let mut indexes = SemanticIndexes::default(); + assert!(indexes.is_empty()); + assert_eq!(indexes.encode("public"), None); + + // The same metric name in two schemas must not collapse to `mixed`. + indexes + .index_for(None) + .record_scalar("cpu_usage", SEMANTIC_METRIC_TYPE, "counter"); + indexes.index_for(Some("tenant_b")).record_scalar( + "cpu_usage", + SEMANTIC_METRIC_TYPE, + "gauge", + ); + let parsed: Decoded = serde_json::from_str(&indexes.encode("public").unwrap()).unwrap(); + assert_eq!( + parsed["public"]["cpu_usage"][SEMANTIC_METRIC_TYPE], + "counter" + ); + assert_eq!( + parsed["tenant_b"]["cpu_usage"][SEMANTIC_METRIC_TYPE], + "gauge" + ); + + // An override naming the default schema merges into it — and a real + // conflict then collapses. + indexes + .index_for(Some("public")) + .record_scalar("cpu_usage", SEMANTIC_METRIC_TYPE, "gauge"); + let parsed: Decoded = serde_json::from_str(&indexes.encode("public").unwrap()).unwrap(); + assert_eq!(parsed["public"]["cpu_usage"][SEMANTIC_METRIC_TYPE], "mixed"); + } + + #[test] + fn test_openmetrics_unit_mapping() { + assert_eq!(openmetrics_unit_to_ucum("seconds"), Some("s")); + assert_eq!(openmetrics_unit_to_ucum("bytes"), Some("By")); + assert_eq!(openmetrics_unit_to_ucum("ratios"), Some("1")); + // Outside the OpenMetrics base set: dropped, not passed through. + assert_eq!(openmetrics_unit_to_ucum("requests"), None); + assert_eq!(openmetrics_unit_to_ucum(""), None); + // No fuzzy matching: UCUM codes are not OpenMetrics names. + assert_eq!(openmetrics_unit_to_ucum("By"), None); + } +} diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index aa60109a616..9d2f1eb8cb5 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -18,8 +18,10 @@ use std::str::FromStr; use std::time::Duration; use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; +use api::greptime_proto::io::prometheus::write::v2::metadata::MetricType as RemoteWriteV2MetricType; use api::greptime_proto::io::prometheus::write::v2::{ - BucketSpan, Histogram, Sample as RemoteWriteV2Sample, TimeSeries as RemoteWriteV2TimeSeries, + BucketSpan, Histogram, Metadata as RemoteWriteV2Metadata, Sample as RemoteWriteV2Sample, + TimeSeries as RemoteWriteV2TimeSeries, }; use api::prom_store::remote::label_matcher::Type as MatcherType; use api::prom_store::remote::{ @@ -2693,6 +2695,52 @@ pub async fn test_prometheus_remote_write_v2(store_type: StorageType) { ) .await; + // A series carrying inline metadata (v2 senders SHOULD) upgrades its + // table's quality to declared, with the metric type and the unit + // canonicalised from the OpenMetrics word to UCUM. + let mut write_request = remote_write_v2::request_with_labels_and_samples( + vec![ + (prom_store::METRIC_NAME_LABEL, "remote_write_v2_typed_total"), + ("job", "api"), + ], + vec![RemoteWriteV2Sample { + value: 1.0, + timestamp: 1000, + start_timestamp: 0, + }], + ); + let unit_ref = write_request.symbols.len() as u32; + write_request.symbols.push("seconds".to_string()); + write_request.timeseries[0].metadata = Some(RemoteWriteV2Metadata { + r#type: RemoteWriteV2MetricType::Counter as i32, + help_ref: 0, + unit_ref, + }); + let compressed_request = prom_store::snappy_compress(&write_request.encode_to_vec()) + .expect("failed to encode snappy"); + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .header( + "Content-Type", + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + ) + .body(compressed_request) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + validate_data( + "prometheus_remote_write_v2_declared_metadata", + &client, + "select count(*) from information_schema.tables where table_name = 'remote_write_v2_typed_total' \ + and create_options like '%greptime.semantic.metric.type=counter%' \ + and create_options like '%greptime.semantic.metric.unit=s%' \ + and create_options like '%greptime.semantic.metric.metadata_quality=declared%';", + "[[1]]", + ) + .await; + guard.remove_all().await; } @@ -2892,6 +2940,53 @@ pub async fn test_prometheus_remote_write_batched(store_type: StorageType) { ) .await; + // The batched path bypasses the operator's auto-create: v2 per-series + // metadata must reach its table options through the batch create too. + let mut write_request = remote_write_v2::request_with_labels_and_samples( + vec![ + ( + prom_store::METRIC_NAME_LABEL, + "remote_write_batched_typed_total", + ), + ("job", "api"), + ], + vec![RemoteWriteV2Sample { + value: 1.0, + timestamp: 1000, + start_timestamp: 0, + }], + ); + let unit_ref = write_request.symbols.len() as u32; + write_request.symbols.push("bytes".to_string()); + write_request.timeseries[0].metadata = Some(RemoteWriteV2Metadata { + r#type: RemoteWriteV2MetricType::Gauge as i32, + help_ref: 0, + unit_ref, + }); + let compressed_request = prom_store::snappy_compress(&write_request.encode_to_vec()) + .expect("failed to encode snappy"); + let res = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .header( + "Content-Type", + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + ) + .body(compressed_request) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + wait_for_data( + &client, + "select count(*) from information_schema.tables where table_name = 'remote_write_batched_typed_total' \ + and create_options like '%greptime.semantic.metric.type=gauge%' \ + and create_options like '%greptime.semantic.metric.unit=By%' \ + and create_options like '%greptime.semantic.metric.metadata_quality=declared%'", + "[[1]]", + ) + .await; + guard.remove_all().await; } From 32e215cad0aa0803991d751bc7b4f44ac1865e62 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 11 Aug 2026 15:12:45 +0800 Subject: [PATCH 20/92] fix(event): preserve procedure lifecycle locators (#8787) * fix(event): preserve procedure lifecycle locators Signed-off-by: WenyXu * fix(event): preserve dropped table lifecycle locators Signed-off-by: WenyXu * test(event): cover lifecycle locators Signed-off-by: WenyXu * test(event): fix lifecycle context expectations Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/common/meta/src/ddl/alter_database.rs | 2 +- .../meta/src/ddl/alter_logical_tables.rs | 13 +- src/common/meta/src/ddl/alter_table.rs | 9 +- src/common/meta/src/ddl/create_database.rs | 2 +- src/common/meta/src/ddl/create_flow.rs | 11 +- .../meta/src/ddl/create_logical_tables.rs | 66 +++++----- src/common/meta/src/ddl/create_table.rs | 17 +-- src/common/meta/src/ddl/create_view.rs | 29 ++++- src/common/meta/src/ddl/drop_database.rs | 2 +- src/common/meta/src/ddl/drop_flow.rs | 6 +- src/common/meta/src/ddl/drop_table.rs | 20 ++- src/common/meta/src/ddl/drop_view.rs | 10 +- src/common/meta/src/ddl/event/database.rs | 18 +-- src/common/meta/src/ddl/event/flow.rs | 31 +++-- src/common/meta/src/ddl/event/table.rs | 31 +++-- src/common/meta/src/ddl/event/view.rs | 64 +++++++--- .../meta/src/ddl/purge_dropped_table.rs | 29 +++-- .../meta/src/ddl/tests/event/database.rs | 29 +++-- src/common/meta/src/ddl/tests/event/flow.rs | 43 +++++-- src/common/meta/src/ddl/tests/event/table.rs | 44 +++++-- src/common/meta/src/ddl/tests/event/view.rs | 114 +++++++++++++----- src/common/meta/src/ddl/truncate_table.rs | 20 ++- src/common/meta/src/ddl/undrop_table.rs | 33 ++--- src/meta-srv/src/event/repartition.rs | 54 +++++---- src/meta-srv/src/procedure/repartition.rs | 2 +- .../src/procedure/repartition/group.rs | 2 +- tests-integration/tests/database_ddl_event.rs | 4 +- tests-integration/tests/flow_ddl_event.rs | 10 +- tests-integration/tests/repartition_event.rs | 49 ++++++++ tests-integration/tests/table_ddl_event.rs | 99 +++++++-------- tests-integration/tests/view_ddl_event.rs | 14 +-- 31 files changed, 562 insertions(+), 315 deletions(-) diff --git a/src/common/meta/src/ddl/alter_database.rs b/src/common/meta/src/ddl/alter_database.rs index 7e16275da3c..f9344213025 100644 --- a/src/common/meta/src/ddl/alter_database.rs +++ b/src/common/meta/src/ddl/alter_database.rs @@ -196,7 +196,7 @@ impl Procedure for AlterDatabaseProcedure { self.data.event_context.clone(), ) } else { - DatabaseDdlEvent::alter_lifecycle() + DatabaseDdlEvent::alter_lifecycle(self.data.catalog(), self.data.schema()) }; Some(Box::new(event)) } diff --git a/src/common/meta/src/ddl/alter_logical_tables.rs b/src/common/meta/src/ddl/alter_logical_tables.rs index 816db0a6afb..f2f15430933 100644 --- a/src/common/meta/src/ddl/alter_logical_tables.rs +++ b/src/common/meta/src/ddl/alter_logical_tables.rs @@ -335,6 +335,11 @@ impl Procedure for AlterLogicalTablesProcedure { if ctx.trigger != EventTrigger::Submitted { return Some(Box::new(TableDdlEvent::lifecycle( TableDdlEventType::AlterLogicalTables, + self.data.tasks.iter().map(|task| { + let table_ref = task.table_ref(); + TableDdlLocator::new(table_ref.catalog, table_ref.schema, table_ref.table) + .with_physical_table_id(self.data.physical_table_id) + }), ))); } @@ -375,12 +380,12 @@ pub struct AlterTablesData { } impl AlterTablesData { - /// Clears all data fields except `state` and `table_cache_keys_to_invalidate` after metadata update. - /// This is done to avoid persisting unnecessary data after the update metadata step. + /// Clears metadata snapshots after the update metadata step. + /// + /// Keep the tasks and physical table ID until the procedure finishes: lifecycle + /// events use them to retain the logical table locator. fn clear_metadata_fields(&mut self) { - self.tasks.clear(); self.table_info_values.clear(); - self.physical_table_id = 0; self.physical_table_info = None; self.physical_columns.clear(); } diff --git a/src/common/meta/src/ddl/alter_table.rs b/src/common/meta/src/ddl/alter_table.rs index acc02c4d1a8..21975fab714 100644 --- a/src/common/meta/src/ddl/alter_table.rs +++ b/src/common/meta/src/ddl/alter_table.rs @@ -579,12 +579,11 @@ impl Procedure for AlterTableProcedure { { return None; } + let table_ref = self.data.table_ref(); + let locator = TableDdlLocator::new(table_ref.catalog, table_ref.schema, table_ref.table) + .with_table_id(self.data.table_id()); let event = match &ctx.trigger { EventTrigger::Submitted => { - let table_ref = self.data.table_ref(); - let locator = - TableDdlLocator::new(table_ref.catalog, table_ref.schema, table_ref.table) - .with_table_id(self.data.table_id()); let kind = self .data .task @@ -594,7 +593,7 @@ impl Procedure for AlterTableProcedure { .and_then(alter_table_kind_name); TableDdlEvent::alter_table_submitted(locator, kind, self.data.event_context.clone()) } - _ => TableDdlEvent::lifecycle(TableDdlEventType::AlterTable), + _ => TableDdlEvent::lifecycle(TableDdlEventType::AlterTable, [locator]), }; Some(Box::new(event)) diff --git a/src/common/meta/src/ddl/create_database.rs b/src/common/meta/src/ddl/create_database.rs index 8f9c5bee1a5..b48e700a519 100644 --- a/src/common/meta/src/ddl/create_database.rs +++ b/src/common/meta/src/ddl/create_database.rs @@ -297,7 +297,7 @@ impl Procedure for CreateDatabaseProcedure { self.data.event_context.clone(), ) } else { - DatabaseDdlEvent::create_lifecycle() + DatabaseDdlEvent::create_lifecycle(&self.data.catalog, &self.data.schema) }; Some(Box::new(event)) } diff --git a/src/common/meta/src/ddl/create_flow.rs b/src/common/meta/src/ddl/create_flow.rs index 66083f5949c..3ef4dadbc63 100644 --- a/src/common/meta/src/ddl/create_flow.rs +++ b/src/common/meta/src/ddl/create_flow.rs @@ -440,9 +440,16 @@ impl Procedure for CreateFlowProcedure { .or(self.data.flow_id), _ => self.data.flow_id, }; - FlowDdlEvent::create_succeeded(flow_id) + FlowDdlEvent::create_succeeded( + &self.data.task.catalog_name, + &self.data.task.flow_name, + flow_id, + ) } - _ => FlowDdlEvent::create_lifecycle(), + _ => FlowDdlEvent::create_lifecycle( + &self.data.task.catalog_name, + &self.data.task.flow_name, + ), }; Some(Box::new(event)) diff --git a/src/common/meta/src/ddl/create_logical_tables.rs b/src/common/meta/src/ddl/create_logical_tables.rs index 8b93ecfcc0f..1c9efb865e4 100644 --- a/src/common/meta/src/ddl/create_logical_tables.rs +++ b/src/common/meta/src/ddl/create_logical_tables.rs @@ -212,6 +212,19 @@ impl CreateLogicalTablesProcedure { } } +impl CreateLogicalTablesProcedure { + fn event_locators(&self) -> impl Iterator + '_ { + self.data.tasks.iter().map(|task| { + TableDdlLocator::new( + &task.create_table.catalog_name, + &task.create_table.schema_name, + &task.create_table.table_name, + ) + .with_physical_table_id(self.data.physical_table_id) + }) + } +} + #[async_trait] impl Procedure for CreateLogicalTablesProcedure { fn type_name(&self) -> &str { @@ -268,21 +281,11 @@ impl Procedure for CreateLogicalTablesProcedure { return None; } let event = match &ctx.trigger { - EventTrigger::Submitted => { - let locators = self.data.tasks.iter().map(|task| { - TableDdlLocator::new( - &task.create_table.catalog_name, - &task.create_table.schema_name, - &task.create_table.table_name, - ) - .with_physical_table_id(self.data.physical_table_id) - }); - TableDdlEvent::create_logical_tables_submitted( - locators, - self.data.tasks.len(), - self.data.event_context.clone(), - ) - } + EventTrigger::Submitted => TableDdlEvent::create_logical_tables_submitted( + self.event_locators(), + self.data.tasks.len(), + self.data.event_context.clone(), + ), EventTrigger::Succeeded => match ctx.lifecycle_state { ProcedureState::Done { output: Some(output), @@ -290,28 +293,27 @@ impl Procedure for CreateLogicalTablesProcedure { .downcast_ref::>() .map(|table_ids| { debug_assert_eq!(self.data.tasks.len(), table_ids.len()); - let locators = - self.data - .tasks - .iter() - .zip(table_ids) - .map(|(task, table_id)| { - TableDdlLocator::new( - &task.create_table.catalog_name, - &task.create_table.schema_name, - &task.create_table.table_name, - ) - .with_table_id(*table_id) - .with_physical_table_id(self.data.physical_table_id) - }); + let locators = self + .event_locators() + .zip(table_ids) + .map(|(locator, table_id)| locator.with_table_id(*table_id)); TableDdlEvent::create_logical_tables_succeeded(locators) }) .unwrap_or_else(|| { - TableDdlEvent::lifecycle(TableDdlEventType::CreateLogicalTables) + TableDdlEvent::lifecycle( + TableDdlEventType::CreateLogicalTables, + self.event_locators(), + ) }), - _ => TableDdlEvent::lifecycle(TableDdlEventType::CreateLogicalTables), + _ => TableDdlEvent::lifecycle( + TableDdlEventType::CreateLogicalTables, + self.event_locators(), + ), }, - _ => TableDdlEvent::lifecycle(TableDdlEventType::CreateLogicalTables), + _ => TableDdlEvent::lifecycle( + TableDdlEventType::CreateLogicalTables, + self.event_locators(), + ), }; Some(Box::new(event)) diff --git a/src/common/meta/src/ddl/create_table.rs b/src/common/meta/src/ddl/create_table.rs index c1690416a2e..c63d7f7a34d 100644 --- a/src/common/meta/src/ddl/create_table.rs +++ b/src/common/meta/src/ddl/create_table.rs @@ -406,11 +406,10 @@ impl Procedure for CreateTableProcedure { { return None; } + let table_ref = self.data.table_ref(); + let locator = TableDdlLocator::new(table_ref.catalog, table_ref.schema, table_ref.table); let event = match &ctx.trigger { EventTrigger::Submitted => { - let table_ref = self.data.table_ref(); - let locator = - TableDdlLocator::new(table_ref.catalog, table_ref.schema, table_ref.table); let create_table = &self.data.task.create_table; TableDdlEvent::create_table_submitted( locator, @@ -425,11 +424,15 @@ impl Procedure for CreateTableProcedure { } => output .downcast_ref::() .copied() - .map(TableDdlEvent::create_table_succeeded) - .unwrap_or_else(|| TableDdlEvent::lifecycle(TableDdlEventType::CreateTable)), - _ => TableDdlEvent::lifecycle(TableDdlEventType::CreateTable), + .map(|table_id| { + TableDdlEvent::create_table_succeeded(locator.clone(), table_id) + }) + .unwrap_or_else(|| { + TableDdlEvent::lifecycle(TableDdlEventType::CreateTable, [locator.clone()]) + }), + _ => TableDdlEvent::lifecycle(TableDdlEventType::CreateTable, [locator.clone()]), }, - _ => TableDdlEvent::lifecycle(TableDdlEventType::CreateTable), + _ => TableDdlEvent::lifecycle(TableDdlEventType::CreateTable, [locator]), }; Some(Box::new(event)) diff --git a/src/common/meta/src/ddl/create_view.rs b/src/common/meta/src/ddl/create_view.rs index 0cb622774f0..a7175854b76 100644 --- a/src/common/meta/src/ddl/create_view.rs +++ b/src/common/meta/src/ddl/create_view.rs @@ -297,12 +297,33 @@ impl Procedure for CreateViewProcedure { ProcedureState::Done { output: Some(output), } => output.downcast_ref::().copied().map_or_else( - ViewDdlEvent::create_lifecycle, - ViewDdlEvent::create_succeeded, + || { + ViewDdlEvent::create_lifecycle( + &self.data.task.create_view.catalog_name, + &self.data.task.create_view.schema_name, + &self.data.task.create_view.view_name, + ) + }, + |view_id| { + ViewDdlEvent::create_succeeded( + &self.data.task.create_view.catalog_name, + &self.data.task.create_view.schema_name, + &self.data.task.create_view.view_name, + view_id, + ) + }, + ), + _ => ViewDdlEvent::create_lifecycle( + &self.data.task.create_view.catalog_name, + &self.data.task.create_view.schema_name, + &self.data.task.create_view.view_name, ), - _ => ViewDdlEvent::create_lifecycle(), }, - _ => ViewDdlEvent::create_lifecycle(), + _ => ViewDdlEvent::create_lifecycle( + &self.data.task.create_view.catalog_name, + &self.data.task.create_view.schema_name, + &self.data.task.create_view.view_name, + ), }; Some(Box::new(event)) diff --git a/src/common/meta/src/ddl/drop_database.rs b/src/common/meta/src/ddl/drop_database.rs index d84a5e6dc86..7772e966cb2 100644 --- a/src/common/meta/src/ddl/drop_database.rs +++ b/src/common/meta/src/ddl/drop_database.rs @@ -202,7 +202,7 @@ impl Procedure for DropDatabaseProcedure { self.event_context.clone(), ) } else { - DatabaseDdlEvent::drop_lifecycle() + DatabaseDdlEvent::drop_lifecycle(&self.context.catalog, &self.context.schema) }; Some(Box::new(event)) } diff --git a/src/common/meta/src/ddl/drop_flow.rs b/src/common/meta/src/ddl/drop_flow.rs index badded386ad..99c18ebc13b 100644 --- a/src/common/meta/src/ddl/drop_flow.rs +++ b/src/common/meta/src/ddl/drop_flow.rs @@ -239,7 +239,11 @@ impl Procedure for DropFlowProcedure { self.data.task.drop_if_exists, self.data.event_context.clone(), ), - _ => FlowDdlEvent::drop_lifecycle(), + _ => FlowDdlEvent::drop_lifecycle( + &self.data.task.catalog_name, + &self.data.task.flow_name, + self.data.task.flow_id, + ), }; Some(Box::new(event)) diff --git a/src/common/meta/src/ddl/drop_table.rs b/src/common/meta/src/ddl/drop_table.rs index 2dcfa2f7d82..7c8dec55e49 100644 --- a/src/common/meta/src/ddl/drop_table.rs +++ b/src/common/meta/src/ddl/drop_table.rs @@ -421,18 +421,16 @@ impl Procedure for DropTableProcedure { { return None; } + let task = &self.data.task; + let locator = TableDdlLocator::new(&task.catalog, &task.schema, &task.table) + .with_table_id(task.table_id); let event = match &ctx.trigger { - EventTrigger::Submitted => { - let task = &self.data.task; - let locator = TableDdlLocator::new(&task.catalog, &task.schema, &task.table) - .with_table_id(task.table_id); - TableDdlEvent::drop_table_submitted( - locator, - task.drop_if_exists, - self.data.event_context.clone(), - ) - } - _ => TableDdlEvent::lifecycle(TableDdlEventType::DropTable), + EventTrigger::Submitted => TableDdlEvent::drop_table_submitted( + locator, + task.drop_if_exists, + self.data.event_context.clone(), + ), + _ => TableDdlEvent::lifecycle(TableDdlEventType::DropTable, [locator]), }; Some(Box::new(event)) diff --git a/src/common/meta/src/ddl/drop_view.rs b/src/common/meta/src/ddl/drop_view.rs index 53a62e76159..ece1d2b06b7 100644 --- a/src/common/meta/src/ddl/drop_view.rs +++ b/src/common/meta/src/ddl/drop_view.rs @@ -231,7 +231,15 @@ impl Procedure for DropViewProcedure { self.data.event_context.clone(), ) } - _ => ViewDdlEvent::drop_lifecycle(), + _ => { + let table_ref = self.data.table_ref(); + ViewDdlEvent::drop_lifecycle( + table_ref.catalog, + table_ref.schema, + table_ref.table, + self.data.view_id(), + ) + } }; Some(Box::new(event)) diff --git a/src/common/meta/src/ddl/event/database.rs b/src/common/meta/src/ddl/event/database.rs index 9d70ea35a5e..b24611799c6 100644 --- a/src/common/meta/src/ddl/event/database.rs +++ b/src/common/meta/src/ddl/event/database.rs @@ -179,16 +179,16 @@ impl DatabaseDdlEvent { ) } - pub(crate) fn create_lifecycle() -> Self { - Self::lifecycle(CREATE_DATABASE_EVENT_TYPE) + pub(crate) fn create_lifecycle(catalog_name: &str, schema_name: &str) -> Self { + Self::lifecycle(CREATE_DATABASE_EVENT_TYPE, catalog_name, schema_name) } - pub(crate) fn alter_lifecycle() -> Self { - Self::lifecycle(ALTER_DATABASE_EVENT_TYPE) + pub(crate) fn alter_lifecycle(catalog_name: &str, schema_name: &str) -> Self { + Self::lifecycle(ALTER_DATABASE_EVENT_TYPE, catalog_name, schema_name) } - pub(crate) fn drop_lifecycle() -> Self { - Self::lifecycle(DROP_DATABASE_EVENT_TYPE) + pub(crate) fn drop_lifecycle(catalog_name: &str, schema_name: &str) -> Self { + Self::lifecycle(DROP_DATABASE_EVENT_TYPE, catalog_name, schema_name) } fn submitted( @@ -207,11 +207,11 @@ impl DatabaseDdlEvent { } } - fn lifecycle(event_type: &'static str) -> Self { + fn lifecycle(event_type: &'static str, catalog_name: &str, schema_name: &str) -> Self { Self { event_type, - catalog_name: None, - schema_name: None, + catalog_name: Some(catalog_name.to_string()), + schema_name: Some(schema_name.to_string()), payload: None, event_context: None, } diff --git a/src/common/meta/src/ddl/event/flow.rs b/src/common/meta/src/ddl/event/flow.rs index 8a52486ddfe..7d0d6cb539b 100644 --- a/src/common/meta/src/ddl/event/flow.rs +++ b/src/common/meta/src/ddl/event/flow.rs @@ -117,29 +117,36 @@ impl FlowDdlEvent { } } - /// Builds a lightweight Create Flow lifecycle event. - pub(crate) fn create_lifecycle() -> Self { - Self::lifecycle(CREATE_FLOW_EVENT_TYPE) + /// Builds a Create Flow lifecycle event with its submitted locator. + pub(crate) fn create_lifecycle(catalog_name: &str, flow_name: &str) -> Self { + Self::lifecycle(CREATE_FLOW_EVENT_TYPE, catalog_name, flow_name) } - /// Builds a successful Create Flow event containing only a resolved ID. - pub(crate) fn create_succeeded(flow_id: Option) -> Self { + /// Builds a successful Create Flow event with its submitted locator and resolved ID. + pub(crate) fn create_succeeded( + catalog_name: &str, + flow_name: &str, + flow_id: Option, + ) -> Self { Self { flow_id, - ..Self::lifecycle(CREATE_FLOW_EVENT_TYPE) + ..Self::lifecycle(CREATE_FLOW_EVENT_TYPE, catalog_name, flow_name) } } - /// Builds a lightweight Drop Flow lifecycle event. - pub(crate) fn drop_lifecycle() -> Self { - Self::lifecycle(DROP_FLOW_EVENT_TYPE) + /// Builds a Drop Flow lifecycle event with its submitted locator. + pub(crate) fn drop_lifecycle(catalog_name: &str, flow_name: &str, flow_id: u32) -> Self { + Self { + flow_id: Some(flow_id), + ..Self::lifecycle(DROP_FLOW_EVENT_TYPE, catalog_name, flow_name) + } } - fn lifecycle(event_type: &'static str) -> Self { + fn lifecycle(event_type: &'static str, catalog_name: &str, flow_name: &str) -> Self { Self { event_type, - catalog_name: None, - flow_name: None, + catalog_name: Some(catalog_name.to_string()), + flow_name: Some(flow_name.to_string()), flow_id: None, payload: None, event_context: None, diff --git a/src/common/meta/src/ddl/event/table.rs b/src/common/meta/src/ddl/event/table.rs index 8608b3151a6..9ae6f530761 100644 --- a/src/common/meta/src/ddl/event/table.rs +++ b/src/common/meta/src/ddl/event/table.rs @@ -103,6 +103,7 @@ impl TableDdlLocator { } /// Creates a locator containing only a table ID. + #[cfg(feature = "enterprise")] pub(crate) fn from_table_id(table_id: TableId) -> Self { Self { table_id: Some(table_id), @@ -361,36 +362,32 @@ impl TableDdlEvent { ) } - /// Builds a lightweight lifecycle event with null domain columns and payload. - pub(crate) fn lifecycle(event_type: TableDdlEventType) -> Self { + /// Builds a lifecycle event with stable object locators and no intent payload. + pub(crate) fn lifecycle( + event_type: TableDdlEventType, + locators: impl IntoIterator, + ) -> Self { Self { event_type, - locators: vec![TableDdlLocator::default()], + locators: locators.into_iter().collect(), payload: None, event_context: None, } } - /// Builds a Create Table success event containing only the allocated table ID. - pub(crate) fn create_table_succeeded(table_id: TableId) -> Self { - Self { - event_type: TableDdlEventType::CreateTable, - locators: vec![TableDdlLocator::from_table_id(table_id)], - payload: None, - event_context: None, - } + /// Builds a Create Table success event containing the submitted locator and allocated ID. + pub(crate) fn create_table_succeeded(locator: TableDdlLocator, table_id: TableId) -> Self { + Self::lifecycle( + TableDdlEventType::CreateTable, + [locator.with_table_id(table_id)], + ) } /// Builds Create Logical Tables success rows from their allocated locators. pub(crate) fn create_logical_tables_succeeded( locators: impl IntoIterator, ) -> Self { - Self { - event_type: TableDdlEventType::CreateLogicalTables, - locators: locators.into_iter().collect(), - payload: None, - event_context: None, - } + Self::lifecycle(TableDdlEventType::CreateLogicalTables, locators) } fn submitted( diff --git a/src/common/meta/src/ddl/event/view.rs b/src/common/meta/src/ddl/event/view.rs index d6dcf2c5f17..c1d5846fcf2 100644 --- a/src/common/meta/src/ddl/event/view.rs +++ b/src/common/meta/src/ddl/event/view.rs @@ -121,19 +121,38 @@ impl ViewDdlEvent { ) } - /// Builds a lightweight create-view lifecycle event with no locator data. - pub(crate) fn create_lifecycle() -> Self { - Self::lifecycle(CREATE_VIEW_EVENT_TYPE) + /// Builds a create-view lifecycle event with its submitted locator. + pub(crate) fn create_lifecycle(catalog_name: &str, schema_name: &str, view_name: &str) -> Self { + Self::lifecycle(CREATE_VIEW_EVENT_TYPE, catalog_name, schema_name, view_name) } - /// Builds the successful create-view row that carries only the allocated id. - pub(crate) fn create_succeeded(view_id: u32) -> Self { - Self::succeeded(CREATE_VIEW_EVENT_TYPE, view_id) + /// Builds the successful create-view row with its submitted locator and allocated ID. + pub(crate) fn create_succeeded( + catalog_name: &str, + schema_name: &str, + view_name: &str, + view_id: u32, + ) -> Self { + Self::succeeded( + CREATE_VIEW_EVENT_TYPE, + catalog_name, + schema_name, + view_name, + view_id, + ) } - /// Builds a lightweight drop-view lifecycle event with no locator data. - pub(crate) fn drop_lifecycle() -> Self { - Self::lifecycle(DROP_VIEW_EVENT_TYPE) + /// Builds a drop-view lifecycle event with its submitted locator and ID. + pub(crate) fn drop_lifecycle( + catalog_name: &str, + schema_name: &str, + view_name: &str, + view_id: u32, + ) -> Self { + Self { + view_id: Some(view_id), + ..Self::lifecycle(DROP_VIEW_EVENT_TYPE, catalog_name, schema_name, view_name) + } } fn submitted( @@ -156,24 +175,35 @@ impl ViewDdlEvent { } } - fn lifecycle(event_type: &'static str) -> Self { + fn lifecycle( + event_type: &'static str, + catalog_name: &str, + schema_name: &str, + view_name: &str, + ) -> Self { Self { event_type, - catalog_name: None, - schema_name: None, - view_name: None, + catalog_name: Some(catalog_name.to_string()), + schema_name: Some(schema_name.to_string()), + view_name: Some(view_name.to_string()), view_id: None, payload: None, event_context: None, } } - fn succeeded(event_type: &'static str, view_id: u32) -> Self { + fn succeeded( + event_type: &'static str, + catalog_name: &str, + schema_name: &str, + view_name: &str, + view_id: u32, + ) -> Self { Self { event_type, - catalog_name: None, - schema_name: None, - view_name: None, + catalog_name: Some(catalog_name.to_string()), + schema_name: Some(schema_name.to_string()), + view_name: Some(view_name.to_string()), view_id: Some(view_id), payload: None, event_context: None, diff --git a/src/common/meta/src/ddl/purge_dropped_table.rs b/src/common/meta/src/ddl/purge_dropped_table.rs index 111b79cabe8..c81aa043dc2 100644 --- a/src/common/meta/src/ddl/purge_dropped_table.rs +++ b/src/common/meta/src/ddl/purge_dropped_table.rs @@ -226,6 +226,20 @@ impl PurgeDroppedTableProcedure { fn executor(&self) -> DropTableExecutor { DropTableExecutor::new(self.data.table_name().clone(), self.data.table_id(), false) } + + fn event_locator(&self) -> TableDdlLocator { + let table_id = self.data.task.table_id; + let Some(table_name) = self.data.table_name.as_ref() else { + return TableDdlLocator::from_table_id(table_id); + }; + + TableDdlLocator::new( + &table_name.catalog_name, + &table_name.schema_name, + &table_name.table_name, + ) + .with_table_id(table_id) + } } #[async_trait] @@ -274,17 +288,12 @@ impl Procedure for PurgeDroppedTableProcedure { { return None; } - let event = match &ctx.trigger { - EventTrigger::Submitted => { - let locator = TableDdlLocator::from_table_id(self.data.task.table_id); - TableDdlEvent::purge_dropped_table_submitted( - locator, - self.data.event_context.clone(), - ) - } - _ => TableDdlEvent::lifecycle(TableDdlEventType::PurgeDroppedTable), + let locator = self.event_locator(); + let event = if ctx.trigger == EventTrigger::Submitted { + TableDdlEvent::purge_dropped_table_submitted(locator, self.data.event_context.clone()) + } else { + TableDdlEvent::lifecycle(TableDdlEventType::PurgeDroppedTable, [locator]) }; - Some(Box::new(event)) } } diff --git a/src/common/meta/src/ddl/tests/event/database.rs b/src/common/meta/src/ddl/tests/event/database.rs index 8aa85c2fbbd..29cb48f7b80 100644 --- a/src/common/meta/src/ddl/tests/event/database.rs +++ b/src/common/meta/src/ddl/tests/event/database.rs @@ -156,30 +156,33 @@ fn test_drop_database_submitted_event_contract() { } #[test] -fn test_database_lifecycle_events_have_fixed_schema_and_null_intent() { +fn test_database_lifecycle_events_preserve_locator_and_null_intent() { for (event, event_type) in [ ( - DatabaseDdlEvent::create_lifecycle(), + DatabaseDdlEvent::create_lifecycle("greptime", "metrics"), CREATE_DATABASE_EVENT_TYPE, ), ( - DatabaseDdlEvent::alter_lifecycle(), + DatabaseDdlEvent::alter_lifecycle("greptime", "metrics"), ALTER_DATABASE_EVENT_TYPE, ), - (DatabaseDdlEvent::drop_lifecycle(), DROP_DATABASE_EVENT_TYPE), + ( + DatabaseDdlEvent::drop_lifecycle("greptime", "metrics"), + DROP_DATABASE_EVENT_TYPE, + ), ] { - assert_event_locator(&event, event_type, None, None); + assert_event_locator(&event, event_type, Some("greptime"), Some("metrics")); assert_eq!(event.json_payload().unwrap(), serde_json::Value::Null); } assert_eq!( - DatabaseDdlEvent::create_lifecycle().extra_schema(), + DatabaseDdlEvent::create_lifecycle("greptime", "metrics").extra_schema(), DatabaseDdlEvent::create_submitted( "c", "s", false, &HashMap::new(), - EventContext::default() + EventContext::default(), ) .extra_schema() ); @@ -202,7 +205,7 @@ fn test_database_events_preserve_procedure_envelope_contract() { ); let lifecycle = ProcedureEvent::new( procedure_id, - Box::new(DatabaseDdlEvent::create_lifecycle()), + Box::new(DatabaseDdlEvent::create_lifecycle("greptime", "metrics")), ProcedureState::Done { output: None }, EventTrigger::Succeeded, ); @@ -220,8 +223,8 @@ fn test_database_events_preserve_procedure_envelope_contract() { CREATE_DATABASE_EVENT_TYPE, "Done", "Succeeded", - None, - None, + Some("greptime"), + Some("metrics"), ); assert_eq!(lifecycle.json_payload().unwrap(), serde_json::Value::Null); } @@ -282,6 +285,7 @@ fn assert_event_locator( catalog_name: Option<&str>, schema_name: Option<&str>, ) { + let has_event_context = !event.json_payload().unwrap().is_null(); assert_event_contract( event, event_type, @@ -298,7 +302,7 @@ fn assert_event_locator( Value { value_data: schema_name.map(|value| ValueData::StringValue(value.to_string())), }, - if catalog_name.is_some() { + if has_event_context { default_event_context_value() } else { Value { value_data: None } @@ -316,6 +320,7 @@ fn assert_procedure_event_contract( catalog_name: Option<&str>, schema_name: Option<&str>, ) { + let has_event_context = procedure_trigger == "Submitted"; assert_event_contract( event, event_type, @@ -348,7 +353,7 @@ fn assert_procedure_event_contract( Value { value_data: schema_name.map(|value| ValueData::StringValue(value.to_string())), }, - if catalog_name.is_some() { + if has_event_context { default_event_context_value() } else { Value { value_data: None } diff --git a/src/common/meta/src/ddl/tests/event/flow.rs b/src/common/meta/src/ddl/tests/event/flow.rs index b2758197917..39f24b74e51 100644 --- a/src/common/meta/src/ddl/tests/event/flow.rs +++ b/src/common/meta/src/ddl/tests/event/flow.rs @@ -99,9 +99,18 @@ fn test_flow_submitted_event_contracts() { #[test] fn test_flow_lifecycle_events_have_fixed_schema_and_null_intent() { for (event, event_type) in [ - (FlowDdlEvent::create_lifecycle(), CREATE_FLOW_EVENT_TYPE), - (FlowDdlEvent::create_succeeded(None), CREATE_FLOW_EVENT_TYPE), - (FlowDdlEvent::drop_lifecycle(), DROP_FLOW_EVENT_TYPE), + ( + FlowDdlEvent::create_lifecycle("greptime", "metrics"), + CREATE_FLOW_EVENT_TYPE, + ), + ( + FlowDdlEvent::create_succeeded("greptime", "metrics", None), + CREATE_FLOW_EVENT_TYPE, + ), + ( + FlowDdlEvent::drop_lifecycle("greptime", "metrics", 42), + DROP_FLOW_EVENT_TYPE, + ), ] { assert_event_contract( &event, @@ -109,9 +118,13 @@ fn test_flow_lifecycle_events_have_fixed_schema_and_null_intent() { &flow_schema(), &[Row { values: vec![ - Value { value_data: None }, - Value { value_data: None }, - Value { value_data: None }, + ValueData::StringValue("greptime".to_string()).into(), + ValueData::StringValue("metrics".to_string()).into(), + if event_type == DROP_FLOW_EVENT_TYPE { + ValueData::U32Value(42).into() + } else { + Value { value_data: None } + }, Value { value_data: None }, ], }], @@ -119,15 +132,15 @@ fn test_flow_lifecycle_events_have_fixed_schema_and_null_intent() { assert_eq!(event.json_payload().unwrap(), serde_json::Value::Null); } - let event = FlowDdlEvent::create_succeeded(Some(42)); + let event = FlowDdlEvent::create_succeeded("greptime", "metrics", Some(42)); assert_event_contract( &event, CREATE_FLOW_EVENT_TYPE, &flow_schema(), &[Row { values: vec![ - Value { value_data: None }, - Value { value_data: None }, + ValueData::StringValue("greptime".to_string()).into(), + ValueData::StringValue("metrics".to_string()).into(), ValueData::U32Value(42).into(), Value { value_data: None }, ], @@ -156,7 +169,11 @@ fn test_flow_events_preserve_procedure_envelope_contract() { ); let succeeded = ProcedureEvent::new( procedure_id, - Box::new(FlowDdlEvent::create_succeeded(Some(42))), + Box::new(FlowDdlEvent::create_succeeded( + "greptime", + "metrics", + Some(42), + )), ProcedureState::Done { output: None }, EventTrigger::Succeeded, ); @@ -178,8 +195,8 @@ fn test_flow_events_preserve_procedure_envelope_contract() { "Done", "Succeeded", FlowEventLocator { - catalog_name: None, - flow_name: None, + catalog_name: Some("greptime"), + flow_name: Some("metrics"), flow_id: Some(42), }, ); @@ -257,7 +274,7 @@ fn assert_procedure_event_contract( .map(ValueData::U32Value) .map(Into::into) .unwrap_or(Value { value_data: None }), - if locator.catalog_name.is_some() { + if trigger == "Submitted" { default_event_context_value() } else { Value { value_data: None } diff --git a/src/common/meta/src/ddl/tests/event/table.rs b/src/common/meta/src/ddl/tests/event/table.rs index 90e7d83b3f2..3326cb06808 100644 --- a/src/common/meta/src/ddl/tests/event/table.rs +++ b/src/common/meta/src/ddl/tests/event/table.rs @@ -108,13 +108,26 @@ fn submitted_event_contracts_are_bounded_and_fixed() { .all(|column| column.semantic_type == SemanticType::Field as i32) ); - let lifecycle = TableDdlEvent::lifecycle(case.event_type); + let lifecycle = TableDdlEvent::lifecycle( + case.event_type, + [TableDdlLocator::new( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + "lifecycle", + )], + ); assert_eq!(lifecycle.extra_schema(), schema); assert_eq!(lifecycle.json_payload().unwrap(), JsonValue::Null); - assert_eq!( - lifecycle.extra_rows().unwrap()[0].values, - vec![Value::default(); schema.len()] - ); + assert_eq!(lifecycle.extra_rows().unwrap()[0].values, { + let mut values = table_locator_values(Some("lifecycle"), None); + if matches!( + case.event_type, + TableDdlEventType::CreateLogicalTables | TableDdlEventType::AlterLogicalTables + ) { + values.push(Value::default()); + } + values + }); } } @@ -158,6 +171,16 @@ fn later_lifecycle_events_are_uniform() { for case in procedure_cases() { let submitted = event_for(case.procedure.as_ref(), EventTrigger::Submitted); let schema = submitted.extra_schema(); + let expected_rows = submitted + .extra_rows() + .unwrap() + .into_iter() + .map(|mut row| { + // Lifecycle events retain locators but not submitted event context. + *row.values.last_mut().unwrap() = Value::default(); + row.values + }) + .collect::>(); for trigger in &triggers { let event = event_for(case.procedure.as_ref(), trigger.clone()); @@ -166,8 +189,13 @@ fn later_lifecycle_events_are_uniform() { assert_eq!(event.extra_schema(), schema); assert_eq!(event.json_payload().unwrap(), JsonValue::Null); assert_eq!( - event.extra_rows().unwrap()[0].values, - vec![Value::default(); schema.len()] + event + .extra_rows() + .unwrap() + .into_iter() + .map(|row| row.values) + .collect::>(), + expected_rows ); } } @@ -193,7 +221,7 @@ fn create_success_events_keep_allocated_ids() { assert_eq!(event.json_payload().unwrap(), JsonValue::Null); assert_eq!( event.extra_rows().unwrap()[0].values, - table_locator_values(None, Some(42)) + table_locator_values(Some("create_success"), Some(42)) ); let logical_tables = CreateLogicalTablesProcedure::new( diff --git a/src/common/meta/src/ddl/tests/event/view.rs b/src/common/meta/src/ddl/tests/event/view.rs index 10947e45ee8..cabe7418ffb 100644 --- a/src/common/meta/src/ddl/tests/event/view.rs +++ b/src/common/meta/src/ddl/tests/event/view.rs @@ -102,26 +102,44 @@ fn test_view_submitted_event_contracts() { #[test] fn test_view_lifecycle_event_contracts() { for (event, event_type) in [ - (ViewDdlEvent::create_lifecycle(), CREATE_VIEW_EVENT_TYPE), - (ViewDdlEvent::drop_lifecycle(), DROP_VIEW_EVENT_TYPE), + ( + ViewDdlEvent::create_lifecycle("greptime", "public", "v_metrics"), + CREATE_VIEW_EVENT_TYPE, + ), + ( + ViewDdlEvent::drop_lifecycle("greptime", "public", "v_metrics", 42), + DROP_VIEW_EVENT_TYPE, + ), ] { - assert_lightweight_event(&event, event_type); + assert_eq!(event.json_payload().unwrap(), serde_json::Value::Null); + assert_view_event_contract( + &event, + event_type, + ViewEventLocator { + catalog_name: Some("greptime"), + schema_name: Some("public"), + view_name: Some("v_metrics"), + view_id: (event_type == DROP_VIEW_EVENT_TYPE).then_some(42), + }, + ); } - let event = ViewDdlEvent::create_succeeded(84); + let event = ViewDdlEvent::create_succeeded("greptime", "public", "v_metrics", 84); assert_view_event_contract( &event, CREATE_VIEW_EVENT_TYPE, ViewEventLocator { + catalog_name: Some("greptime"), + schema_name: Some("public"), + view_name: Some("v_metrics"), view_id: Some(84), - ..Default::default() }, ); assert_eq!(event.json_payload().unwrap(), serde_json::Value::Null); } #[test] -fn test_view_procedures_emit_lightweight_lifecycle_events() { +fn test_view_procedures_preserve_lifecycle_locators() { let create = CreateViewProcedure::new( test_create_view_task("view_name"), EventContext::default(), @@ -147,18 +165,36 @@ fn test_view_procedures_emit_lightweight_lifecycle_events() { EventTrigger::Poisoned, ]; - for (procedure, event_type) in [ - (&create as &dyn Procedure, CREATE_VIEW_EVENT_TYPE), - (&drop as &dyn Procedure, DROP_VIEW_EVENT_TYPE), + for (procedure, event_type, view_id) in [ + (&create as &dyn Procedure, CREATE_VIEW_EVENT_TYPE, None), + (&drop as &dyn Procedure, DROP_VIEW_EVENT_TYPE, Some(42)), ] { for trigger in &triggers { let event = event_for(procedure, trigger.clone()); - assert_lightweight_event(event.as_ref(), event_type); + assert_view_event_contract( + event.as_ref(), + event_type, + ViewEventLocator { + catalog_name: Some("greptime"), + schema_name: Some("public"), + view_name: Some("view_name"), + view_id, + }, + ); } } let event = event_for(&drop, EventTrigger::Succeeded); - assert_lightweight_event(event.as_ref(), DROP_VIEW_EVENT_TYPE); + assert_view_event_contract( + event.as_ref(), + DROP_VIEW_EVENT_TYPE, + ViewEventLocator { + catalog_name: Some("greptime"), + schema_name: Some("public"), + view_name: Some("view_name"), + view_id: Some(42), + }, + ); } #[test] @@ -177,8 +213,10 @@ fn test_create_view_succeeded_output_mapping() { event.as_ref(), CREATE_VIEW_EVENT_TYPE, ViewEventLocator { + catalog_name: Some("greptime"), + schema_name: Some("public"), + view_name: Some("view_name"), view_id: Some(84), - ..Default::default() }, ); assert_eq!(event.json_payload().unwrap(), serde_json::Value::Null); @@ -187,7 +225,16 @@ fn test_create_view_succeeded_output_mapping() { for output in invalid_outputs { let state = ProcedureState::Done { output }; let event = event_for_state(&procedure, EventTrigger::Succeeded, &state); - assert_lightweight_event(event.as_ref(), CREATE_VIEW_EVENT_TYPE); + assert_view_event_contract( + event.as_ref(), + CREATE_VIEW_EVENT_TYPE, + ViewEventLocator { + catalog_name: Some("greptime"), + schema_name: Some("public"), + view_name: Some("view_name"), + view_id: None, + }, + ); } } @@ -233,7 +280,12 @@ fn test_view_event_procedure_envelope_contract() { ); let succeeded = ProcedureEvent::new( procedure_id, - Box::new(ViewDdlEvent::create_succeeded(42)), + Box::new(ViewDdlEvent::create_succeeded( + "greptime", + "public", + "view_name", + 42, + )), ProcedureState::Done { output: None }, EventTrigger::Succeeded, ); @@ -256,8 +308,10 @@ fn test_view_event_procedure_envelope_contract() { "Done", "Succeeded", ViewEventLocator { + catalog_name: Some("greptime"), + schema_name: Some("public"), + view_name: Some("view_name"), view_id: Some(42), - ..Default::default() }, ); } @@ -280,11 +334,6 @@ impl ViewEventLocator<'_> { .map(ValueData::U32Value) .map(Into::into) .unwrap_or_default(), - if self.catalog_name.is_some() { - default_event_context_value() - } else { - Value { value_data: None } - }, ] } } @@ -300,19 +349,13 @@ fn view_schema() -> Vec { } fn assert_view_event_contract(event: &dyn Event, event_type: &str, locator: ViewEventLocator<'_>) { - assert_event_contract( - event, - event_type, - &view_schema(), - &[Row { - values: locator.values(), - }], - ); -} - -fn assert_lightweight_event(event: &dyn Event, event_type: &str) { - assert_view_event_contract(event, event_type, ViewEventLocator::default()); - assert_eq!(event.json_payload().unwrap(), serde_json::Value::Null); + let mut values = locator.values(); + values.push(if event.json_payload().unwrap().is_null() { + Value { value_data: None } + } else { + default_event_context_value() + }); + assert_event_contract(event, event_type, &view_schema(), &[Row { values }]); } fn assert_procedure_event_contract( @@ -337,6 +380,11 @@ fn assert_procedure_event_contract( procedure_trigger_value(trigger), ]; values.extend(locator.values()); + values.push(if trigger == "Submitted" { + default_event_context_value() + } else { + Value { value_data: None } + }); assert_event_contract(event, event_type, &schema, &[Row { values }]); } diff --git a/src/common/meta/src/ddl/truncate_table.rs b/src/common/meta/src/ddl/truncate_table.rs index 9d0ed09e8c8..6a561b62afb 100644 --- a/src/common/meta/src/ddl/truncate_table.rs +++ b/src/common/meta/src/ddl/truncate_table.rs @@ -99,18 +99,16 @@ impl Procedure for TruncateTableProcedure { { return None; } + let task = &self.data.task; + let locator = TableDdlLocator::new(&task.catalog, &task.schema, &task.table) + .with_table_id(task.table_id); let event = match &ctx.trigger { - EventTrigger::Submitted => { - let task = &self.data.task; - let locator = TableDdlLocator::new(&task.catalog, &task.schema, &task.table) - .with_table_id(task.table_id); - TableDdlEvent::truncate_table_submitted( - locator, - task.time_ranges.len(), - self.data.event_context.clone(), - ) - } - _ => TableDdlEvent::lifecycle(TableDdlEventType::TruncateTable), + EventTrigger::Submitted => TableDdlEvent::truncate_table_submitted( + locator, + task.time_ranges.len(), + self.data.event_context.clone(), + ), + _ => TableDdlEvent::lifecycle(TableDdlEventType::TruncateTable, [locator]), }; Some(Box::new(event)) diff --git a/src/common/meta/src/ddl/undrop_table.rs b/src/common/meta/src/ddl/undrop_table.rs index 635d15ee6a5..3295b1163e2 100644 --- a/src/common/meta/src/ddl/undrop_table.rs +++ b/src/common/meta/src/ddl/undrop_table.rs @@ -380,6 +380,23 @@ impl UndropTableProcedure { } } +impl UndropTableProcedure { + fn event_locator(&self) -> TableDdlLocator { + self.data + .table_name + .as_ref() + .map(|table_name| { + TableDdlLocator::new( + &table_name.catalog_name, + &table_name.schema_name, + &table_name.table_name, + ) + }) + .unwrap_or_default() + .with_table_id(self.data.task.table_id) + } +} + #[async_trait] impl Procedure for UndropTableProcedure { fn type_name(&self) -> &str { @@ -437,24 +454,12 @@ impl Procedure for UndropTableProcedure { { return None; } + let locator = self.event_locator(); let event = match &ctx.trigger { EventTrigger::Submitted => { - let locator = self - .data - .table_name - .as_ref() - .map(|table_name| { - TableDdlLocator::new( - &table_name.catalog_name, - &table_name.schema_name, - &table_name.table_name, - ) - }) - .unwrap_or_default() - .with_table_id(self.data.task.table_id); TableDdlEvent::undrop_table_submitted(locator, self.data.event_context.clone()) } - _ => TableDdlEvent::lifecycle(TableDdlEventType::UndropTable), + _ => TableDdlEvent::lifecycle(TableDdlEventType::UndropTable, [locator]), }; Some(Box::new(event)) diff --git a/src/meta-srv/src/event/repartition.rs b/src/meta-srv/src/event/repartition.rs index ad47fbf9d42..958bc4e11d8 100644 --- a/src/meta-srv/src/event/repartition.rs +++ b/src/meta-srv/src/event/repartition.rs @@ -95,12 +95,12 @@ impl RepartitionEvent { } } - pub(crate) fn lifecycle() -> Self { + pub(crate) fn lifecycle(persistent_ctx: &RepartitionPersistentContext) -> Self { Self { - catalog_name: None, - schema_name: None, - table_name: None, - table_id: None, + catalog_name: Some(persistent_ctx.catalog_name.clone()), + schema_name: Some(persistent_ctx.schema_name.clone()), + table_name: Some(persistent_ctx.table_name.clone()), + table_id: Some(persistent_ctx.table_id), payload: None, event_context: None, } @@ -279,14 +279,14 @@ impl RepartitionGroupEvent { } } - pub(crate) fn lifecycle() -> Self { + pub(crate) fn lifecycle(persistent_ctx: &GroupPersistentContext) -> Self { Self { - catalog_name: None, - schema_name: None, - table_name: None, - table_id: None, - parent_procedure_id: None, - group_id: None, + catalog_name: Some(persistent_ctx.catalog_name.clone()), + schema_name: Some(persistent_ctx.schema_name.clone()), + table_name: persistent_ctx.table_name.clone(), + table_id: Some(persistent_ctx.table_id), + parent_procedure_id: persistent_ctx.parent_procedure_id.map(|id| id.to_string()), + group_id: Some(persistent_ctx.group_id.to_string()), topology: None, payload: None, } @@ -377,7 +377,7 @@ mod tests { use std::time::Duration; use api::v1::value::ValueData; - use api::v1::{ColumnSchema, Row, Value}; + use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::event_table::{ PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, @@ -641,22 +641,32 @@ mod tests { } #[test] - fn test_lifecycle_events_are_lightweight() { - let parent = RepartitionEvent::lifecycle(); + fn test_lifecycle_events_preserve_locators_and_null_payloads() { + let parent_ctx = parent_persistent_ctx(); + let parent = RepartitionEvent::lifecycle(&parent_ctx); assert_event_contract( &parent, REPARTITION_EVENT_TYPE, &parent_schema(), - &[null_row(parent_schema().len())], + &[Row { + values: vec![ + ValueData::StringValue("greptime".to_string()).into(), + ValueData::StringValue("public".to_string()).into(), + ValueData::StringValue("repartition_events".to_string()).into(), + ValueData::U32Value(1024).into(), + Default::default(), + ], + }], ); assert_eq!(parent.json_payload().unwrap(), serde_json::Value::Null); - let group = RepartitionGroupEvent::lifecycle(); + let group_ctx = new_persistent_context(1024, vec![], vec![]); + let group = RepartitionGroupEvent::lifecycle(&group_ctx); assert_event_contract( &group, REPARTITION_GROUP_EVENT_TYPE, &group_schema(), - &[null_row(group_schema().len())], + &[group.extra_row(None)], ); assert_eq!(group.json_payload().unwrap(), serde_json::Value::Null); } @@ -762,12 +772,4 @@ mod tests { ], } } - - fn null_row(column_count: usize) -> Row { - Row { - values: (0..column_count) - .map(|_| Value { value_data: None }) - .collect(), - } - } } diff --git a/src/meta-srv/src/procedure/repartition.rs b/src/meta-srv/src/procedure/repartition.rs index 62f9f26aabd..d8b585b28ee 100644 --- a/src/meta-srv/src/procedure/repartition.rs +++ b/src/meta-srv/src/procedure/repartition.rs @@ -809,7 +809,7 @@ impl Procedure for RepartitionProcedure { let start = self.state.as_any().downcast_ref::()?; RepartitionEvent::submitted(&self.context.persistent_ctx, start) } else { - RepartitionEvent::lifecycle() + RepartitionEvent::lifecycle(&self.context.persistent_ctx) }; Some(Box::new(event)) } diff --git a/src/meta-srv/src/procedure/repartition/group.rs b/src/meta-srv/src/procedure/repartition/group.rs index 700c2fdcd1d..a3c38b59557 100644 --- a/src/meta-srv/src/procedure/repartition/group.rs +++ b/src/meta-srv/src/procedure/repartition/group.rs @@ -273,7 +273,7 @@ impl Procedure for RepartitionGroupProcedure { let event = if matches!(ctx.trigger, EventTrigger::Submitted) { RepartitionGroupEvent::submitted(&self.context.persistent_ctx) } else { - RepartitionGroupEvent::lifecycle() + RepartitionGroupEvent::lifecycle(&self.context.persistent_ctx) }; Some(Box::new(event)) } diff --git a/tests-integration/tests/database_ddl_event.rs b/tests-integration/tests/database_ddl_event.rs index 79b741b977f..6cbc65a2673 100644 --- a/tests-integration/tests/database_ddl_event.rs +++ b/tests-integration/tests/database_ddl_event.rs @@ -138,8 +138,8 @@ FROM greptime_private.events WHERE type = '{event_type}' AND procedure_state = 'Done' AND json_path_match(procedure_trigger, '$.type == "Succeeded"') - AND catalog_name IS NULL - AND schema_name IS NULL + AND catalog_name = 'greptime' + AND schema_name = '{DATABASE_NAME}' AND json_is_null(payload)"# ); assert_single_event(instance, &lifecycle).await; diff --git a/tests-integration/tests/flow_ddl_event.rs b/tests-integration/tests/flow_ddl_event.rs index 004d8382ed2..c5fc9be94c9 100644 --- a/tests-integration/tests/flow_ddl_event.rs +++ b/tests-integration/tests/flow_ddl_event.rs @@ -98,8 +98,8 @@ WHERE type = '{CREATE_FLOW_EVENT_TYPE}' AND procedure_id = '{procedure_id}' AND procedure_state = 'Done' AND json_path_match(procedure_trigger, '$.type == "Succeeded"') - AND catalog_name IS NULL - AND flow_name IS NULL + AND catalog_name = 'greptime' + AND flow_name = '{flow}' AND flow_id IS NOT NULL AND json_is_null(payload)"#, ), @@ -152,9 +152,9 @@ WHERE type = '{DROP_FLOW_EVENT_TYPE}' AND procedure_id = '{procedure_id}' AND procedure_state = 'Done' AND json_path_match(procedure_trigger, '$.type == "Succeeded"') - AND catalog_name IS NULL - AND flow_name IS NULL - AND flow_id IS NULL + AND catalog_name = 'greptime' + AND flow_name = '{flow}' + AND flow_id IS NOT NULL AND json_is_null(payload)"#, ), ) diff --git a/tests-integration/tests/repartition_event.rs b/tests-integration/tests/repartition_event.rs index 3871406cfa2..65503a4a8ee 100644 --- a/tests-integration/tests/repartition_event.rs +++ b/tests-integration/tests/repartition_event.rs @@ -129,9 +129,32 @@ WHERE type = 'repartition' +-------------+--------------+-------------+--------------------+----------+-----------------+----------------------+---------------+"; assert_eventually_eq(instance, &query, expected).await; + assert_repartition_lifecycle_event(instance, &procedure_id).await; procedure_id } +async fn assert_repartition_lifecycle_event( + instance: &Arc, + procedure_id: &str, +) { + let query = format!( + r#"SELECT type, catalog_name, schema_name, table_name, table_id +FROM greptime_private.events +WHERE type = 'repartition' + AND procedure_id = '{procedure_id}' + AND json_path_match(procedure_trigger, '$.type == "Succeeded"') + AND json_is_null(payload)"# + ); + let expected = "\ ++-------------+--------------+-------------+--------------------+----------+ +| type | catalog_name | schema_name | table_name | table_id | ++-------------+--------------+-------------+--------------------+----------+ +| repartition | greptime | public | repartition_events | 1024 | ++-------------+--------------+-------------+--------------------+----------+"; + + assert_eventually_eq(instance, &query, expected).await; +} + async fn assert_repartition_group_event( instance: &Arc, parent_procedure_id: &str, @@ -170,6 +193,32 @@ ORDER BY target_region_id"# | repartition_group | greptime | public | repartition_events | 1024 | true | true | 4398046511104 | 0 | true | 4398046511105 | 1 | host >= 10 | +-------------------+--------------+-------------+--------------------+----------+-----------------------------+--------------+------------------+----------------------+----------------+------------------+----------------------+-----------------------+"; + assert_eventually_eq(instance, &query, expected).await; + assert_repartition_group_lifecycle_event(instance, &procedure_id, parent_procedure_id).await; +} + +async fn assert_repartition_group_lifecycle_event( + instance: &Arc, + procedure_id: &str, + parent_procedure_id: &str, +) { + let query = format!( + r#"SELECT type, catalog_name, schema_name, table_name, table_id, + parent_procedure_id = '{parent_procedure_id}' AS matches_parent_procedure_id, + repartition_group_id IS NOT NULL AS has_group_id +FROM greptime_private.events +WHERE type = 'repartition_group' + AND procedure_id = '{procedure_id}' + AND json_path_match(procedure_trigger, '$.type == "Succeeded"') + AND json_is_null(payload)"# + ); + let expected = "\ ++-------------------+--------------+-------------+--------------------+----------+-----------------------------+--------------+ +| type | catalog_name | schema_name | table_name | table_id | matches_parent_procedure_id | has_group_id | ++-------------------+--------------+-------------+--------------------+----------+-----------------------------+--------------+ +| repartition_group | greptime | public | repartition_events | 1024 | true | true | ++-------------------+--------------+-------------+--------------------+----------+-----------------------------+--------------+"; + assert_eventually_eq(instance, &query, expected).await; } diff --git a/tests-integration/tests/table_ddl_event.rs b/tests-integration/tests/table_ddl_event.rs index c53842d66d6..5c23e101de2 100644 --- a/tests-integration/tests/table_ddl_event.rs +++ b/tests-integration/tests/table_ddl_event.rs @@ -230,23 +230,17 @@ async fn test_table_ddl_procedure_events() { +--------------+------------------+", ) .await; - assert_id_terminal_event( + assert_terminal_event( &frontend, "create_table", &create_table_procedure_id, - table_id, "Succeeded", - "\ -+--------------+ -| type | -+--------------+ -| create_table | -+--------------+", + Some(TABLE), + Some(table_id), ) .await; - // Act / Assert: Alter and Truncate have a rich submitted row and lightweight - // terminal lifecycle row. + // Act / Assert: Alter and Truncate retain their submitted locators on terminal rows. run_sql_with_context( &frontend, &format!("ALTER TABLE {TABLE} ADD COLUMN extra STRING"), @@ -278,11 +272,13 @@ async fn test_table_ddl_procedure_events() { +-------------+------------------+", ) .await; - assert_lightweight_terminal_event( + assert_terminal_event( &frontend, "alter_table", &alter_table_procedure_id, "Succeeded", + Some(TABLE), + Some(table_id), ) .await; run_sql( @@ -322,11 +318,13 @@ async fn test_table_ddl_procedure_events() { +----------------+------------------+", ) .await; - assert_lightweight_terminal_event( + assert_terminal_event( &frontend, "truncate_table", &truncate_table_procedure_id, "Succeeded", + Some(TABLE), + Some(table_id), ) .await; @@ -363,11 +361,13 @@ async fn test_table_ddl_procedure_events() { +------------+------------------+", ) .await; - assert_lightweight_terminal_event( + assert_terminal_event( &frontend, "drop_table", &drop_table_procedure_id, "Succeeded", + Some(TABLE), + Some(table_id), ) .await; #[cfg(feature = "enterprise")] @@ -404,11 +404,13 @@ async fn test_table_ddl_procedure_events() { +--------------+", ) .await; - assert_lightweight_terminal_event( + assert_terminal_event( &frontend, "undrop_table", &undrop_table_procedure_id, "Succeeded", + Some(TABLE), + Some(table_id), ) .await; run_sql(&frontend, &format!("DROP TABLE {TABLE}")).await; @@ -444,11 +446,13 @@ async fn test_table_ddl_procedure_events() { +---------------------+", ) .await; - assert_lightweight_terminal_event( + assert_terminal_event( &frontend, "purge_dropped_table", &purge_table_procedure_id, "Succeeded", + Some(TABLE), + Some(table_id), ) .await; } @@ -518,7 +522,7 @@ async fn test_table_ddl_procedure_events() { .await; // Act / Assert: logical Alter preserves the one-row-per-logical-table submitted - // contract and emits a lightweight terminal row. + // contract and retains the logical-table locator on its terminal row. run_sql_with_context( &frontend, &format!("ALTER TABLE {LOGICAL_TABLE} ADD COLUMN rack STRING PRIMARY KEY"), @@ -551,11 +555,13 @@ async fn test_table_ddl_procedure_events() { +----------------------+--------------------------+", ) .await; - assert_lightweight_terminal_event( + assert_terminal_event( &frontend, "alter_logical_tables", &alter_logical_tables_procedure_id, "Succeeded", + Some(LOGICAL_TABLE), + None, ) .await; } @@ -719,27 +725,46 @@ async fn assert_id_submitted_event( assert_eventually_eq(instance, &query, expected).await; } -async fn assert_id_terminal_event( +async fn assert_terminal_event( instance: &Arc, event_type: &str, procedure_id: &str, - table_id: u32, terminal_trigger: &str, - expected: &str, + table_name: Option<&str>, + table_id: Option, ) { + let table_name_predicate = table_name.map_or_else( + || format!("{} IS NULL", TABLE_NAME_COLUMN.name()), + |table_name| format!("{} = '{table_name}'", TABLE_NAME_COLUMN.name()), + ); + let catalog_and_schema_predicate = table_name.map_or_else( + || { + format!( + "{} IS NULL AND {} IS NULL", + CATALOG_NAME_COLUMN.name(), + SCHEMA_NAME_COLUMN.name(), + ) + }, + |_| { + format!( + "{} = 'greptime' AND {} = 'public'", + CATALOG_NAME_COLUMN.name(), + SCHEMA_NAME_COLUMN.name(), + ) + }, + ); + let table_id_predicate = table_id.map_or_else( + || format!("{} IS NULL", TABLE_ID_COLUMN.name()), + |table_id| format!("{} = {table_id}", TABLE_ID_COLUMN.name()), + ); let query = format!( - "SELECT {} FROM {EVENTS_TABLE} WHERE {} = '{event_type}' AND {} = '{procedure_id}' AND json_path_match({}, '$.type == \"{terminal_trigger}\"') AND {} = {table_id} AND json_is_null({}) AND {} IS NULL AND {} IS NULL AND {} IS NULL", - TYPE_COLUMN.name(), + "SELECT count(*) AS event_count FROM {EVENTS_TABLE} WHERE {} = '{event_type}' AND {} = '{procedure_id}' AND json_path_match({}, '$.type == \"{terminal_trigger}\"') AND json_is_null({}) AND {catalog_and_schema_predicate} AND {table_name_predicate} AND {table_id_predicate}", TYPE_COLUMN.name(), PROCEDURE_ID_COLUMN.name(), PROCEDURE_TRIGGER_COLUMN.name(), - TABLE_ID_COLUMN.name(), PAYLOAD_COLUMN.name(), - CATALOG_NAME_COLUMN.name(), - SCHEMA_NAME_COLUMN.name(), - TABLE_NAME_COLUMN.name(), ); - assert_eventually_eq(instance, &query, expected).await; + assert_single_event(instance, &query).await; } async fn assert_logical_terminal_event( @@ -764,23 +789,3 @@ async fn assert_logical_terminal_event( ); assert_eventually_eq(instance, &query, expected).await; } - -async fn assert_lightweight_terminal_event( - instance: &Arc, - event_type: &str, - procedure_id: &str, - terminal_trigger: &str, -) { - let query = format!( - "SELECT count(*) AS event_count FROM {EVENTS_TABLE} WHERE {} = '{event_type}' AND {} = '{procedure_id}' AND json_path_match({}, '$.type == \"{terminal_trigger}\"') AND json_is_null({}) AND {} IS NULL AND {} IS NULL AND {} IS NULL AND {} IS NULL", - TYPE_COLUMN.name(), - PROCEDURE_ID_COLUMN.name(), - PROCEDURE_TRIGGER_COLUMN.name(), - PAYLOAD_COLUMN.name(), - CATALOG_NAME_COLUMN.name(), - SCHEMA_NAME_COLUMN.name(), - TABLE_NAME_COLUMN.name(), - TABLE_ID_COLUMN.name(), - ); - assert_single_event(instance, &query).await; -} diff --git a/tests-integration/tests/view_ddl_event.rs b/tests-integration/tests/view_ddl_event.rs index a25b58b684a..fce33d09bc8 100644 --- a/tests-integration/tests/view_ddl_event.rs +++ b/tests-integration/tests/view_ddl_event.rs @@ -134,9 +134,9 @@ WHERE type = '{CREATE_VIEW_EVENT_TYPE}' AND procedure_id = '{procedure_id}' AND procedure_state = 'Done' AND json_path_match(procedure_trigger, '$.type == "Succeeded"') - AND catalog_name IS NULL - AND schema_name IS NULL - AND view_name IS NULL + AND catalog_name = 'greptime' + AND schema_name = 'public' + AND view_name = '{view}' AND view_id IS NOT NULL AND json_is_null(payload)"#, ), @@ -190,10 +190,10 @@ WHERE type = '{DROP_VIEW_EVENT_TYPE}' AND procedure_id = '{procedure_id}' AND procedure_state = 'Done' AND json_path_match(procedure_trigger, '$.type == "Succeeded"') - AND catalog_name IS NULL - AND schema_name IS NULL - AND view_name IS NULL - AND view_id IS NULL + AND catalog_name = 'greptime' + AND schema_name = 'public' + AND view_name = '{view}' + AND view_id IS NOT NULL AND json_is_null(payload)"#, ), ) From c851430349faf7be24ed1ed530b4f361d1588b18 Mon Sep 17 00:00:00 2001 From: discord9 Date: Tue, 11 Aug 2026 15:16:06 +0800 Subject: [PATCH 21/92] chore(deps): bump datafusion to 452cb4b (support Dictionary literals in substrait) (#8839) Bump the GreptimeTeam/datafusion fork rev from 6d6ae9a to 452cb4b, which includes fix(substrait): support Dictionary literals in producer. This fixes flow queries against dictionary-encoded PK string columns (metric tables) failing with: Failed to encode DataFusion plan: NotImplemented("Unsupported literal: Dictionary(UInt32, Utf8(...))") The substrait producer now encodes ScalarValue::Dictionary as its inner value wrapped in a cast to the dictionary type, so the original SQL works without CAST workarounds. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 116 ++++++++++++++++++++++++++--------------------------- Cargo.toml | 30 +++++++------- 2 files changed, 73 insertions(+), 73 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d6385dd7657..2e6fe8ae92c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2226,7 +2226,7 @@ checksum = "fe6d2e5af09e8c8ad56c969f2157a3d4238cebc7c55f0a517728c38f7b200f81" dependencies = [ "serde", "termcolor", - "unicode-width 0.2.1", + "unicode-width 0.1.14", ] [[package]] @@ -2237,7 +2237,7 @@ checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" dependencies = [ "serde", "termcolor", - "unicode-width 0.2.1", + "unicode-width 0.1.14", ] [[package]] @@ -2253,7 +2253,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" dependencies = [ "lazy_static", - "windows-sys 0.59.0", + "windows-sys 0.48.0", ] [[package]] @@ -3726,7 +3726,7 @@ checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" [[package]] name = "datafusion" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "arrow-schema 58.3.0", @@ -3780,7 +3780,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3804,7 +3804,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3826,7 +3826,7 @@ dependencies = [ [[package]] name = "datafusion-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -3850,7 +3850,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "futures", "log", @@ -3860,7 +3860,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "async-compression", @@ -3894,7 +3894,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "arrow-ipc 58.3.0", @@ -3917,7 +3917,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3939,7 +3939,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3962,7 +3962,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3991,12 +3991,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" [[package]] name = "datafusion-execution" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "arrow-buffer 58.3.0", @@ -4018,7 +4018,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "async-trait", @@ -4040,7 +4040,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4052,7 +4052,7 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "arrow-buffer 58.3.0", @@ -4083,7 +4083,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4104,7 +4104,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4116,7 +4116,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "arrow-ord 58.3.0", @@ -4140,7 +4140,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "async-trait", @@ -4155,7 +4155,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4172,7 +4172,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -4181,7 +4181,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "datafusion-doc", "quote", @@ -4191,7 +4191,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "chrono", @@ -4240,7 +4240,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4263,7 +4263,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4277,7 +4277,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4293,7 +4293,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4311,7 +4311,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4342,7 +4342,7 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "chrono", @@ -4369,7 +4369,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4379,7 +4379,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4395,7 +4395,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "async-trait", "datafusion-common", @@ -4408,7 +4408,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "arrow 58.3.0", "bigdecimal 0.4.8", @@ -4426,7 +4426,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=6d6ae9a8073900bdf89d098d491d4dc0269699b0#6d6ae9a8073900bdf89d098d491d4dc0269699b0" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" dependencies = [ "async-recursion", "async-trait", @@ -5119,7 +5119,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.52.0", ] [[package]] @@ -5681,7 +5681,7 @@ checksum = "94e7099f6313ecacbe1256e8ff9d617b75d1bcb16a6fddef94866d225a01a14a" dependencies = [ "io-lifetimes 2.0.4", "rustix 1.0.7", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -6578,7 +6578,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.6.4", + "socket2 0.5.10", "tokio", "tower-service", "tracing", @@ -6649,7 +6649,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core 0.61.2", + "windows-core 0.57.0", ] [[package]] @@ -7077,7 +7077,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20fd6de4ccfcc187e38bc21cfa543cb5a302cb86a8b114eb7f0bf0dc9f8ac00f" dependencies = [ "io-lifetimes 3.0.1", - "windows-sys 0.60.2", + "windows-sys 0.52.0", ] [[package]] @@ -7115,7 +7115,7 @@ checksum = "e04d7f318608d35d4b61ddd75cbdaee86b023ebe2bd5a66ee0915f0bf93095a9" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -7818,7 +7818,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667" dependencies = [ "cfg-if", - "windows-targets 0.53.5", + "windows-targets 0.48.5", ] [[package]] @@ -11212,7 +11212,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4" dependencies = [ "bytes", - "heck 0.5.0", + "heck 0.4.1", "itertools 0.12.1", "log", "multimap", @@ -11232,8 +11232,8 @@ version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ - "heck 0.5.0", - "itertools 0.14.0", + "heck 0.4.1", + "itertools 0.10.5", "log", "multimap", "once_cell", @@ -11281,7 +11281,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.10.5", "proc-macro2", "quote", "syn 2.0.117", @@ -11294,7 +11294,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9120690fafc389a67ba3803df527d0ec9cbbc9cc45e4cc20b332996dfb672425" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.10.5", "proc-macro2", "quote", "syn 2.0.117", @@ -11760,7 +11760,7 @@ dependencies = [ "once_cell", "socket2 0.5.10", "tracing", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -12709,7 +12709,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.4.15", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -12722,7 +12722,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.9.4", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -12813,7 +12813,7 @@ dependencies = [ "security-framework 3.7.0", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -13697,7 +13697,7 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1961e2ef424c1424204d3a5d6975f934f56b6d50ff5732382d84ebf460e147f7" dependencies = [ - "heck 0.5.0", + "heck 0.4.1", "proc-macro2", "quote", "syn 2.0.117", @@ -14112,7 +14112,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -14722,7 +14722,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix 1.0.7", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -16470,7 +16470,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.48.0", ] [[package]] @@ -16914,7 +16914,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f3fd376f71958b862e7afb20cfe5a22830e1963462f3a17f49d82a6c1d1f42d" dependencies = [ "bitflags 2.12.1", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 64e6dba2f7f..a138a880ec6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -345,21 +345,21 @@ git = "https://github.com/GreptimeTeam/greptime-meter.git" rev = "5618e779cf2bb4755b499c630fba4c35e91898cb" [patch.crates-io] -datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } -datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "6d6ae9a8073900bdf89d098d491d4dc0269699b0" } +datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "2aefa08a8d69c96eec2d6d6703598a009bba6e4c" } # on branch v0.61.x [profile.release] From f27f27da661c69033de7b2e4a5b70b4f058f8dd2 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 11 Aug 2026 15:59:12 +0800 Subject: [PATCH 22/92] fix(ci): grant pull-requests write and stop counting drafts (#8844) Posting to `/issues/{n}/comments` is authorized against the target object, and that object is a pull request, so `issues: write` alone is refused with 403 and the warning comment never lands. Drafts are no longer counted and no longer warned about. `ready_for_review` is added to the trigger types so that opening as a draft and flipping it to ready still goes through the check. Signed-off-by: Dennis Zhuang --- .github/scripts/pr-open-limit.js | 12 +++++++++--- .github/workflows/pr-open-limit.yml | 5 +++-- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/.github/scripts/pr-open-limit.js b/.github/scripts/pr-open-limit.js index dd1650b0d64..46f770074f3 100644 --- a/.github/scripts/pr-open-limit.js +++ b/.github/scripts/pr-open-limit.js @@ -1,7 +1,8 @@ // Warns when a repository member has too many open pull requests. // // Review capacity is the bottleneck, so authors are expected to land or close -// existing work before opening more. Drafts count: they still occupy attention. +// existing work before opening more. Drafts are not counted and never warned +// about: they are not asking for review yet. // This is advisory only — it never closes a pull request or fails the job. import { appendFileSync } from "node:fs"; @@ -47,7 +48,7 @@ async function isTeamMember(octokit, owner, repo, username) { function buildComment(author, total, limit, openPrs) { const listed = openPrs .slice(0, MAX_LISTED_PRS) - .map((pr) => `- #${pr.number} ${pr.title}${pr.draft ? " _(draft)_" : ""}`); + .map((pr) => `- #${pr.number} ${pr.title}`); const hidden = openPrs.length - listed.length; const list = hidden > 0 ? `${listed.join("\n")}\n- ...and ${hidden} more` : listed.join("\n"); @@ -99,6 +100,11 @@ async function upsertComment(octokit, params, body) { return; } + if (process.env.PR_IS_DRAFT === "true") { + summary(`Skipping #${prNumber}: draft.`); + return; + } + const octokit = new Octokit({ auth: process.env.GITHUB_TOKEN }); if (!(await isTeamMember(octokit, owner, repo, author))) { @@ -117,7 +123,7 @@ async function upsertComment(octokit, params, body) { // The pull request that triggered this run is already open, so exclude it and // report the total separately. const others = allOpen.filter( - (pr) => pr.user?.login === author && pr.number !== prNumber + (pr) => pr.user?.login === author && pr.number !== prNumber && !pr.draft ); const total = others.length + 1; diff --git a/.github/workflows/pr-open-limit.yml b/.github/workflows/pr-open-limit.yml index de5b723cb25..081454cb00e 100644 --- a/.github/workflows/pr-open-limit.yml +++ b/.github/workflows/pr-open-limit.yml @@ -5,15 +5,15 @@ on: types: - opened - reopened + - ready_for_review concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number }} cancel-in-progress: true -# The warning comment goes through the issues API, hence `issues: write`. permissions: contents: read - pull-requests: read + pull-requests: write issues: write jobs: @@ -44,4 +44,5 @@ jobs: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} PR_NUMBER: ${{ github.event.pull_request.number }} PR_AUTHOR: ${{ github.event.pull_request.user.login }} + PR_IS_DRAFT: ${{ github.event.pull_request.draft }} MAX_OPEN_PRS: ${{ vars.MAX_OPEN_PRS_PER_AUTHOR || 5 }} From 97cbf79fb6dfe9216e75ab2f22bbca32b61f6f67 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Tue, 11 Aug 2026 16:01:06 +0800 Subject: [PATCH 23/92] feat(promql): support native histogram vector operators (#8798) * feat(promql): support native histogram vector operators Signed-off-by: shuiyisong * fix: CR issue Signed-off-by: shuiyisong * chore: rebase main Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong * fix: cr issue Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- .../src/extension_plan/instant_manipulate.rs | 72 +- src/promql/src/functions/native_histogram.rs | 20 +- src/query/src/promql/planner.rs | 1741 +++++++++++++++-- 3 files changed, 1668 insertions(+), 165 deletions(-) diff --git a/src/promql/src/extension_plan/instant_manipulate.rs b/src/promql/src/extension_plan/instant_manipulate.rs index f92a4e4490f..4851916a017 100644 --- a/src/promql/src/extension_plan/instant_manipulate.rs +++ b/src/promql/src/extension_plan/instant_manipulate.rs @@ -18,6 +18,7 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; +use common_query::prelude::{greptime_native_histogram, greptime_value}; use datafusion::arrow::array::{Array, TimestampMillisecondArray, UInt64Array}; use datafusion::arrow::datatypes::{DataType, SchemaRef}; use datafusion::arrow::record_batch::RecordBatch; @@ -52,6 +53,15 @@ use crate::metrics::PROMQL_SERIES_COUNT; const MAX_INSTANT_MANIPULATE_OUTPUT_POINTS: usize = 1_000_000; +fn mixed_sample_fields(field: Option<&str>) -> [Option<&str>; 2] { + let companion = match field { + Some(field) if field == greptime_value() => Some(greptime_native_histogram()), + Some(field) if field == greptime_native_histogram() => Some(greptime_value()), + _ => None, + }; + [field, companion] +} + /// Manipulate the input record batch to make it suitable for Instant Operator. /// /// This plan will try to align the input time series, for every timestamp between @@ -66,7 +76,7 @@ pub struct InstantManipulate { time_index_column: String, // Planner-provided tag-column hint for execution fast paths. tag_columns: Vec, - /// A optional column for validating staleness + /// Primary sample column used to derive the columns checked for staleness. field_column: Option, input: LogicalPlan, unfix: Option, @@ -97,9 +107,7 @@ impl UserDefinedLogicalNodeCore for InstantManipulate { } let mut exprs = vec![col(&self.time_index_column)]; - if let Some(field) = &self.field_column { - exprs.push(col(field)); - } + exprs.extend(self.staleness_field_columns().map(col)); exprs } @@ -116,7 +124,7 @@ impl UserDefinedLogicalNodeCore for InstantManipulate { let mut required = output_columns.to_vec(); required.push(input_schema.index_of_column_by_name(None, &self.time_index_column)?); - if let Some(field) = &self.field_column { + for field in self.staleness_field_columns() { required.push(input_schema.index_of_column_by_name(None, field)?); } @@ -223,6 +231,21 @@ impl InstantManipulate { "InstantManipulate" } + fn staleness_field_columns(&self) -> impl Iterator { + let [field, companion] = mixed_sample_fields(self.field_column.as_deref()); + [ + field, + companion.filter(|companion| { + self.input + .schema() + .index_of_column_by_name(None, companion) + .is_some() + }), + ] + .into_iter() + .flatten() + } + fn resolve_tag_columns(input: &LogicalPlan, tag_columns: &[String]) -> Vec { if !tag_columns.is_empty() { return tag_columns.to_vec(); @@ -385,11 +408,9 @@ impl ExecutionPlan for InstantManipulateExec { .column_with_name(&self.time_index_column) .expect("time index column not found") .0; - let field_index = self - .field_column - .as_ref() - .and_then(|name| schema.column_with_name(name)) - .map(|x| x.0); + let field_indices = mixed_sample_fields(self.field_column.as_deref()).map(|field| { + field.and_then(|field| schema.column_with_name(field).map(|(index, _)| index)) + }); let tsid_index = schema .column_with_name("__tsid") .filter(|(_, field)| field.data_type() == &DataType::UInt64) @@ -400,7 +421,7 @@ impl ExecutionPlan for InstantManipulateExec { lookback_delta: self.lookback_delta, interval: self.interval, time_index, - field_index, + field_indices, tsid_index, reuse_tsid_column: self.reuse_tsid_column && tsid_index.is_some(), schema, @@ -467,7 +488,7 @@ pub struct InstantManipulateStream { interval: Millisecond, // Column index of TIME INDEX column's position in schema time_index: usize, - field_index: Option, + field_indices: [Option; 2], tsid_index: Option, reuse_tsid_column: bool, @@ -532,11 +553,16 @@ impl InstantManipulateStream { return Ok(input); } - // Field column for staleness checks, classified once per batch. - let stale_sample_column = self - .field_index - .map(|index| input.column(index).as_ref()) - .and_then(prometheus_stale_sample_column); + // Field columns for staleness checks, classified once per batch. + let stale_sample_columns = self.field_indices.map(|index| { + index.and_then(|index| prometheus_stale_sample_column(input.column(index).as_ref())) + }); + let is_stale = |row| { + stale_sample_columns + .iter() + .flatten() + .any(|column| is_prometheus_stale_sample(*column, row)) + }; // Optimize iteration range based on actual data bounds let first_ts = ts_column.value(0); @@ -581,9 +607,7 @@ impl InstantManipulateStream { let curr = ts_column.value(cursor); match curr.cmp(&expected_ts) { Ordering::Equal => { - if stale_sample_column - .is_some_and(|column| is_prometheus_stale_sample(column, cursor)) - { + if is_stale(cursor) { // Ignore the stale marker. } else { take_indices.push(cursor as u64); @@ -615,9 +639,7 @@ impl InstantManipulateStream { let prev_ts = ts_column.value(prev_cursor); if prev_ts + self.lookback_delta > expected_ts { // only use the point in the time range - if stale_sample_column - .is_some_and(|column| is_prometheus_stale_sample(column, prev_cursor)) - { + if is_stale(prev_cursor) { // Do not use a stale marker as the newest value. continue; } @@ -626,9 +648,7 @@ impl InstantManipulateStream { aligned_ts.push(expected_ts); } } - } else if stale_sample_column - .is_some_and(|column| is_prometheus_stale_sample(column, cursor)) - { + } else if is_stale(cursor) { // Do not use a stale marker as the newest value. } else { // use this point diff --git a/src/promql/src/functions/native_histogram.rs b/src/promql/src/functions/native_histogram.rs index 1863e663b82..e6c9e84d59f 100644 --- a/src/promql/src/functions/native_histogram.rs +++ b/src/promql/src/functions/native_histogram.rs @@ -173,12 +173,20 @@ impl ScalarUDFImpl for NativeHistogramAnnotationUdf { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DfResult { - if let Some(collector) = args - .config_options - .extensions - .get::() - .cloned() - .or_else(|| self.collector.clone()) + let has_dropped_sample = !args.args.is_empty() + && (0..args.number_rows).any(|row| { + args.args.iter().all(|arg| match arg { + ColumnarValue::Array(array) => array.is_valid(row), + ColumnarValue::Scalar(value) => !value.is_null(), + }) + }); + if has_dropped_sample + && let Some(collector) = args + .config_options + .extensions + .get::() + .cloned() + .or_else(|| self.collector.clone()) { collector.record_info(self.message.clone()); } diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index d21a90bf123..10462ea1828 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -64,14 +64,17 @@ use promql::extension_plan::{ use promql::functions::{ AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, MixedRange, - NativeHistogramAbsentOverTime, NativeHistogramAvg, NativeHistogramAvgOverTime, - NativeHistogramChanges, NativeHistogramCount, NativeHistogramCountOverTime, - NativeHistogramDelta, NativeHistogramDrop, NativeHistogramFraction, NativeHistogramIDelta, + NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAvg, + NativeHistogramAvgOverTime, NativeHistogramChanges, NativeHistogramCount, + NativeHistogramCountOverTime, NativeHistogramDelta, NativeHistogramDivScalar, + NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease, NativeHistogramLastOverTime, + NativeHistogramMulScalar, NativeHistogramNeg, NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile, NativeHistogramRate, - NativeHistogramResets, NativeHistogramStddev, NativeHistogramStdvar, NativeHistogramSum, - NativeHistogramSumOverTime, PredictLinear, PresentOverTime, QuantileOverTime, Rate, Resets, - Round, StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf, + NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar, + NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, PredictLinear, + PresentOverTime, QuantileOverTime, Rate, Resets, Round, StddevOverTime, StdvarOverTime, + SumOverTime, quantile_udaf, }; use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers}; use promql_parser::parser::token::TokenType; @@ -392,6 +395,8 @@ pub struct PromPlanner { promql_annotations: Option, } +type BinaryFieldPair<'a> = (&'a String, &'a String); + impl PromPlanner { pub async fn stmt_to_plan( table_provider: DfTableSourceProvider, @@ -823,8 +828,20 @@ impl PromPlanner { let UnaryExpr { expr } = unary_expr; // Unary Expr in PromQL implys the `-` operator let input = self.prom_expr_to_plan(expr, query_engine_state).await?; + self.negate_field_columns(input) + } + + fn negate_field_columns(&mut self, input: LogicalPlan) -> Result { + let input_schema = input.schema().clone(); self.projection_for_each_field_column(input, |col| { - Ok(DfExpr::Negative(Box::new(DfExpr::Column(col.into())))) + if Self::field_column_is_native_histogram(&input_schema, col) { + Ok(DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramNeg::scalar_udf()), + args: vec![DfExpr::Column(col.into())], + })) + } else { + Ok(DfExpr::Negative(Box::new(DfExpr::Column(col.into())))) + } }) } @@ -866,6 +883,16 @@ impl PromPlanner { }); } + if planned_leaves.iter().any(|leaf| { + Self::field_columns_contain_native_histogram( + leaf.plan.schema(), + &leaf.ctx.field_columns, + ) + }) { + self.ctx = original_ctx; + return Ok(None); + } + if !Self::binary_island_join_contexts_supported(&planned_leaves) { self.ctx = original_ctx; return Ok(None); @@ -1214,10 +1241,46 @@ impl PromPlanner { if let Some(time_expr) = self.try_build_special_time_expr_with_context(lhs) { expr = time_expr } + let input_schema = input.schema().clone(); + let preserve_any_value = Self::field_columns_are_alternative_samples( + &input_schema, + &self.ctx.field_columns, + ); + let has_native_histogram = Self::field_columns_contain_native_histogram( + &input_schema, + &self.ctx.field_columns, + ); + let retain_field_columns = self + .ctx + .field_columns + .iter() + .map(|col| { + Self::binary_result_is_histogram( + *op, + false, + Self::field_column_is_native_histogram(&input_schema, col), + ) + .is_some() + }) + .collect(); + let promql_annotations = self.promql_annotations.clone(); let bin_expr_builder = |col: &String| { let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; - let mut binary_expr = - binary_expr_builder(expr.clone(), DfExpr::Column(col.into()))?; + let rhs_is_histogram = + Self::field_column_is_native_histogram(&input_schema, col); + let rhs = DfExpr::Column(col.into()); + let mut binary_expr = match Self::native_histogram_binary_expr( + *op, + expr.clone(), + false, + rhs.clone(), + rhs_is_histogram, + is_comparison_op && !should_return_bool, + promql_annotations.clone(), + )? { + Some(expr) => expr, + None => binary_expr_builder(expr.clone(), rhs)?, + }; if is_comparison_op && should_return_bool { binary_expr = DfExpr::Cast(Cast { @@ -1230,7 +1293,14 @@ impl PromPlanner { if is_comparison_op && !should_return_bool { self.filter_on_field_column(input, bin_expr_builder) } else { - self.projection_for_each_field_column(input, bin_expr_builder) + let projected = + self.projection_for_each_field_column(input, bin_expr_builder)?; + self.filter_binary_projection( + projected, + has_native_histogram, + preserve_any_value, + retain_field_columns, + ) } } // lhs is a column, rhs is a literal @@ -1240,10 +1310,46 @@ impl PromPlanner { if let Some(time_expr) = self.try_build_special_time_expr_with_context(rhs) { expr = time_expr } + let input_schema = input.schema().clone(); + let preserve_any_value = Self::field_columns_are_alternative_samples( + &input_schema, + &self.ctx.field_columns, + ); + let has_native_histogram = Self::field_columns_contain_native_histogram( + &input_schema, + &self.ctx.field_columns, + ); + let retain_field_columns = self + .ctx + .field_columns + .iter() + .map(|col| { + Self::binary_result_is_histogram( + *op, + Self::field_column_is_native_histogram(&input_schema, col), + false, + ) + .is_some() + }) + .collect(); + let promql_annotations = self.promql_annotations.clone(); let bin_expr_builder = |col: &String| { let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; - let mut binary_expr = - binary_expr_builder(DfExpr::Column(col.into()), expr.clone())?; + let lhs_is_histogram = + Self::field_column_is_native_histogram(&input_schema, col); + let lhs = DfExpr::Column(col.into()); + let mut binary_expr = match Self::native_histogram_binary_expr( + *op, + lhs.clone(), + lhs_is_histogram, + expr.clone(), + false, + is_comparison_op && !should_return_bool, + promql_annotations.clone(), + )? { + Some(expr) => expr, + None => binary_expr_builder(lhs, expr.clone())?, + }; if is_comparison_op && should_return_bool { binary_expr = DfExpr::Cast(Cast { @@ -1256,7 +1362,14 @@ impl PromPlanner { if is_comparison_op && !should_return_bool { self.filter_on_field_column(input, bin_expr_builder) } else { - self.projection_for_each_field_column(input, bin_expr_builder) + let projected = + self.projection_for_each_field_column(input, bin_expr_builder)?; + self.filter_binary_projection( + projected, + has_native_histogram, + preserve_any_value, + retain_field_columns, + ) } } // both are columns. join them on time index @@ -1291,6 +1404,14 @@ impl PromPlanner { ); } + let has_native_histogram = Self::field_columns_contain_native_histogram( + left_input.schema(), + &left_field_columns, + ) || Self::field_columns_contain_native_histogram( + right_input.schema(), + &right_field_columns, + ); + // normal join if left_table_ref == right_table_ref { // rename table references to avoid ambiguity @@ -1308,21 +1429,41 @@ impl PromPlanner { self.ctx.table_name = Some("rhs".to_string()); } } - let (output_field_columns, field_columns) = - Self::align_binary_field_columns(&left_field_columns, &right_field_columns); - let left_aligned_field_columns = field_columns + // Computed scalars reach this join path instead of the literal projection paths. + // Broadcast them for arithmetic in the same way as literal scalars. + let broadcast_scalar = !is_comparison_op; + let (field_groups, invalid_field_pairs) = Self::align_binary_field_columns( + left_input.schema(), + right_input.schema(), + &left_field_columns, + &right_field_columns, + *op, + broadcast_scalar && lhs.value_type() == ValueType::Scalar, + broadcast_scalar && rhs.value_type() == ValueType::Scalar, + ); + let left_aligned_field_columns = field_groups .iter() - .map(|(left_col_name, _)| (*left_col_name).clone()) + .flat_map(|(_, pairs)| { + pairs + .iter() + .map(|(left_col_name, _)| (*left_col_name).clone()) + }) .collect::>(); - let right_aligned_field_columns = field_columns + let right_aligned_field_columns = field_groups .iter() - .map(|(_, right_col_name)| (*right_col_name).clone()) + .flat_map(|(_, pairs)| { + pairs + .iter() + .map(|(_, right_col_name)| (*right_col_name).clone()) + }) .collect::>(); - // PromQL binary arithmetic only combines the shared prefix of value columns. - // Keep the output field count aligned with that zipped prefix so planning - // remains stable even when the two sides have uneven multi-field schemas. - self.ctx.field_columns = output_field_columns; - let mut field_columns = field_columns.into_iter(); + // Regular multi-field vectors combine their shared prefix. Alternative + // float/histogram lanes instead align by valid PromQL sample combinations. + self.ctx.field_columns = field_groups + .iter() + .map(|(output, _)| output.clone()) + .collect(); + let mut field_groups = field_groups.into_iter(); let join_plan = self.join_on_non_field_columns( left_input, @@ -1339,28 +1480,103 @@ impl PromPlanner { &right_context, )?; let join_plan_schema = join_plan.schema().clone(); + let promql_annotations = self.promql_annotations.clone(); + // These predicates always pass; they only evaluate otherwise-discarded pairs + // while collecting annotations. + let invalid_pair_predicates = invalid_field_pairs + .into_iter() + .filter(|_| promql_annotations.is_some()) + .map(|(left_col_name, right_col_name)| { + let left_field = join_plan_schema + .qualified_field_with_name(Some(&left_table_ref), left_col_name) + .context(DataFusionPlanningSnafu)?; + let right_field = join_plan_schema + .qualified_field_with_name(Some(&right_table_ref), right_col_name) + .context(DataFusionPlanningSnafu)?; + let left_is_histogram = + left_field.1.data_type() == &Self::native_histogram_arrow_type(); + let right_is_histogram = + right_field.1.data_type() == &Self::native_histogram_arrow_type(); + let drop_expr = Self::native_histogram_binary_expr( + *op, + DfExpr::Column(left_field.into()), + left_is_histogram, + DfExpr::Column(right_field.into()), + right_is_histogram, + true, + promql_annotations.clone(), + )? + .with_context(|| UnexpectedPlanExprSnafu { + desc: "invalid native histogram pair produced no drop expression", + })?; + Ok(DfExpr::Not(Box::new(drop_expr))) + }) + .collect::>>()?; + let join_plan = if let Some(predicate) = conjunction(invalid_pair_predicates) { + LogicalPlanBuilder::from(join_plan) + .filter(predicate) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)? + } else { + join_plan + }; let bin_expr_builder = |_: &String| { - let (left_col_name, right_col_name) = field_columns.next().unwrap(); - let left_col = join_plan_schema - .qualified_field_with_name(Some(&left_table_ref), left_col_name) - .context(DataFusionPlanningSnafu)? - .into(); - let right_col = join_plan_schema - .qualified_field_with_name(Some(&right_table_ref), right_col_name) - .context(DataFusionPlanningSnafu)? - .into(); + let (_, field_pairs) = + field_groups + .next() + .with_context(|| UnexpectedPlanExprSnafu { + desc: "missing binary field group", + })?; + let binary_exprs = field_pairs + .into_iter() + .map(|(left_col_name, right_col_name)| { + let left_field = join_plan_schema + .qualified_field_with_name(Some(&left_table_ref), left_col_name) + .context(DataFusionPlanningSnafu)?; + let right_field = join_plan_schema + .qualified_field_with_name(Some(&right_table_ref), right_col_name) + .context(DataFusionPlanningSnafu)?; + let left_is_histogram = + left_field.1.data_type() == &Self::native_histogram_arrow_type(); + let right_is_histogram = + right_field.1.data_type() == &Self::native_histogram_arrow_type(); + let left_col = left_field.into(); + let right_col = right_field.into(); - let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; - let mut binary_expr = - binary_expr_builder(DfExpr::Column(left_col), DfExpr::Column(right_col))?; - if is_comparison_op && should_return_bool { - binary_expr = DfExpr::Cast(Cast { - expr: Box::new(binary_expr), - data_type: ArrowDataType::Float64, - }); + let binary_expr_builder = Self::prom_token_to_binary_expr_builder(*op)?; + let lhs = DfExpr::Column(left_col); + let rhs = DfExpr::Column(right_col); + let mut binary_expr = match Self::native_histogram_binary_expr( + *op, + lhs.clone(), + left_is_histogram, + rhs.clone(), + right_is_histogram, + is_comparison_op && !should_return_bool, + promql_annotations.clone(), + )? { + Some(expr) => expr, + None => binary_expr_builder(lhs, rhs)?, + }; + if is_comparison_op && should_return_bool { + binary_expr = DfExpr::Cast(Cast { + expr: Box::new(binary_expr), + data_type: ArrowDataType::Float64, + }); + } + Ok(binary_expr) + }) + .collect::>>()?; + if let [binary_expr] = binary_exprs.as_slice() { + Ok(binary_expr.clone()) + } else { + Ok(DfExpr::ScalarFunction(ScalarFunction { + func: coalesce(), + args: binary_exprs, + })) } - Ok(binary_expr) }; if is_comparison_op && !should_return_bool { // PromQL comparison operators without `bool` are filters: @@ -1386,12 +1602,90 @@ impl PromPlanner { project_context.field_columns = project_field_columns; self.project_binary_join_side(filtered, project_table_ref, &project_context) } else { - self.projection_for_each_field_column(join_plan, bin_expr_builder) + let projected = + self.projection_for_each_field_column(join_plan, bin_expr_builder)?; + let preserve_any_value = Self::field_columns_are_alternative_samples( + projected.schema(), + &self.ctx.field_columns, + ); + let retain_field_columns = vec![true; self.ctx.field_columns.len()]; + self.filter_binary_projection( + projected, + has_native_histogram, + preserve_any_value, + retain_field_columns, + ) } } } } + fn filter_binary_projection( + &mut self, + input: LogicalPlan, + has_native_histogram: bool, + preserve_any_value: bool, + retain_field_columns: Vec, + ) -> Result { + if !has_native_histogram { + return Ok(input); + } + + ensure!( + retain_field_columns.len() == self.ctx.field_columns.len(), + UnexpectedPlanExprSnafu { + desc: "binary output field count changed unexpectedly", + } + ); + + let filtered = LogicalPlanBuilder::from(input) + .filter(self.create_empty_values_filter_expr(preserve_any_value)?) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + if retain_field_columns.iter().all(|retain| *retain) { + return Ok(filtered); + } + + let retained = self + .ctx + .field_columns + .iter() + .zip(retain_field_columns) + .filter(|(_, retain)| *retain) + .map(|(field, _)| field.clone()) + .collect::>(); + if retained.is_empty() { + return Ok(filtered); + } + self.ctx.field_columns = retained; + + let mut output_columns = self + .ctx + .field_columns + .iter() + .chain(&self.ctx.tag_columns) + .cloned() + .collect::>(); + output_columns.extend(self.ctx.time_index_column.iter().cloned()); + if self.ctx.use_tsid { + output_columns.insert(DATA_SCHEMA_TSID_COLUMN_NAME.to_string()); + } + let project_exprs = filtered + .schema() + .iter() + .filter(|(_, field)| output_columns.contains(field.name())) + .map(|(qualifier, field)| { + DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) + }) + .collect::>(); + LogicalPlanBuilder::from(filtered) + .project(project_exprs) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + fn project_binary_join_side( &mut self, input: LogicalPlan, @@ -4416,6 +4710,69 @@ impl PromPlanner { } } + fn native_histogram_binary_expr( + token: TokenType, + lhs: DfExpr, + lhs_is_histogram: bool, + rhs: DfExpr, + rhs_is_histogram: bool, + filter_context: bool, + promql_annotations: Option, + ) -> Result> { + if !lhs_is_histogram && !rhs_is_histogram { + return Ok(None); + } + + let scalar_fn = |func: ScalarUdfDef, args| { + DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(func), + args, + }) + }; + let invalid_expr = || { + let message = format!( + "{}: dropped native histogram samples because this binary operation is not supported for native histograms", + token + ); + let func = if filter_context { + NativeHistogramDrop::bool_false_udf(message, promql_annotations.clone()) + } else { + NativeHistogramDrop::float_null_udf(message, promql_annotations.clone()) + }; + let args = vec![lhs.clone(), rhs.clone()]; + scalar_fn(func, args) + }; + + let expr = match (token.id(), lhs_is_histogram, rhs_is_histogram) { + (token::T_ADD, true, true) => scalar_fn( + NativeHistogramAdd::scalar_udf_with_collector(promql_annotations.clone()), + vec![lhs, rhs], + ), + (token::T_SUB, true, true) => scalar_fn( + NativeHistogramSub::scalar_udf_with_collector(promql_annotations.clone()), + vec![lhs, rhs], + ), + (token::T_MUL, true, false) => { + scalar_fn(NativeHistogramMulScalar::scalar_udf(), vec![lhs, rhs]) + } + (token::T_MUL, false, true) => { + scalar_fn(NativeHistogramScalarMul::scalar_udf(), vec![lhs, rhs]) + } + (token::T_DIV, true, false) => { + scalar_fn(NativeHistogramDivScalar::scalar_udf(), vec![lhs, rhs]) + } + (token::T_EQLC, true, true) => { + scalar_fn(NativeHistogramEq::scalar_udf(), vec![lhs, rhs]) + } + (token::T_NEQ, true, true) => { + scalar_fn(NativeHistogramNotEq::scalar_udf(), vec![lhs, rhs]) + } + _ => invalid_expr(), + }; + + Ok(Some(expr)) + } + /// Return a lambda to build binary expression from token. /// Because some binary operator are function in DataFusion like `atan2` or `^`. #[allow(clippy::type_complexity)] @@ -4501,18 +4858,117 @@ impl PromPlanner { } fn align_binary_field_columns<'a>( + left_schema: &DFSchemaRef, + right_schema: &DFSchemaRef, left_field_columns: &'a [String], right_field_columns: &'a [String], - ) -> (Vec, Vec<(&'a String, &'a String)>) { - let field_pairs = left_field_columns - .iter() - .zip(right_field_columns.iter()) - .collect::>(); - let output_field_columns = field_pairs - .iter() - .map(|(left_col_name, _)| (*left_col_name).clone()) - .collect(); - (output_field_columns, field_pairs) + op: TokenType, + left_is_scalar: bool, + right_is_scalar: bool, + ) -> ( + Vec<(String, Vec>)>, + Vec>, + ) { + // Mixed vectors store mutually exclusive float and histogram samples in two columns. + // Retain each valid sample combination and group expressions by their output lane. + let left_alternative = Self::alternative_sample_columns(left_schema, left_field_columns); + let right_alternative = Self::alternative_sample_columns(right_schema, right_field_columns); + let alternative_alignment = match (left_alternative, right_alternative) { + (Some(output_names), Some(_)) => Some(( + output_names, + left_field_columns + .iter() + .flat_map(|left| right_field_columns.iter().map(move |right| (left, right))) + .collect::>(), + )), + (Some(output_names), None) if right_field_columns.len() == 1 => Some(( + output_names, + left_field_columns + .iter() + .map(|left| (left, &right_field_columns[0])) + .collect::>(), + )), + (None, Some(output_names)) if left_field_columns.len() == 1 => Some(( + output_names, + right_field_columns + .iter() + .map(|right| (&left_field_columns[0], right)) + .collect::>(), + )), + _ => None, + }; + let mut invalid_pairs = Vec::new(); + if let Some(((float_output, histogram_output), field_pairs)) = alternative_alignment { + let mut float_pairs = Vec::new(); + let mut histogram_pairs = Vec::new(); + for (left, right) in field_pairs { + let left_is_histogram = Self::field_column_is_native_histogram(left_schema, left); + let right_is_histogram = + Self::field_column_is_native_histogram(right_schema, right); + match Self::binary_result_is_histogram(op, left_is_histogram, right_is_histogram) { + Some(false) => float_pairs.push((left, right)), + Some(true) => histogram_pairs.push((left, right)), + None => invalid_pairs.push((left, right)), + } + } + if !float_pairs.is_empty() || !histogram_pairs.is_empty() { + return ( + [ + (!float_pairs.is_empty()).then(|| (float_output.to_string(), float_pairs)), + (!histogram_pairs.is_empty()) + .then(|| (histogram_output.to_string(), histogram_pairs)), + ] + .into_iter() + .flatten() + .collect(), + invalid_pairs, + ); + } + } + + if left_is_scalar && !right_is_scalar && left_field_columns.len() == 1 { + return ( + right_field_columns + .iter() + .map(|right| (right.clone(), vec![(&left_field_columns[0], right)])) + .collect(), + invalid_pairs, + ); + } + if right_is_scalar && !left_is_scalar && right_field_columns.len() == 1 { + return ( + left_field_columns + .iter() + .map(|left| (left.clone(), vec![(left, &right_field_columns[0])])) + .collect(), + invalid_pairs, + ); + } + + ( + left_field_columns + .iter() + .zip(right_field_columns.iter()) + .map(|(left, right)| (left.clone(), vec![(left, right)])) + .collect(), + invalid_pairs, + ) + } + + fn binary_result_is_histogram( + token: TokenType, + lhs_is_histogram: bool, + rhs_is_histogram: bool, + ) -> Option { + match (token.id(), lhs_is_histogram, rhs_is_histogram) { + (_, false, false) => Some(false), + (token::T_ADD | token::T_SUB, true, true) + | (token::T_MUL, true, false) + | (token::T_MUL, false, true) + | (token::T_DIV, true, false) => Some(true), + (token::T_EQLC | token::T_NEQ, true, true) => Some(false), + _ => None, + } } fn plan_has_tsid_column(plan: &LogicalPlan) -> bool { @@ -4540,6 +4996,15 @@ impl PromPlanner { .is_some_and(|data_type| data_type == &Self::native_histogram_arrow_type()) } + fn field_columns_contain_native_histogram( + schema: &DFSchemaRef, + field_columns: &[String], + ) -> bool { + field_columns + .iter() + .any(|field| Self::field_column_is_native_histogram(schema, field)) + } + fn field_column_is_float_range(schema: &DFSchemaRef, field_column: &str) -> bool { Self::field_column_type(schema, field_column).is_some_and(|data_type| { matches!( @@ -4940,9 +5405,13 @@ impl PromPlanner { } ensure!( - left_context.field_columns.len() == 1, + left_context.field_columns.len() == 1 + || Self::field_columns_are_alternative_samples( + left.schema(), + &left_context.field_columns, + ), MultiFieldsNotSupportedSnafu { - operator: "AND operator" + operator: "AND/UNLESS operator" } ); // Generate join plan. @@ -5072,24 +5541,28 @@ impl PromPlanner { (false, false) => {} } - // checks ensure!( - left_context.field_columns.len() == right_context.field_columns.len(), - CombineTableColumnMismatchSnafu { - left: left_context.field_columns.clone(), - right: right_context.field_columns.clone() + !left.schema().fields().is_empty() && !right.schema().fields().is_empty(), + UnexpectedPlanExprSnafu { + desc: "OR operator input has zero columns", } ); + let left_has_alternative_samples = + Self::field_columns_are_alternative_samples(left.schema(), &left_context.field_columns); + let right_has_alternative_samples = Self::field_columns_are_alternative_samples( + right.schema(), + &right_context.field_columns, + ); ensure!( - left_context.field_columns.len() == 1, + left_context.field_columns.len() == 1 || left_has_alternative_samples, MultiFieldsNotSupportedSnafu { operator: "OR operator" } ); ensure!( - !left.schema().fields().is_empty() && !right.schema().fields().is_empty(), - UnexpectedPlanExprSnafu { - desc: "OR operator input has zero columns", + right_context.field_columns.len() == 1 || right_has_alternative_samples, + MultiFieldsNotSupportedSnafu { + operator: "OR operator" } ); @@ -5122,62 +5595,125 @@ impl PromPlanner { .with_context(|| TimeIndexNotFoundSnafu { table: right_qualifier_string.clone(), })?; - // Take the name of first field column. The length is checked above. - let left_field_col = left_context.field_columns.first().unwrap(); - let right_field_col = right_context.field_columns.first().unwrap(); - let left_field = left - .schema() + let native_histogram_type = Self::native_histogram_arrow_type(); + let is_numeric = |data_type: &ArrowDataType| { + matches!( + data_type, + ArrowDataType::Int8 + | ArrowDataType::Int16 + | ArrowDataType::Int32 + | ArrowDataType::Int64 + | ArrowDataType::UInt8 + | ArrowDataType::UInt16 + | ArrowDataType::UInt32 + | ArrowDataType::UInt64 + | ArrowDataType::Float32 + | ArrowDataType::Float64 + ) + }; + let left_fields = left_context + .field_columns .iter() - .find(|(_, field)| field.name() == left_field_col) - .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) - .with_context(|| ColumnNotFoundSnafu { - col: left_field_col.clone(), - })?; - let right_field = right - .schema() + .map(|name| { + left.schema() + .iter() + .find(|(_, field)| field.name() == name) + .map(|(qualifier, field)| { + (name.clone(), qualifier.cloned(), field.data_type().clone()) + }) + .with_context(|| ColumnNotFoundSnafu { col: name.clone() }) + }) + .collect::>>()?; + let right_fields = right_context + .field_columns .iter() - .find(|(_, field)| field.name() == right_field_col) - .map(|(qualifier, field)| (qualifier.cloned(), field.data_type().clone())) - .with_context(|| ColumnNotFoundSnafu { - col: right_field_col.clone(), - })?; - let target_field_type = if left_field.1 == right_field.1 { - left_field.1.clone() - } else if matches!( - left_field.1, - ArrowDataType::Int8 - | ArrowDataType::Int16 - | ArrowDataType::Int32 - | ArrowDataType::Int64 - | ArrowDataType::UInt8 - | ArrowDataType::UInt16 - | ArrowDataType::UInt32 - | ArrowDataType::UInt64 - | ArrowDataType::Float32 - | ArrowDataType::Float64 - ) && matches!( - right_field.1, - ArrowDataType::Int8 - | ArrowDataType::Int16 - | ArrowDataType::Int32 - | ArrowDataType::Int64 - | ArrowDataType::UInt8 - | ArrowDataType::UInt16 - | ArrowDataType::UInt32 - | ArrowDataType::UInt64 - | ArrowDataType::Float32 - | ArrowDataType::Float64 - ) { + .map(|name| { + right + .schema() + .iter() + .find(|(_, field)| field.name() == name) + .map(|(qualifier, field)| { + (name.clone(), qualifier.cloned(), field.data_type().clone()) + }) + .with_context(|| ColumnNotFoundSnafu { col: name.clone() }) + }) + .collect::>>()?; + let left_field = &left_fields[0]; + let right_field = &right_fields[0]; + let left_field_col = &left_field.0; + let right_field_col = &right_field.0; + let fields_are_samples = |fields: &[(String, Option, ArrowDataType)]| { + fields.iter().all(|(_, _, data_type)| { + is_numeric(data_type) || data_type == &native_histogram_type + }) + }; + let mixed_sample_types = if left_has_alternative_samples || right_has_alternative_samples { + if !fields_are_samples(&left_fields) || !fields_are_samples(&right_fields) { + return UnexpectedPlanExprSnafu { + desc: format!( + "OR value fields have incompatible types: {:?} and {:?}", + left_fields + .iter() + .map(|(_, _, data_type)| data_type) + .collect::>(), + right_fields + .iter() + .map(|(_, _, data_type)| data_type) + .collect::>() + ), + } + .fail(); + } + true + } else { + (left_field.2 == native_histogram_type && is_numeric(&right_field.2)) + || (right_field.2 == native_histogram_type && is_numeric(&left_field.2)) + }; + let target_field_type = if mixed_sample_types { + // Mixed vectors use the existing response representation: one nullable float column + // and one nullable native-histogram column. + ArrowDataType::Float64 + } else if left_field.2 == right_field.2 { + left_field.2.clone() + } else if is_numeric(&left_field.2) && is_numeric(&right_field.2) { ArrowDataType::Float64 } else { return UnexpectedPlanExprSnafu { desc: format!( "OR value fields have incompatible types: {:?} and {:?}", - left_field.1, right_field.1 + left_field.2, right_field.2 ), } .fail(); }; + let (mixed_float_field_col, mixed_histogram_field_col) = if mixed_sample_types { + let mut reserved_names = left + .schema() + .fields() + .iter() + .chain(right.schema().fields().iter()) + .map(|field| field.name().clone()) + .collect::>(); + for (name, _, _) in left_fields.iter().chain(&right_fields) { + reserved_names.remove(name); + } + reserved_names.extend(all_tags.iter().cloned()); + let unique_name = |prefix: &str, reserved_names: &mut HashSet| { + let mut index = 0; + loop { + let name = format!("{prefix}{index}"); + index += 1; + if reserved_names.insert(name.clone()) { + break name; + } + } + }; + let float_field = unique_name(OR_FLOAT_FIELD_PREFIX, &mut reserved_names); + let histogram_field = unique_name(OR_HISTOGRAM_FIELD_PREFIX, &mut reserved_names); + (float_field, histogram_field) + } else { + (left_field_col.clone(), String::new()) + }; let left_tag_types = left_tag_cols_set .iter() .map(|label| { @@ -5244,8 +5780,15 @@ impl PromPlanner { // remove time index column all_columns_set.remove(&left_time_index_column); all_columns_set.remove(&right_time_index_column); - // remove field column in the right - if left_field_col != right_field_col { + if mixed_sample_types { + for (name, _, _) in left_fields.iter().chain(&right_fields) { + all_columns_set.remove(name); + } + all_columns_set.extend(all_tags.iter().cloned()); + all_columns_set.insert(mixed_float_field_col.clone()); + all_columns_set.insert(mixed_histogram_field_col.clone()); + } else if left_field_col != right_field_col { + // remove field column in the right all_columns_set.remove(right_field_col); } let mut all_columns = all_columns_set.into_iter().collect::>(); @@ -5284,11 +5827,52 @@ impl PromPlanner { .alias(col.clone()) } }; + let null_histogram = + ScalarValue::try_new_null(&native_histogram_type).context(DataFusionPlanningSnafu)?; + let mixed_value_expr = |fields: &[(String, Option, ArrowDataType)], + output_col: &String| { + if output_col == &mixed_float_field_col { + if let Some((name, qualifier, data_type)) = fields + .iter() + .find(|(_, _, data_type)| is_numeric(data_type)) + { + let expr = DfExpr::Column(Column::new(qualifier.clone(), name)); + if data_type == &ArrowDataType::Float64 { + expr.alias(output_col) + } else { + DfExpr::Cast(Cast { + expr: Box::new(expr), + data_type: ArrowDataType::Float64, + }) + .alias(output_col) + } + } else { + DfExpr::Literal(ScalarValue::Float64(None), None).alias(output_col) + } + } else { + fields + .iter() + .find(|(_, _, data_type)| data_type == &native_histogram_type) + .map(|(name, qualifier, _)| { + DfExpr::Column(Column::new(qualifier.clone(), name)).alias(output_col) + }) + .unwrap_or_else(|| { + DfExpr::Literal(null_histogram.clone(), None).alias(output_col) + }) + } + }; let left_proj_exprs = all_columns.iter().map(|col| { - if col == left_field_col && left_field.1 != target_field_type { + if mixed_sample_types + && (col == &mixed_float_field_col || col == &mixed_histogram_field_col) + { + mixed_value_expr(&left_fields, col) + } else if !mixed_sample_types + && col == left_field_col + && left_field.2 != target_field_type + { DfExpr::Cast(Cast { expr: Box::new(DfExpr::Column(Column::new( - left_field.0.clone(), + left_field.1.clone(), left_field_col, ))), data_type: target_field_type.clone(), @@ -5310,9 +5894,13 @@ impl PromPlanner { // `skip（1)` to skip the time index column let right_proj_exprs_without_time_index = all_columns.iter().skip(1).map(|col| { // expr - if col == left_field_col { - let expr = DfExpr::Column(Column::new(right_field.0.clone(), right_field_col)); - if right_field.1 != target_field_type { + if mixed_sample_types + && (col == &mixed_float_field_col || col == &mixed_histogram_field_col) + { + mixed_value_expr(&right_fields, col) + } else if !mixed_sample_types && col == left_field_col { + let expr = DfExpr::Column(Column::new(right_field.1.clone(), right_field_col)); + if right_field.2 != target_field_type { DfExpr::Cast(Cast { expr: Box::new(expr), data_type: target_field_type.clone(), @@ -5529,7 +6117,11 @@ impl PromPlanner { visible_tags.sort_unstable(); output_context.time_index_column = Some(left_time_index_column); output_context.tag_columns = visible_tags; - output_context.field_columns = vec![output_field_col]; + output_context.field_columns = if mixed_sample_types { + vec![mixed_float_field_col, mixed_histogram_field_col] + } else { + vec![output_field_col] + }; output_context.use_tsid = left_has_tsid && right_has_tsid; self.ctx = output_context; @@ -5551,6 +6143,10 @@ impl PromPlanner { where F: FnMut(&String) -> Result, { + // Keep the generated float/histogram lane names while an element-wise operation + // preserves both sample types, so downstream operators still recognize the pair. + let preserve_field_names = + Self::field_columns_are_alternative_samples(input.schema(), &self.ctx.field_columns); let table_ref = self.ctx.table_name.clone().map(TableReference::bare); let non_field_columns_iter = self .ctx @@ -5572,10 +6168,12 @@ impl PromPlanner { .collect::>>()?; // alias the computation exprs to remove qualifier - self.ctx.field_columns = result_field_columns - .iter() - .map(|expr| expr.schema_name().to_string()) - .collect(); + if !preserve_field_names { + self.ctx.field_columns = result_field_columns + .iter() + .map(|expr| expr.schema_name().to_string()) + .collect(); + } let field_columns_iter = result_field_columns .into_iter() .zip(self.ctx.field_columns.iter()) @@ -5594,24 +6192,32 @@ impl PromPlanner { .context(DataFusionPlanningSnafu) } - /// Build a filter plan that filter on value column. Notice that only one value column - /// is expected. - fn filter_on_field_column( - &self, - input: LogicalPlan, - mut name_to_expr: F, - ) -> Result + /// Build a filter plan on one value column or a float/histogram alternative pair. + fn filter_on_field_column(&self, input: LogicalPlan, name_to_expr: F) -> Result where F: FnMut(&String) -> Result, { ensure!( - self.ctx.field_columns.len() == 1, + self.ctx.field_columns.len() == 1 + || Self::field_columns_are_alternative_samples( + input.schema(), + &self.ctx.field_columns, + ), UnsupportedExprSnafu { name: "filter on multi-value input" } ); - let field_column_filter = name_to_expr(&self.ctx.field_columns[0])?; + let field_column_filters = self + .ctx + .field_columns + .iter() + .map(name_to_expr) + .collect::>>()?; + let field_column_filter = + disjunction(field_column_filters).context(UnsupportedExprSnafu { + name: "filter on empty input", + })?; LogicalPlanBuilder::from(input) .filter(field_column_filter) @@ -5743,7 +6349,9 @@ mod test { CounterResetHint, NativeHistogram, build_histogram_array, }; use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; + use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; use common_query::test_util::DummyDecoder; + use common_recordbatch::RecordBatch as GreptimeRecordBatch; use datafusion::arrow::array::{ Array, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, }; @@ -5761,8 +6369,9 @@ mod test { use promql_parser::parser; use session::context::QueryContext; use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan}; - use table::metadata::{TableInfoBuilder, TableMetaBuilder}; - use table::test_util::EmptyTable; + use table::Table; + use table::metadata::{FilterPushDownType, TableInfoBuilder, TableMetaBuilder}; + use table::test_util::{EmptyTable, MemTable as GreptimeMemTable}; use super::*; use crate::QueryEngineContext; @@ -5896,6 +6505,7 @@ mod test { enum DirectOrValue { Float64(f64), Int64(i64), + NativeHistogram(NativeHistogram), Utf8(&'static str), } @@ -5904,6 +6514,7 @@ mod test { match self { Self::Float64(_) => ArrowDataType::Float64, Self::Int64(_) => ArrowDataType::Int64, + Self::NativeHistogram(_) => native_histogram_value_type().as_arrow_type(), Self::Utf8(_) => ArrowDataType::Utf8, } } @@ -5911,6 +6522,7 @@ mod test { match self { Self::Float64(v) => Arc::new(Float64Array::from(vec![*v])), Self::Int64(v) => Arc::new(Int64Array::from(vec![*v])), + Self::NativeHistogram(v) => build_histogram_array(&[Some(v.clone())]), Self::Utf8(v) => Arc::new(StringArray::from(vec![*v])), } } @@ -5933,6 +6545,121 @@ mod test { } } + fn operator_metric_table( + name: &str, + table_id: u32, + tag: &str, + value: DirectOrValue, + ) -> table::TableRef { + let value_type = match &value { + DirectOrValue::Float64(_) => ConcreteDataType::float64_datatype(), + DirectOrValue::Int64(_) => ConcreteDataType::int64_datatype(), + DirectOrValue::NativeHistogram(_) => native_histogram_value_type().clone(), + DirectOrValue::Utf8(_) => ConcreteDataType::string_datatype(), + }; + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "tag".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + "ts".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("v".to_string(), value_type, true), + ])); + let batch = RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![ + Arc::new(StringArray::from(vec![tag])) as Arc, + Arc::new(TimestampMillisecondArray::from(vec![1_000])), + value.array(), + ], + ) + .unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + name, + GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch), + table_id, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .value_indices(vec![2]) + .next_column_id(3) + .build() + .unwrap(); + let info = Arc::new( + TableInfoBuilder::default() + .table_id(table_id) + .name(name) + .meta(meta) + .build() + .unwrap(), + ); + Arc::new(Table::new( + info, + FilterPushDownType::Unsupported, + backing.data_source(), + )) + } + + fn operator_table_provider() -> DfTableSourceProvider { + let catalog = MemoryCatalogManager::with_default_setup(); + let tables = [ + operator_metric_table("lf", 2_001, "a", DirectOrValue::Float64(2.0)), + operator_metric_table( + "lh", + 2_002, + "b", + DirectOrValue::NativeHistogram(direct_or_histogram()), + ), + operator_metric_table("rf", 2_003, "b", DirectOrValue::Float64(3.0)), + operator_metric_table( + "rh", + 2_004, + "a", + DirectOrValue::NativeHistogram(direct_or_histogram()), + ), + operator_metric_table("fallback", 2_005, "c", DirectOrValue::Float64(7.0)), + ]; + for table in tables { + let info = table.table_info(); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: info.name.clone(), + table_id: info.ident.table_id, + table, + }) + .unwrap(); + } + DfTableSourceProvider::new( + catalog, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) + } + + fn operator_eval_stmt(expr: &str) -> EvalStmt { + let time = UNIX_EPOCH.checked_add(Duration::from_secs(1)).unwrap(); + EvalStmt { + expr: parser::parse(expr).unwrap(), + start: time, + end: time, + interval: Duration::from_secs(1), + lookback_delta: Duration::from_secs(5), + } + } + struct DirectOrSource { name: &'static str, empty: bool, @@ -6093,6 +6820,66 @@ mod test { execute(plan, &build_query_engine_state()).await } + async fn mixed_direct_or(histogram_on_left: bool) -> (PromPlanner, LogicalPlan) { + let sample = |histogram: bool| { + if histogram { + DirectOrValue::NativeHistogram(direct_or_histogram()) + } else { + DirectOrValue::Float64(1.25) + } + }; + let left = tagged_source( + "lhs", + false, + ( + "k", + Some(if histogram_on_left { + "histogram" + } else { + "float" + }), + ), + sample(histogram_on_left), + ); + let right = tagged_source( + "rhs", + false, + ( + "k", + Some(if histogram_on_left { + "float" + } else { + "histogram" + }), + ), + sample(!histogram_on_left), + ); + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::default(), + promql_annotations: None, + }; + let left_context = direct_or_context("lhs", &["job", "k"], "v"); + let right_context = direct_or_context("rhs", &["job", "k"], "v"); + let plan = planner + .or_operator( + scan(&left), + scan(&right), + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + &or_modifier("lhs or on(k) rhs"), + ) + .unwrap(); + (planner, plan) + } + fn assert_no_internal_or_keys(schema: &DFSchema) { assert!( schema @@ -6119,6 +6906,23 @@ mod test { .collect() } + fn histograms(batches: &[RecordBatch], column: &str) -> Vec { + batches + .iter() + .flat_map(|batch| { + let values = batch + .column_by_name(column) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..values.len()).filter_map(|row| { + common_query::native_histogram::read_histogram(values, row).unwrap() + }) + }) + .collect() + } + fn rows(batches: &[RecordBatch]) -> Vec<(f64, Option)> { let mut rows = batches .iter() @@ -8833,6 +9637,110 @@ mod test { assert!(!plan.contains("PromHistogramFold"), "{plan}"); } + #[tokio::test] + async fn timestamp_filters_native_histogram_stale_marker_before_projection() { + let mut stale = direct_or_histogram(); + stale.sum = f64::from_bits(PROMETHEUS_STALE_NAN_BITS); + let table = operator_metric_table( + "stale_histogram", + 2_100, + "a", + DirectOrValue::NativeHistogram(stale), + ); + let catalog = MemoryCatalogManager::with_default_setup(); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "stale_histogram".to_string(), + table_id: 2_100, + table, + }) + .unwrap(); + let provider = DfTableSourceProvider::new( + catalog, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ); + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + provider, + &operator_eval_stmt("timestamp(stale_histogram)"), + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert!(plan_text.contains(TIMESTAMP_VALUE_PREFIX), "{plan_text}"); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + } + + #[tokio::test] + async fn timestamp_filters_stale_marker_from_mixed_sample_companion() { + let histograms = build_histogram_array(&[None]); + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "timestamp", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + ), + Field::new( + greptime_native_histogram(), + histograms.data_type().clone(), + true, + ), + Field::new(greptime_value(), ArrowDataType::Float64, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000])), + histograms, + Arc::new(Float64Array::from(vec![f64::from_bits( + PROMETHEUS_STALE_NAN_BITS, + )])), + ], + ) + .unwrap(); + let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); + let input = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) + .unwrap() + .build() + .unwrap(); + let input = LogicalPlan::Extension(Extension { + node: Arc::new(SeriesDivide::new( + Vec::new(), + "timestamp".to_string(), + input, + )), + }); + let input = LogicalPlan::Extension(Extension { + node: Arc::new(InstantManipulate::new( + 1_000, + 1_000, + 5_000, + 1_000, + "timestamp".to_string(), + Vec::new(), + Some(greptime_native_histogram().to_string()), + input, + )), + }); + // Match timestamp()'s parent projection, which otherwise prunes the companion lane. + let plan = LogicalPlanBuilder::from(input) + .project([col("timestamp")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + } + #[tokio::test] async fn native_histogram_rate_can_feed_count() { let plan = native_histogram_plan("histogram_count(rate(some_metric[5m]))").await; @@ -10654,7 +11562,6 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom .contains("OR value fields have incompatible types") ); } - #[tokio::test] async fn test_or_with_histogram_quantile_missing_le_column() { let case = r#"histogram_quantile(0.99, non_existent_histogram_bucket) or normal_metric"#; @@ -10871,4 +11778,572 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom "{plan:?}" ); } + + #[tokio::test] + async fn test_direct_or_preserves_float_and_native_histogram_samples() { + for histogram_on_left in [false, true] { + let (planner, plan) = mixed_direct_or(histogram_on_left).await; + + let float_field = &planner.ctx.field_columns[0]; + let histogram_field = &planner.ctx.field_columns[1]; + assert!(float_field.starts_with(OR_FLOAT_FIELD_PREFIX)); + assert!(histogram_field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)); + assert_eq!( + plan.schema() + .field_with_name(None, float_field) + .unwrap() + .data_type(), + &ArrowDataType::Float64 + ); + assert_eq!( + plan.schema() + .field_with_name(None, histogram_field) + .unwrap() + .data_type(), + &native_histogram_value_type().as_arrow_type() + ); + + let (optimized, batches) = execute(plan, &build_query_engine_state()).await; + assert_no_internal_or_keys(optimized.schema()); + let mut sample_kinds = batches + .iter() + .flat_map(|batch| { + let values = batch.column_by_name(float_field).unwrap(); + let histograms = batch.column_by_name(histogram_field).unwrap(); + (0..batch.num_rows()) + .map(|row| (values.is_valid(row), histograms.is_valid(row))) + }) + .collect::>(); + sample_kinds.sort_unstable(); + assert_eq!(sample_kinds, vec![(false, true), (true, false)]); + } + } + + #[tokio::test] + async fn test_mixed_binary_operator_aligns_both_alternative_inputs() { + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("(lf or on(tag) lh) * on(tag) (rf or on(tag) rh)"), + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert!( + plan_text.contains("prom_native_histogram_mul_scalar"), + "{plan_text}" + ); + assert!( + plan_text.contains("prom_native_histogram_scalar_mul"), + "{plan_text}" + ); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + let histogram_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + assert!(values(&batches, &float_field).is_empty()); + let mut sums = histograms(&batches, &histogram_field) + .into_iter() + .map(|histogram| histogram.sum) + .collect::>(); + sums.sort_by(f64::total_cmp); + assert_eq!(sums, vec![2.0, 3.0]); + } + + #[tokio::test] + async fn test_mixed_binary_operator_reports_only_dropped_samples() { + for (query, expected_rows, expected_infos) in [ + ("(lf or on(tag) lh) + on(tag) (rf or on(tag) rh)", 0, 1), + ("(lf or on(tag) lh) + on(tag) (lf or on(tag) lh)", 2, 0), + ("(lf or on(tag) lh) % on(tag) lh", 0, 1), + ] { + let state = build_query_engine_state(); + let annotations = PromqlAnnotationCollector::default(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + operator_table_provider(), + &operator_eval_stmt(query), + &state, + Some(annotations.clone()), + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected_rows, + "{query}" + ); + let mut warnings = vec![]; + let mut infos = vec![]; + annotations.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty(), "{query}: {warnings:?}"); + assert_eq!(infos.len(), expected_infos, "{query}: {infos:?}"); + } + } + + #[tokio::test] + async fn test_mixed_or_can_feed_another_or() { + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("lf or on(tag) lh or on(tag) fallback"), + &state, + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_FLOAT_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + let histogram_field = plan + .schema() + .fields() + .iter() + .find(|field| field.name().starts_with(OR_HISTOGRAM_FIELD_PREFIX)) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); + let mut float_values = values(&batches, &float_field); + float_values.sort_by(f64::total_cmp); + assert_eq!(float_values, vec![2.0, 7.0]); + assert_eq!(histograms(&batches, &histogram_field).len(), 1); + } + + #[tokio::test] + async fn test_mixed_fields_align_with_single_float_vector() { + let (planner, mixed) = mixed_direct_or(false).await; + let scale = tagged_source( + "scale", + false, + ("k", Some("float")), + DirectOrValue::Float64(2.0), + ); + let scale = scan(&scale); + let scale_fields = vec!["v".to_string()]; + let PromExpr::Binary(binary) = parser::parse("lhs * rhs").unwrap() else { + unreachable!() + }; + + let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns( + mixed.schema(), + scale.schema(), + &planner.ctx.field_columns, + &scale_fields, + binary.op, + false, + false, + ); + assert!(invalid_pairs.is_empty()); + assert_eq!( + groups + .iter() + .map(|(output, _)| output.clone()) + .collect::>(), + planner.ctx.field_columns + ); + assert_eq!(groups.len(), 2); + assert!( + groups + .iter() + .flat_map(|(_, pairs)| pairs) + .all(|(_, right)| *right == &scale_fields[0]) + ); + + let (groups, invalid_pairs) = PromPlanner::align_binary_field_columns( + scale.schema(), + mixed.schema(), + &scale_fields, + &planner.ctx.field_columns, + binary.op, + false, + false, + ); + assert!(invalid_pairs.is_empty()); + assert_eq!( + groups + .iter() + .map(|(output, _)| output.clone()) + .collect::>(), + planner.ctx.field_columns + ); + assert_eq!(groups.len(), 2); + assert!( + groups + .iter() + .flat_map(|(_, pairs)| pairs) + .all(|(left, _)| *left == &scale_fields[0]) + ); + } + + #[tokio::test] + async fn test_non_bool_comparison_filters_mixed_sample_lanes() { + let (planner, input) = mixed_direct_or(false).await; + let input_schema = input.schema().clone(); + let plan = planner + .filter_on_field_column(input, |field| { + if PromPlanner::field_column_is_native_histogram(&input_schema, field) { + Ok(lit(false)) + } else { + Ok(col(field).gt(lit(0.0))) + } + }) + .unwrap(); + let float_field = planner.ctx.field_columns[0].clone(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, &float_field), vec![1.25]); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + } + + #[tokio::test] + async fn test_mixed_left_and_unless_preserve_sample_lanes() { + for (expression, expected_sample_kind) in [ + ("lhs and on(k) mask", (false, true)), + ("lhs unless on(k) mask", (true, false)), + ] { + let (mut planner, left) = mixed_direct_or(false).await; + let left_context = planner.ctx.clone(); + let float_field = left_context.field_columns[0].clone(); + let histogram_field = left_context.field_columns[1].clone(); + let mask = tagged_source( + "mask", + false, + ("k", Some("histogram")), + DirectOrValue::Float64(1.0), + ); + let PromExpr::Binary(binary) = parser::parse(expression).unwrap() else { + unreachable!() + }; + let plan = planner + .set_op_on_non_field_columns( + left, + scan(&mask), + left_context, + direct_or_context("mask", &["job", "k"], "v"), + binary.op, + &binary.modifier, + ) + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + let sample_kinds = batches + .iter() + .flat_map(|batch| { + let floats = batch.column_by_name(&float_field).unwrap(); + let histograms = batch.column_by_name(&histogram_field).unwrap(); + (0..batch.num_rows()) + .map(|row| (floats.is_valid(row), histograms.is_valid(row))) + }) + .collect::>(); + assert_eq!(sample_kinds, vec![expected_sample_kind], "{expression}"); + } + } + + #[tokio::test] + async fn test_mixed_fields_arithmetic_broadcasts_computed_scalar() { + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &build_eval_stmt("some_metric * scalar(vector(2))"), + &build_query_engine_state(), + ) + .await + .unwrap(); + let schema = plan.schema(); + assert_eq!( + schema + .field_with_unqualified_name(greptime_value()) + .unwrap() + .data_type(), + &ArrowDataType::Float64 + ); + assert_eq!( + schema + .field_with_unqualified_name(greptime_native_histogram()) + .unwrap() + .data_type(), + &native_histogram_value_type().as_arrow_type() + ); + assert!( + plan.display_indent_schema() + .to_string() + .contains("prom_native_histogram_mul_scalar"), + "{plan:?}" + ); + } + + #[tokio::test] + async fn test_unsupported_histogram_binary_does_not_block_or_fallback() { + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("((lf or on(tag) lh) % 2) or on(tag) lh"), + &state, + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + let histogram_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &native_histogram_value_type().as_arrow_type()) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(values(&batches, &float_field), vec![0.0]); + assert_eq!(histograms(&batches, &histogram_field).len(), 1); + } + + #[tokio::test] + async fn test_unary_negates_mixed_float_and_native_histogram_samples() { + for histogram_on_left in [false, true] { + let (mut planner, input) = mixed_direct_or(histogram_on_left).await; + let plan = planner.negate_field_columns(input).unwrap(); + assert!(PromPlanner::field_columns_are_alternative_samples( + plan.schema(), + &planner.ctx.field_columns + )); + let float_field = planner + .ctx + .field_columns + .iter() + .find(|field| field.starts_with(OR_FLOAT_FIELD_PREFIX)) + .unwrap(); + let histogram_field = planner + .ctx + .field_columns + .iter() + .find(|field| field.starts_with(OR_HISTOGRAM_FIELD_PREFIX)) + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, float_field), vec![-1.25]); + let histogram = batches + .iter() + .find_map(|batch| { + let values = batch + .column_by_name(histogram_field) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..values.len()).find_map(|row| { + common_query::native_histogram::read_histogram(values, row).unwrap() + }) + }) + .unwrap(); + assert_eq!(histogram.count, -1.0); + assert_eq!(histogram.sum, -1.0); + assert_eq!(histogram.reset_hint, CounterResetHint::Gauge); + } + } + + #[tokio::test] + async fn test_mixed_or_value_aliases_do_not_replace_labels() { + let left = source( + "lhs", + false, + 1, + vec![("job", Some("job")), ("k", Some("float"))], + DirectOrValue::Float64(1.0), + ); + let right = source( + "rhs", + false, + 1, + vec![ + ("job", Some("job")), + ("k", Some("histogram")), + (greptime_value(), Some("value-label")), + ], + DirectOrValue::NativeHistogram(direct_or_histogram()), + ); + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::default(), + promql_annotations: None, + }; + let left = LogicalPlanBuilder::from(scan(&left)) + .project(vec![ + col("ts"), + col("job"), + col("k"), + col("v").alias(greptime_value()), + ]) + .unwrap() + .build() + .unwrap(); + let left_context = direct_or_context("lhs", &["job", "k"], greptime_value()); + let right_context = direct_or_context("rhs", &["job", "k", greptime_value()], "v"); + let plan = planner + .or_operator( + left, + scan(&right), + left_context.tag_columns.iter().cloned().collect(), + right_context.tag_columns.iter().cloned().collect(), + left_context, + right_context, + &or_modifier("lhs or on(k) rhs"), + ) + .unwrap(); + + assert_eq!( + plan.schema() + .field_with_name(None, greptime_value()) + .unwrap() + .data_type(), + &ArrowDataType::Utf8 + ); + assert!( + planner + .ctx + .field_columns + .iter() + .all(|field| { field != greptime_value() && field != greptime_native_histogram() }) + ); + assert!(PromPlanner::field_columns_are_alternative_samples( + plan.schema(), + &planner.ctx.field_columns + )); + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + let labels = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name(greptime_value()) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + }) + .collect::>(); + assert_eq!(labels, vec!["value-label"]); + } + + #[tokio::test] + async fn test_mixed_or_routes_float_histogram_and_label_functions() { + for (function, expected) in [("abs", 1.25), ("round", 1.0), ("histogram_count", 1.0)] { + let (mut planner, input) = mixed_direct_or(false).await; + let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( + input.schema(), + &planner.ctx.field_columns, + ); + let PromExpr::Call(call) = parser::parse(&format!("{function}(lhs)")).unwrap() else { + unreachable!() + }; + let state = build_query_engine_state(); + let (mut exprs, _) = planner + .create_function_expr(&call.func, vec![], input.schema(), &state) + .unwrap(); + exprs.insert(0, planner.create_time_index_column_expr().unwrap()); + exprs.extend(planner.create_tag_column_exprs().unwrap()); + let plan = LogicalPlanBuilder::from(input) + .project(exprs) + .unwrap() + .filter( + planner + .create_empty_values_filter_expr(preserve_any_value) + .unwrap(), + ) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(plan, &state).await; + let values = batches + .iter() + .flat_map(|batch| { + batch + .schema() + .fields() + .iter() + .position(|field| field.data_type() == &ArrowDataType::Float64) + .map(|index| { + batch + .column(index) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + }) + .into_iter() + .flatten() + }) + .collect::>(); + assert_eq!(values, vec![expected], "{function}"); + } + + let (mut planner, input) = mixed_direct_or(false).await; + let preserve_any_value = PromPlanner::field_columns_are_alternative_samples( + input.schema(), + &planner.ctx.field_columns, + ); + let PromExpr::Call(call) = + parser::parse(r#"label_replace(lhs, "copy", "$1", "k", "(.*)")"#).unwrap() + else { + unreachable!() + }; + let args = planner.create_function_args(&call.args.args).unwrap(); + let state = build_query_engine_state(); + let (mut exprs, _) = planner + .create_function_expr(&call.func, args.literals, input.schema(), &state) + .unwrap(); + exprs.insert(0, planner.create_time_index_column_expr().unwrap()); + exprs.extend(planner.create_tag_column_exprs().unwrap()); + let plan = LogicalPlanBuilder::from(input) + .project(exprs) + .unwrap() + .filter( + planner + .create_empty_values_filter_expr(preserve_any_value) + .unwrap(), + ) + .unwrap() + .build() + .unwrap(); + let (_, batches) = execute(plan, &state).await; + let sample_count = batches.iter().map(RecordBatch::num_rows).sum::(); + assert_eq!(sample_count, 2); + } } From 7375be06352146e84d8eab89313eabb341cb92c3 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Tue, 11 Aug 2026 16:27:03 +0800 Subject: [PATCH 24/92] ci: update code owners (#8843) Signed-off-by: evenyag --- .github/CODEOWNERS | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 549d6300ea1..e9e4d6ea561 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -6,8 +6,8 @@ ## [Module] Database Engine /src/index @evenyag @discord9 @WenyXu -/src/mito2 @evenyag @v0y4g3r @waynexia -/src/query @evenyag @waynexia @discord9 +/src/mito2 @evenyag @v0y4g3r +/src/query @evenyag @discord9 ## [Module] Distributed /src/common/meta @MichaelScofield @WenyXu @@ -20,8 +20,8 @@ /src/store-api @v0y4g3r @evenyag ## [Module] Metrics Engine -/src/metric-engine @waynexia @WenyXu -/src/promql @waynexia @evenyag @discord9 +/src/metric-engine @WenyXu @v0y4g3r +/src/promql @evenyag @discord9 ## [Module] Flow -/src/flow @discord9 @waynexia +/src/flow @discord9 @WenyXu From 133eda683670cfc497470ca8a827fad8cc879eb1 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 11 Aug 2026 16:28:23 +0800 Subject: [PATCH 25/92] refactor(operator): split semantic_graph relationship builders into a submodule (#8841) Signed-off-by: Dennis Zhuang --- src/operator/src/statement/semantic_graph.rs | 950 +----------------- .../statement/semantic_graph/relationships.rs | 905 +++++++++++++++++ 2 files changed, 959 insertions(+), 896 deletions(-) create mode 100644 src/operator/src/statement/semantic_graph/relationships.rs diff --git a/src/operator/src/statement/semantic_graph.rs b/src/operator/src/statement/semantic_graph.rs index 15f76905555..af9108198d4 100644 --- a/src/operator/src/statement/semantic_graph.rs +++ b/src/operator/src/statement/semantic_graph.rs @@ -25,6 +25,8 @@ //! including filter pushdown into the source table scans. See //! `docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md`. +mod relationships; + use std::sync::{Arc, LazyLock}; use api::v1::column_data_type_extension::TypeExt; @@ -33,21 +35,17 @@ use api::v1::{ SemanticType, }; use common_catalog::consts::{ - DEFAULT_PRIVATE_SCHEMA_NAME, DURATION_NANO_COLUMN, PARENT_SPAN_ID_COLUMN, - SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SPAN_ID_COLUMN, SPAN_KIND_CLIENT, SPAN_KIND_COLUMN, - SPAN_KIND_SERVER, SPAN_STATUS_CODE_COLUMN, SPAN_STATUS_ERROR, TRACE_ID_COLUMN, - TRACE_TIMESTAMP_COLUMN, + DEFAULT_PRIVATE_SCHEMA_NAME, SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, }; use common_function::function::FunctionContext; use common_function::function_registry::FUNCTION_REGISTRY; use datafusion::arrow::datatypes::{DataType, TimeUnit}; use datafusion::dataframe::DataFrame; use datafusion::functions::{core as core_fns, datetime as datetime_fns, string as string_fns}; -use datafusion::functions_aggregate::expr_fn::{count, sum}; use datafusion::functions_nested::expr_fn::make_array; -use datafusion::functions_window::expr_fn::row_number; use datafusion_common::{Column, Result as DfResult, ScalarValue}; -use datafusion_expr::{Expr, ExprFunctionExt, JoinType, LogicalPlan, ScalarUDF, cast, ident, lit}; +use datafusion_expr::{Expr, LogicalPlan, ScalarUDF, cast, ident, lit}; +pub use relationships::{CallsSource, DeclaredSource, build_relationships_plan}; use store_api::mito_engine_options::{APPEND_MODE_KEY, MERGE_MODE_KEY, TTL_KEY}; /// Whether an existing declared-edge table still matches the canonical @@ -619,19 +617,6 @@ pub struct RegistrySource { pub scan: DataFrame, } -/// A trace table's scan paired with the `service` declaration it derives -/// `calls` edges for — a unit of `build_relationships_plan`. -pub struct CallsSource { - pub service: EntityDeclaration, - pub scan: DataFrame, -} - -/// The declared-edge table's scan (`semantic_relationships_declared`), whose -/// rows `build_relationships_plan` unions into the edge set. -pub struct DeclaredSource { - pub scan: DataFrame, -} - /// Builds the `semantic_entities` registry plan: one branch and source scan per /// declaring table, filtered to `window`, then `UNION ALL` across source tables. /// Returns `None` when nothing declared an entity, so the computed table streams @@ -650,266 +635,71 @@ pub fn build_registry_plan( Ok(union_df.map(DataFrame::into_unoptimized_plan)) } -/// The projected columns of `semantic_relationships`, in order. Every derived -/// branch and the declared-edge branch must project exactly these so the -/// top-level `UNION ALL` type-aligns; `build_relationships_plan` re-selects -/// them over the union to enforce the contract. (The physical declared table -/// additionally stores `valid_from`/`valid_until`, which feed the validity -/// filter and the projected window columns.) -const RELATIONSHIP_COLUMNS: [&str; 16] = [ - "observed_at", - "window_start", - "window_end", - "fresh_until", - "src_type", - "src_id", - "dst_type", - "dst_id", - "rel_type", - "provenance", - "confidence", - "request_count", - "error_count", - "duration_sum", - "duration_count", - "attributes", -]; +/// RecordBatch readers shared by this module's tests and the `relationships` +/// tests. +#[cfg(test)] +pub(crate) mod test_util { + use datafusion::arrow::array::{Array, BinaryArray, StringArray, TimestampMillisecondArray}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::prelude::SessionContext; + use datafusion_expr::LogicalPlan; -/// A child server span starts no earlier than 5 minutes before its client span -/// (clock-skew allowance) and no later than 1 hour after it; the bounds keep the -/// join windowed instead of pairing arbitrarily distant spans of a long-lived -/// trace. -const CHILD_SPAN_EARLY_NANOS: i64 = 5 * 60 * 1_000_000_000; -const CHILD_SPAN_LATE_NANOS: i64 = 60 * 60 * 1_000_000_000; - -/// Builds the `semantic_relationships` plan: the trace-derived `calls` branch -/// unioned with the declared-edge branch, re-projected to the 16-column -/// contract. Returns `None` when there is neither a trace table nor a -/// declared-edge table, so the computed table streams empty. -pub fn build_relationships_plan( - traces: Vec, - declared: Option, - window: &GraphQueryWindow, -) -> DfResult> { - let mut union_df = calls_branch(traces, window)?; - if let Some(declared) = declared { - union_df = union_all(union_df, declared_edges(declared.scan, window)?)?; - } - union_df - .map(|df| { - Ok(df - .select_columns(&RELATIONSHIP_COLUMNS)? - .into_unoptimized_plan()) - }) - .transpose() -} - -/// Ranking column used to pick the latest declared-edge revision per edge key. -const DECLARED_REVISION_COLUMN: &str = "__declared_revision"; - -/// The declared-edge branch: the latest revision per edge key *as of the -/// queried window* (mito dedups on primary key + `observed_at`, so re-asserting -/// an edge stores a new revision), filtered by business-validity overlap. -/// -/// `valid_from` defaults to the declaration time; a NULL `valid_until` means -/// the edge holds while its row exists (TTL retires it), so `window_end` / -/// `fresh_until` take the window's upper bound. The output `observed_at` is -/// synthesized inside the queried range: the scan's filters are re-applied -/// above the computed table, and the physical revision time would fail them. -fn declared_edges(scan: DataFrame, window: &GraphQueryWindow) -> DfResult { - let eff_valid_from = - core_fns::coalesce().call(vec![ident("valid_from"), ident(OBSERVED_AT_COLUMN)]); - let eff_valid_until = - core_fns::coalesce().call(vec![ident("valid_until"), window.observed_end()]); - - // As-of the queried window: a revision recorded after its end, or whose - // validity starts after it, was not the edge's state inside the window and - // must not outrank (and thereby hide) the revision that was. - let as_of = ident(OBSERVED_AT_COLUMN) - .lt(window.observed_end()) - .and(eff_valid_from.clone().lt(window.observed_end())); - - let revision = row_number() - .partition_by( - DECLARED_EDGE_IDENTITY_COLUMNS - .iter() - .map(|c| ident(*c)) - .collect(), - ) - // generation_id/scope break observed_at ties deterministically (rows - // differing only in them are not merged by the storage dedup). - .order_by(vec![ - ident(OBSERVED_AT_COLUMN).sort(false, false), - ident("generation_id").sort(false, false), - ident("scope").sort(false, false), - ]) - .build()? - .alias(DECLARED_REVISION_COLUMN); - - // The as-of filter already bounds eff_valid_from below the window's end; - // only the retirement side of the overlap remains to check. - let still_valid = ident("valid_until") - .is_null() - .or(ident("valid_until").gt(window.observed_start())); - // Tags come out of the storage engine dictionary-encoded; cast to plain - // strings so the union with the derived branches type-aligns. - let tag_utf8 = |name: &str| cast(ident(name), DataType::Utf8).alias(name); - - scan.filter(as_of)? - .window(vec![revision])? - .filter(ident(DECLARED_REVISION_COLUMN).eq(lit(1_u64)))? - .filter(still_valid)? - .select(vec![ - core_fns::greatest() - .call(vec![eff_valid_from.clone(), window.observed_start()]) - .alias(OBSERVED_AT_COLUMN), - eff_valid_from.alias("window_start"), - eff_valid_until.clone().alias("window_end"), - eff_valid_until.alias("fresh_until"), - tag_utf8("src_type"), - tag_utf8("src_id"), - tag_utf8("dst_type"), - tag_utf8("dst_id"), - tag_utf8("rel_type"), - tag_utf8("provenance"), - ident("confidence"), - ident("request_count"), - ident("error_count"), - ident("duration_sum"), - ident("duration_count"), - ident("attributes"), - ]) -} - -/// The `calls` derivation: pair each client span with its child server span on -/// `trace_id` + `parent_span_id`, project to `service`, union the pairs of all -/// trace tables, and aggregate RED metrics per 60s window in one pass — an edge -/// observed across several trace tables yields one row, not per-table -/// fragments (the plan form of the Tempo servicegraph connector). Column names -/// are the fixed `greptime_trace_v1` schema, the reason -/// `table_data_model = greptime_trace_v1` is required. Returns `None` when -/// there is no trace table. -fn calls_branch( - traces: Vec, - window: &GraphQueryWindow, -) -> DfResult> { - let mut union_df: Option = None; - for trace in traces { - union_df = union_all(union_df, calls_pairs(&trace.service, trace.scan, window)?)?; - } - let Some(pairs) = union_df else { - return Ok(None); - }; - - let df = pairs - .aggregate( - vec![ident("observed_at"), ident("src_id"), ident("dst_id")], - vec![ - count(lit(1)).alias("request_count"), - count(lit(1)) - .filter(ident("status_code").eq(lit(SPAN_STATUS_ERROR))) - .build()? - .alias("error_count"), - sum(ident("duration_nano")).alias("duration_nano_sum"), - ], - )? - .select(vec![ - ident("observed_at"), - ident("observed_at").alias("window_start"), - (ident("observed_at") + bin_interval()).alias("window_end"), - (ident("observed_at") + bin_interval()).alias("fresh_until"), - lit("service").alias("src_type"), - ident("src_id"), - lit("service").alias("dst_type"), - ident("dst_id"), - lit("calls").alias("rel_type"), - lit("trace").alias("provenance"), - lit(1.0_f64).alias("confidence"), - ident("request_count"), - ident("error_count"), - // duration_nano sums in nanoseconds; the contract column is seconds. - (cast(ident("duration_nano_sum"), DataType::Float64) / lit(1e9_f64)) - .alias("duration_sum"), - ident("request_count").alias("duration_count"), - null_json().alias("attributes"), - ])?; - Ok(Some(df)) -} - -/// One trace table's client/server span pairs, projected to the aggregation -/// inputs `(observed_at, src_id, dst_id, status_code, duration_nano)`. -/// Endpoint ids are built from the table's `service` entity declaration, so -/// edges land on exactly the entity ids the registry emits (a composite -/// service identity renders the same sorted `k=v` form on both sides). -fn calls_pairs( - service: &EntityDeclaration, - trace: DataFrame, - window: &GraphQueryWindow, -) -> DfResult { - let mut client_pred = ident(SPAN_KIND_COLUMN) - .eq(lit(SPAN_KIND_CLIENT)) - .and(ident(TRACE_TIMESTAMP_COLUMN).gt_eq(window.source_start())) - .and(ident(TRACE_TIMESTAMP_COLUMN).lt(window.source_end())); - let mut server_pred = ident(SPAN_KIND_COLUMN) - .eq(lit(SPAN_KIND_SERVER)) - // Static bounds implied by the window and the join's time-proximity - // conditions below; the join bounds reference client.timestamp and - // cannot prune the server-side scan. - .and( - ident(TRACE_TIMESTAMP_COLUMN) - .gt_eq(window.source_start() - interval(CHILD_SPAN_EARLY_NANOS)), - ) - .and( - ident(TRACE_TIMESTAMP_COLUMN).lt(window.source_end() + interval(CHILD_SPAN_LATE_NANOS)), - ); - // A NULL identity component identifies nothing, on either endpoint. - for id in &service.id_columns { - client_pred = client_pred.and(ident(id).is_not_null()); - server_pred = server_pred.and(ident(id).is_not_null()); + pub(crate) async fn collect(ctx: &SessionContext, plan: LogicalPlan) -> Vec { + ctx.execute_logical_plan(plan) + .await + .unwrap() + .collect() + .await + .unwrap() } - let client = trace.clone().filter(client_pred)?.alias("client")?; - let server = trace.filter(server_pred)?.alias("server")?; + pub(crate) fn json_texts(batch: &RecordBatch, column: usize) -> Vec> { + let array = batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + (0..array.len()) + .map(|i| { + array + .is_valid(i) + .then(|| jsonb::from_slice(array.value(i)).unwrap().to_string()) + }) + .collect() + } - let join_conditions = vec![ - qcol("client", TRACE_ID_COLUMN).eq(qcol("server", TRACE_ID_COLUMN)), - qcol("server", PARENT_SPAN_ID_COLUMN).eq(qcol("client", SPAN_ID_COLUMN)), - qcol("server", TRACE_TIMESTAMP_COLUMN) - .gt_eq(qcol("client", TRACE_TIMESTAMP_COLUMN) - interval(CHILD_SPAN_EARLY_NANOS)), - qcol("server", TRACE_TIMESTAMP_COLUMN) - .lt_eq(qcol("client", TRACE_TIMESTAMP_COLUMN) + interval(CHILD_SPAN_LATE_NANOS)), - ]; + pub(crate) fn strings(batch: &RecordBatch, column: usize) -> Vec { + let array = batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + (0..array.len()) + .map(|i| array.value(i).to_string()) + .collect() + } - client - .join_on(server, JoinType::Inner, join_conditions)? - .select(vec![ - bin_ms(qcol("client", TRACE_TIMESTAMP_COLUMN)).alias("observed_at"), - // The cast inside entity_id_expr also normalizes tag columns, which - // come out of the storage engine dictionary-encoded. - entity_id_expr(&service.id_columns, &|c| qcol("client", c)).alias("src_id"), - entity_id_expr(&service.id_columns, &|c| qcol("server", c)).alias("dst_id"), - qcol("server", SPAN_STATUS_CODE_COLUMN).alias("status_code"), - qcol("server", DURATION_NANO_COLUMN).alias("duration_nano"), - ])? - // Exclude self-calls on the composed identity: an edge needs two - // distinct service entities. - .filter(ident("src_id").not_eq(ident("dst_id"))) + pub(crate) fn ts_values(batch: &RecordBatch, column: usize) -> Vec { + let array = batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + (0..array.len()).map(|i| array.value(i)).collect() + } } #[cfg(test)] mod tests { use std::sync::Arc; - use common_catalog::consts::SERVICE_NAME_COLUMN; - use datafusion::arrow::array::{ - Array, ArrayRef, BinaryArray, Float64Array, Int64Array, StringArray, - TimestampMillisecondArray, TimestampNanosecondArray, UInt64Array, - }; + use datafusion::arrow::array::{ArrayRef, Int64Array, StringArray, TimestampMillisecondArray}; use datafusion::arrow::datatypes::{Field, Schema}; use datafusion::arrow::record_batch::RecordBatch; use datafusion::datasource::MemTable; use datafusion::prelude::SessionContext; + use super::test_util::{collect, json_texts, strings}; use super::*; fn test_window() -> GraphQueryWindow { @@ -942,41 +732,6 @@ mod tests { assert_eq!((window.source_start_ms, window.source_end_ms), (-60_000, 0)); } - async fn collect(ctx: &SessionContext, plan: LogicalPlan) -> Vec { - ctx.execute_logical_plan(plan) - .await - .unwrap() - .collect() - .await - .unwrap() - } - - fn json_texts(batch: &RecordBatch, column: usize) -> Vec> { - let array = batch - .column(column) - .as_any() - .downcast_ref::() - .unwrap(); - (0..array.len()) - .map(|i| { - array - .is_valid(i) - .then(|| jsonb::from_slice(array.value(i)).unwrap().to_string()) - }) - .collect() - } - - fn strings(batch: &RecordBatch, column: usize) -> Vec { - let array = batch - .column(column) - .as_any() - .downcast_ref::() - .unwrap(); - (0..array.len()) - .map(|i| array.value(i).to_string()) - .collect() - } - /// A metric-like table: ms timestamps, service/pid identity, nullable /// descriptive column with JSON-hostile characters. fn metric_table_ctx() -> SessionContext { @@ -1234,258 +989,6 @@ mod tests { ); } - /// A trace-like table in the fixed `greptime_trace_v1` shape (ns timestamps). - fn trace_table_ctx() -> SessionContext { - let schema = Arc::new(Schema::new(vec![ - Field::new( - TRACE_TIMESTAMP_COLUMN, - DataType::Timestamp(TimeUnit::Nanosecond, None), - false, - ), - Field::new(TRACE_ID_COLUMN, DataType::Utf8, false), - Field::new(SPAN_ID_COLUMN, DataType::Utf8, false), - Field::new(PARENT_SPAN_ID_COLUMN, DataType::Utf8, true), - Field::new(SPAN_KIND_COLUMN, DataType::Utf8, false), - Field::new(SPAN_STATUS_CODE_COLUMN, DataType::Utf8, false), - Field::new(SERVICE_NAME_COLUMN, DataType::Utf8, false), - Field::new("service_namespace", DataType::Utf8, false), - Field::new(DURATION_NANO_COLUMN, DataType::UInt64, false), - ])); - const MS: i64 = 1_000_000; - // Two client->server pairs frontend->cart (one errored), one pair - // cart->cart (self-call, excluded), one unmatched client span. - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(TimestampNanosecondArray::from(vec![ - 1_000 * MS, // client frontend->cart - 1_010 * MS, // server cart - 2_000 * MS, // client frontend->cart (error) - 2_010 * MS, // server cart (error) - 3_000 * MS, // client cart->cart (self-call) - 3_010 * MS, // server cart - 4_000 * MS, // client with no matching server - ])) as ArrayRef, - Arc::new(StringArray::from(vec![ - "t1", "t1", "t2", "t2", "t3", "t3", "t4", - ])), - Arc::new(StringArray::from(vec![ - "c1", "s1", "c2", "s2", "c3", "s3", "c4", - ])), - Arc::new(StringArray::from(vec![ - None, - Some("c1"), - None, - Some("c2"), - None, - Some("c3"), - None, - ])), - Arc::new(StringArray::from(vec![ - "SPAN_KIND_CLIENT", - "SPAN_KIND_SERVER", - "SPAN_KIND_CLIENT", - "SPAN_KIND_SERVER", - "SPAN_KIND_CLIENT", - "SPAN_KIND_SERVER", - "SPAN_KIND_CLIENT", - ])), - Arc::new(StringArray::from(vec![ - "STATUS_CODE_UNSET", - "STATUS_CODE_UNSET", - "STATUS_CODE_UNSET", - "STATUS_CODE_ERROR", - "STATUS_CODE_UNSET", - "STATUS_CODE_UNSET", - "STATUS_CODE_UNSET", - ])), - Arc::new(StringArray::from(vec![ - "frontend", "cart", "frontend", "cart", "cart", "cart", "frontend", - ])), - Arc::new(StringArray::from(vec!["ns1"; 7])), - Arc::new(UInt64Array::from(vec![ - 0, - 500_000_000, // 0.5s - 0, - 1_500_000_000, // 1.5s - 0, - 100, - 0, - ])), - ], - ) - .unwrap(); - let ctx = SessionContext::new(); - ctx.register_table( - "opentelemetry_traces", - Arc::new(MemTable::try_new(schema.clone(), vec![vec![batch.clone()]]).unwrap()), - ) - .unwrap(); - ctx.register_table( - "opentelemetry_traces_2", - Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), - ) - .unwrap(); - ctx - } - - fn trace_service_decl(id_columns: &[&str]) -> EntityDeclaration { - EntityDeclaration { - schema: "public".to_string(), - table: "opentelemetry_traces".to_string(), - time_index: TRACE_TIMESTAMP_COLUMN.to_string(), - entity_type: "service".to_string(), - id_columns: id_columns.iter().map(|s| s.to_string()).collect(), - descriptive_columns: vec![], - scope_columns: vec![], - } - } - - #[tokio::test] - async fn calls_plan_aggregates_red_metrics() { - let ctx = trace_table_ctx(); - let trace = ctx.table("opentelemetry_traces").await.unwrap(); - let plan = build_relationships_plan_for_test( - vec![CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN]), - scan: trace, - }], - &test_window(), - ) - .unwrap() - .unwrap(); - - let names = plan - .schema() - .fields() - .iter() - .map(|f| f.name().as_str()) - .collect::>(); - assert_eq!(names, RELATIONSHIP_COLUMNS); - - let batches = collect(&ctx, plan).await; - let total: usize = batches.iter().map(|b| b.num_rows()).sum(); - // frontend->cart only: the self-call and the unmatched client drop out; - // both pairs land in the same 60s bin. - assert_eq!(total, 1); - let batch = &batches[0]; - assert_eq!(strings(batch, 5), vec!["frontend"]); - assert_eq!(strings(batch, 7), vec!["cart"]); - assert_eq!(strings(batch, 8), vec!["calls"]); - assert_eq!(strings(batch, 9), vec!["trace"]); - - let request_count = batch - .column(11) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(request_count.value(0), 2); - let error_count = batch - .column(12) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(error_count.value(0), 1); - let duration_sum = batch - .column(13) - .as_any() - .downcast_ref::() - .unwrap(); - assert!((duration_sum.value(0) - 2.0).abs() < 1e-9); - let duration_count = batch - .column(14) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(duration_count.value(0), 2); - // Derived calls edges carry no attributes: a typed-JSON NULL. - assert!(json_texts(batch, 15)[0].is_none()); - } - - #[tokio::test] - async fn calls_plan_merges_edges_across_trace_tables() { - let ctx = trace_table_ctx(); - let t1 = ctx.table("opentelemetry_traces").await.unwrap(); - let t2 = ctx.table("opentelemetry_traces_2").await.unwrap(); - let plan = build_relationships_plan_for_test( - vec![ - CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN]), - scan: t1, - }, - CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN]), - scan: t2, - }, - ], - &test_window(), - ) - .unwrap() - .unwrap(); - - let batches = collect(&ctx, plan).await; - let total: usize = batches.iter().map(|b| b.num_rows()).sum(); - // The same frontend->cart edge from both tables folds into one row with - // summed RED metrics. - assert_eq!(total, 1); - let batch = &batches[0]; - let request_count = batch - .column(11) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(request_count.value(0), 4); - let error_count = batch - .column(12) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(error_count.value(0), 2); - let duration_sum = batch - .column(13) - .as_any() - .downcast_ref::() - .unwrap(); - assert!((duration_sum.value(0) - 4.0).abs() < 1e-9); - } - - #[tokio::test] - async fn calls_endpoints_follow_service_declaration() { - let ctx = trace_table_ctx(); - let trace = ctx.table("opentelemetry_traces").await.unwrap(); - let plan = build_relationships_plan_for_test( - vec![CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN, "service_namespace"]), - scan: trace, - }], - &test_window(), - ) - .unwrap() - .unwrap(); - - let batches = collect(&ctx, plan).await; - let total: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(total, 1); - let batch = &batches[0]; - // Composite service identity renders the same sorted `k=v` form the - // registry emits, so edges land on registry entity ids. - assert_eq!( - strings(batch, 5), - vec!["service_name=frontend,service_namespace=ns1"] - ); - assert_eq!( - strings(batch, 7), - vec!["service_name=cart,service_namespace=ns1"] - ); - } - - fn build_relationships_plan_for_test( - traces: Vec, - window: &GraphQueryWindow, - ) -> DfResult> { - build_relationships_plan(traces, None, window) - } - /// A `TableInfo` derived from the canonical declared-edge definition, with /// injection points for each aspect the schema matcher must check. fn declared_table_info( @@ -1616,349 +1119,4 @@ mod tests { ); assert!(!declared_relationships_schema_matches(&wrong_column_type)); } - - fn ts_values(batch: &RecordBatch, column: usize) -> Vec { - let array = batch - .column(column) - .as_any() - .downcast_ref::() - .unwrap(); - (0..array.len()).map(|i| array.value(i)).collect() - } - - /// Registers a declared-edge table holding: - /// - frontend->db: two revisions (0.5 then 1.0), the newer retiring the edge at 300s; - /// - api->cache: expired at 100s; - /// - agent-1->search: open-ended `provenance = 'agent'` with JSON attributes; - /// - webapp->auth: two revisions, the newer declared at 900s; - /// - batch->queue: declared at 100s but only valid from 700s; - /// - gateway->redis: asserted twice at the same instant as generations g1/g2. - fn declared_table(ctx: &SessionContext) { - let ts_field = |name: &str| { - Field::new( - name, - DataType::Timestamp(TimeUnit::Millisecond, None), - name != OBSERVED_AT_COLUMN, - ) - }; - let schema = Arc::new(Schema::new(vec![ - ts_field(OBSERVED_AT_COLUMN), - ts_field("window_start"), - ts_field("window_end"), - ts_field("fresh_until"), - ts_field("valid_from"), - ts_field("valid_until"), - Field::new("src_type", DataType::Utf8, false), - Field::new("src_id", DataType::Utf8, false), - Field::new("rel_type", DataType::Utf8, false), - Field::new("dst_type", DataType::Utf8, false), - Field::new("dst_id", DataType::Utf8, false), - Field::new("provenance", DataType::Utf8, false), - Field::new("scope", DataType::Utf8, false), - Field::new("generation_id", DataType::Utf8, false), - Field::new("confidence", DataType::Float64, true), - Field::new("request_count", DataType::Int64, true), - Field::new("error_count", DataType::Int64, true), - Field::new("duration_sum", DataType::Float64, true), - Field::new("duration_count", DataType::Int64, true), - Field::new("attributes", DataType::Binary, true), - ])); - let attrs = jsonb::parse_value(br#"{"connection_type":"virtual_node"}"#) - .unwrap() - .to_vec(); - let no_ts = || Arc::new(TimestampMillisecondArray::from(vec![None::; 9])) as ArrayRef; - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(TimestampMillisecondArray::from(vec![ - 100_000, 200_000, 50_000, 90_000, 100_000, 900_000, 100_000, 100_000, 100_000, - ])) as ArrayRef, - no_ts(), // window_start - no_ts(), // window_end - no_ts(), // fresh_until - Arc::new(TimestampMillisecondArray::from(vec![ - None, - None, - None, - None, - None, - None, - Some(700_000), - None, - None, - ])), - Arc::new(TimestampMillisecondArray::from(vec![ - None, - Some(300_000), - Some(100_000), - None, - None, - None, - None, - None, - None, - ])), - Arc::new(StringArray::from(vec![ - "service", "service", "service", "agent", "service", "service", "service", - "service", "service", - ])), - Arc::new(StringArray::from(vec![ - "frontend", "frontend", "api", "agent-1", "webapp", "webapp", "batch", - "gateway", "gateway", - ])), - Arc::new(StringArray::from(vec![ - "depends_on", - "depends_on", - "depends_on", - "uses", - "depends_on", - "depends_on", - "depends_on", - "depends_on", - "depends_on", - ])), - Arc::new(StringArray::from(vec![ - "service", "service", "service", "tool", "service", "service", "service", - "service", "service", - ])), - Arc::new(StringArray::from(vec![ - "db", "db", "cache", "search", "auth", "auth", "queue", "redis", "redis", - ])), - Arc::new(StringArray::from(vec![ - "declared", "declared", "declared", "agent", "declared", "declared", - "declared", "declared", "declared", - ])), - Arc::new(StringArray::from(vec![""; 9])), - Arc::new(StringArray::from(vec![ - "", "", "", "", "", "", "", "g1", "g2", - ])), - Arc::new(Float64Array::from(vec![ - Some(0.5), - Some(1.0), - Some(1.0), - Some(0.8), - Some(0.5), - Some(1.0), - Some(1.0), - Some(0.5), - Some(1.0), - ])), - Arc::new(Int64Array::from(vec![None::; 9])), - Arc::new(Int64Array::from(vec![None::; 9])), - Arc::new(Float64Array::from(vec![None::; 9])), - Arc::new(Int64Array::from(vec![None::; 9])), - Arc::new(BinaryArray::from(vec![ - None, - None, - None, - Some(attrs.as_slice()), - None, - None, - None, - None, - None, - ])), - ], - ) - .unwrap(); - ctx.register_table( - SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, - Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), - ) - .unwrap(); - } - - #[tokio::test] - async fn declared_edges_latest_revision_validity_and_synthetic_time() { - let ctx = SessionContext::new(); - declared_table(&ctx); - let scan = ctx - .table(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) - .await - .unwrap(); - - // Queried window [120s, 600s). - let window = GraphQueryWindow::from_observed(120_000, 600_000); - let plan = build_relationships_plan(vec![], Some(DeclaredSource { scan }), &window) - .unwrap() - .unwrap(); - let names = plan - .schema() - .fields() - .iter() - .map(|f| f.name().as_str()) - .collect::>(); - assert_eq!(names, RELATIONSHIP_COLUMNS); - - // (src_id, observed_at, window_start, window_end, fresh_until, - // provenance, attributes) - type DeclaredRow = (String, i64, i64, i64, i64, String, Option); - - let batches = collect(&ctx, plan).await; - let mut rows: Vec = vec![]; - for batch in &batches { - let observed = ts_values(batch, 0); - let window_start = ts_values(batch, 1); - let window_end = ts_values(batch, 2); - let fresh_until = ts_values(batch, 3); - let src = strings(batch, 5); - let provenance = strings(batch, 9); - let attributes = json_texts(batch, 15); - for i in 0..batch.num_rows() { - rows.push(( - src[i].clone(), - observed[i], - window_start[i], - window_end[i], - fresh_until[i], - provenance[i].clone(), - attributes[i].clone(), - )); - } - } - rows.sort(); - - // api->cache expired before the window; batch->queue is not yet valid; - // frontend->db keeps only the latest revision; webapp->auth keeps the - // first revision (the second postdates the window); gateway->redis - // collapses to one row across generations. - assert_eq!(rows.len(), 4); - - // agent->tool: open-ended validity declared at 90s. The synthesized - // observed_at is clamped into the queried window; window_end and - // fresh_until take the window's upper bound. - assert_eq!( - rows[0], - ( - "agent-1".to_string(), - 120_000, - 90_000, - 600_000, - 600_000, - "agent".to_string(), - Some(r#"{"connection_type":"virtual_node"}"#.to_string()), - ) - ); - - // frontend->db: the latest revision (declared at 200s, retired at 300s) - // supersedes the open-ended first revision. - assert_eq!( - rows[1], - ( - "frontend".to_string(), - 200_000, - 200_000, - 300_000, - 300_000, - "declared".to_string(), - None, - ) - ); - - assert_eq!( - rows[2], - ( - "gateway".to_string(), - 120_000, - 100_000, - 600_000, - 600_000, - "declared".to_string(), - None, - ) - ); - - assert_eq!( - rows[3], - ( - "webapp".to_string(), - 120_000, - 100_000, - 600_000, - 600_000, - "declared".to_string(), - None, - ) - ); - } - - #[tokio::test] - async fn declared_edges_pick_the_revision_as_of_the_window() { - let ctx = SessionContext::new(); - declared_table(&ctx); - let scan = ctx - .table(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) - .await - .unwrap(); - - // A historical window that predates frontend->db's second revision: - // the first revision was the edge's state then and must not be hidden - // by the newer one. - let window = GraphQueryWindow::from_observed(0, 150_000); - let plan = build_relationships_plan(vec![], Some(DeclaredSource { scan }), &window) - .unwrap() - .unwrap(); - - let batches = collect(&ctx, plan).await; - let mut rows: Vec<(String, f64)> = vec![]; - for batch in &batches { - let src = strings(batch, 5); - let confidence = batch - .column(10) - .as_any() - .downcast_ref::() - .unwrap(); - for (i, src) in src.into_iter().enumerate() { - rows.push((src, confidence.value(i))); - } - } - rows.sort_by(|a, b| a.0.cmp(&b.0)); - - // api->cache is valid inside this window; batch->queue is not yet. - assert_eq!( - rows, - vec![ - ("agent-1".to_string(), 0.8), - ("api".to_string(), 1.0), - ("frontend".to_string(), 0.5), - ("gateway".to_string(), 1.0), - ("webapp".to_string(), 0.5), - ] - ); - } - - #[tokio::test] - async fn declared_edges_union_calls_branch() { - let ctx = trace_table_ctx(); - declared_table(&ctx); - let trace = ctx.table("opentelemetry_traces").await.unwrap(); - let declared = ctx - .table(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) - .await - .unwrap(); - - let plan = build_relationships_plan( - vec![CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN]), - scan: trace, - }], - Some(DeclaredSource { scan: declared }), - &test_window(), - ) - .unwrap() - .unwrap(); - - let batches = collect(&ctx, plan).await; - let mut provenance: Vec = batches.iter().flat_map(|b| strings(b, 9)).collect(); - provenance.sort(); - // One trace-derived calls edge plus the declared edges valid somewhere - // inside [0s, 600s): frontend, api, webapp, gateway, and the - // agent-asserted one. - assert_eq!( - provenance, - vec![ - "agent", "declared", "declared", "declared", "declared", "trace" - ] - ); - } } diff --git a/src/operator/src/statement/semantic_graph/relationships.rs b/src/operator/src/statement/semantic_graph/relationships.rs new file mode 100644 index 00000000000..7ebf5ae86fd --- /dev/null +++ b/src/operator/src/statement/semantic_graph/relationships.rs @@ -0,0 +1,905 @@ +// Copyright 2023 Greptime Team +// +// 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. + +//! The `semantic_relationships` derivation: the plan builders for every edge +//! branch (trace-derived `calls`, the declared-edge union) behind the computed +//! table. Shared expression helpers and the entity registry live in the parent +//! module. + +use common_catalog::consts::{ + DURATION_NANO_COLUMN, PARENT_SPAN_ID_COLUMN, SPAN_ID_COLUMN, SPAN_KIND_CLIENT, + SPAN_KIND_COLUMN, SPAN_KIND_SERVER, SPAN_STATUS_CODE_COLUMN, SPAN_STATUS_ERROR, + TRACE_ID_COLUMN, TRACE_TIMESTAMP_COLUMN, +}; +use datafusion::arrow::datatypes::DataType; +use datafusion::dataframe::DataFrame; +use datafusion::functions::core as core_fns; +use datafusion::functions_aggregate::expr_fn::{count, sum}; +use datafusion::functions_window::expr_fn::row_number; +use datafusion_common::Result as DfResult; +use datafusion_expr::{ExprFunctionExt, JoinType, LogicalPlan, cast, ident, lit}; + +use crate::statement::semantic_graph::{ + DECLARED_EDGE_IDENTITY_COLUMNS, EntityDeclaration, GraphQueryWindow, OBSERVED_AT_COLUMN, + bin_interval, bin_ms, entity_id_expr, interval, null_json, qcol, union_all, +}; + +/// The projected columns of `semantic_relationships`, in order. Every derived +/// branch and the declared-edge branch must project exactly these so the +/// top-level `UNION ALL` type-aligns; `build_relationships_plan` re-selects +/// them over the union to enforce the contract. (The physical declared table +/// additionally stores `valid_from`/`valid_until`, which feed the validity +/// filter and the projected window columns.) +const RELATIONSHIP_COLUMNS: [&str; 16] = [ + "observed_at", + "window_start", + "window_end", + "fresh_until", + "src_type", + "src_id", + "dst_type", + "dst_id", + "rel_type", + "provenance", + "confidence", + "request_count", + "error_count", + "duration_sum", + "duration_count", + "attributes", +]; + +/// A child server span starts no earlier than 5 minutes before its client span +/// (clock-skew allowance) and no later than 1 hour after it; the bounds keep the +/// join windowed instead of pairing arbitrarily distant spans of a long-lived +/// trace. +const CHILD_SPAN_EARLY_NANOS: i64 = 5 * 60 * 1_000_000_000; +const CHILD_SPAN_LATE_NANOS: i64 = 60 * 60 * 1_000_000_000; + +/// A trace table's scan paired with the `service` declaration it derives +/// `calls` edges for — a unit of `build_relationships_plan`. +pub struct CallsSource { + pub service: EntityDeclaration, + pub scan: DataFrame, +} + +/// The declared-edge table's scan (`semantic_relationships_declared`), whose +/// rows `build_relationships_plan` unions into the edge set. +pub struct DeclaredSource { + pub scan: DataFrame, +} + +/// Builds the `semantic_relationships` plan: the trace-derived `calls` branch +/// unioned with the declared-edge branch, re-projected to the 16-column +/// contract. Returns `None` when there is neither a trace table nor a +/// declared-edge table, so the computed table streams empty. +pub fn build_relationships_plan( + traces: Vec, + declared: Option, + window: &GraphQueryWindow, +) -> DfResult> { + let mut union_df = calls_branch(traces, window)?; + if let Some(declared) = declared { + union_df = union_all(union_df, declared_edges(declared.scan, window)?)?; + } + union_df + .map(|df| { + Ok(df + .select_columns(&RELATIONSHIP_COLUMNS)? + .into_unoptimized_plan()) + }) + .transpose() +} + +/// Ranking column used to pick the latest declared-edge revision per edge key. +const DECLARED_REVISION_COLUMN: &str = "__declared_revision"; + +/// The declared-edge branch: the latest revision per edge key *as of the +/// queried window* (mito dedups on primary key + `observed_at`, so re-asserting +/// an edge stores a new revision), filtered by business-validity overlap. +/// +/// `valid_from` defaults to the declaration time; a NULL `valid_until` means +/// the edge holds while its row exists (TTL retires it), so `window_end` / +/// `fresh_until` take the window's upper bound. The output `observed_at` is +/// synthesized inside the queried range: the scan's filters are re-applied +/// above the computed table, and the physical revision time would fail them. +fn declared_edges(scan: DataFrame, window: &GraphQueryWindow) -> DfResult { + let eff_valid_from = + core_fns::coalesce().call(vec![ident("valid_from"), ident(OBSERVED_AT_COLUMN)]); + let eff_valid_until = + core_fns::coalesce().call(vec![ident("valid_until"), window.observed_end()]); + + // As-of the queried window: a revision recorded after its end, or whose + // validity starts after it, was not the edge's state inside the window and + // must not outrank (and thereby hide) the revision that was. + let as_of = ident(OBSERVED_AT_COLUMN) + .lt(window.observed_end()) + .and(eff_valid_from.clone().lt(window.observed_end())); + + let revision = row_number() + .partition_by( + DECLARED_EDGE_IDENTITY_COLUMNS + .iter() + .map(|c| ident(*c)) + .collect(), + ) + // generation_id/scope break observed_at ties deterministically (rows + // differing only in them are not merged by the storage dedup). + .order_by(vec![ + ident(OBSERVED_AT_COLUMN).sort(false, false), + ident("generation_id").sort(false, false), + ident("scope").sort(false, false), + ]) + .build()? + .alias(DECLARED_REVISION_COLUMN); + + // The as-of filter already bounds eff_valid_from below the window's end; + // only the retirement side of the overlap remains to check. + let still_valid = ident("valid_until") + .is_null() + .or(ident("valid_until").gt(window.observed_start())); + // Tags come out of the storage engine dictionary-encoded; cast to plain + // strings so the union with the derived branches type-aligns. + let tag_utf8 = |name: &str| cast(ident(name), DataType::Utf8).alias(name); + + scan.filter(as_of)? + .window(vec![revision])? + .filter(ident(DECLARED_REVISION_COLUMN).eq(lit(1_u64)))? + .filter(still_valid)? + .select(vec![ + core_fns::greatest() + .call(vec![eff_valid_from.clone(), window.observed_start()]) + .alias(OBSERVED_AT_COLUMN), + eff_valid_from.alias("window_start"), + eff_valid_until.clone().alias("window_end"), + eff_valid_until.alias("fresh_until"), + tag_utf8("src_type"), + tag_utf8("src_id"), + tag_utf8("dst_type"), + tag_utf8("dst_id"), + tag_utf8("rel_type"), + tag_utf8("provenance"), + ident("confidence"), + ident("request_count"), + ident("error_count"), + ident("duration_sum"), + ident("duration_count"), + ident("attributes"), + ]) +} + +/// The `calls` derivation: pair each client span with its child server span on +/// `trace_id` + `parent_span_id`, project to `service`, union the pairs of all +/// trace tables, and aggregate RED metrics per 60s window in one pass — an edge +/// observed across several trace tables yields one row, not per-table +/// fragments (the plan form of the Tempo servicegraph connector). Column names +/// are the fixed `greptime_trace_v1` schema, the reason +/// `table_data_model = greptime_trace_v1` is required. Returns `None` when +/// there is no trace table. +fn calls_branch( + traces: Vec, + window: &GraphQueryWindow, +) -> DfResult> { + let mut union_df: Option = None; + for trace in traces { + union_df = union_all(union_df, calls_pairs(&trace.service, trace.scan, window)?)?; + } + let Some(pairs) = union_df else { + return Ok(None); + }; + + let df = pairs + .aggregate( + vec![ident("observed_at"), ident("src_id"), ident("dst_id")], + vec![ + count(lit(1)).alias("request_count"), + count(lit(1)) + .filter(ident("status_code").eq(lit(SPAN_STATUS_ERROR))) + .build()? + .alias("error_count"), + sum(ident("duration_nano")).alias("duration_nano_sum"), + ], + )? + .select(vec![ + ident("observed_at"), + ident("observed_at").alias("window_start"), + (ident("observed_at") + bin_interval()).alias("window_end"), + (ident("observed_at") + bin_interval()).alias("fresh_until"), + lit("service").alias("src_type"), + ident("src_id"), + lit("service").alias("dst_type"), + ident("dst_id"), + lit("calls").alias("rel_type"), + lit("trace").alias("provenance"), + lit(1.0_f64).alias("confidence"), + ident("request_count"), + ident("error_count"), + // duration_nano sums in nanoseconds; the contract column is seconds. + (cast(ident("duration_nano_sum"), DataType::Float64) / lit(1e9_f64)) + .alias("duration_sum"), + ident("request_count").alias("duration_count"), + null_json().alias("attributes"), + ])?; + Ok(Some(df)) +} + +/// One trace table's client/server span pairs, projected to the aggregation +/// inputs `(observed_at, src_id, dst_id, status_code, duration_nano)`. +/// Endpoint ids are built from the table's `service` entity declaration, so +/// edges land on exactly the entity ids the registry emits (a composite +/// service identity renders the same sorted `k=v` form on both sides). +fn calls_pairs( + service: &EntityDeclaration, + trace: DataFrame, + window: &GraphQueryWindow, +) -> DfResult { + let mut client_pred = ident(SPAN_KIND_COLUMN) + .eq(lit(SPAN_KIND_CLIENT)) + .and(ident(TRACE_TIMESTAMP_COLUMN).gt_eq(window.source_start())) + .and(ident(TRACE_TIMESTAMP_COLUMN).lt(window.source_end())); + let mut server_pred = ident(SPAN_KIND_COLUMN) + .eq(lit(SPAN_KIND_SERVER)) + // Static bounds implied by the window and the join's time-proximity + // conditions below; the join bounds reference client.timestamp and + // cannot prune the server-side scan. + .and( + ident(TRACE_TIMESTAMP_COLUMN) + .gt_eq(window.source_start() - interval(CHILD_SPAN_EARLY_NANOS)), + ) + .and( + ident(TRACE_TIMESTAMP_COLUMN).lt(window.source_end() + interval(CHILD_SPAN_LATE_NANOS)), + ); + // A NULL identity component identifies nothing, on either endpoint. + for id in &service.id_columns { + client_pred = client_pred.and(ident(id).is_not_null()); + server_pred = server_pred.and(ident(id).is_not_null()); + } + + let client = trace.clone().filter(client_pred)?.alias("client")?; + let server = trace.filter(server_pred)?.alias("server")?; + + let join_conditions = vec![ + qcol("client", TRACE_ID_COLUMN).eq(qcol("server", TRACE_ID_COLUMN)), + qcol("server", PARENT_SPAN_ID_COLUMN).eq(qcol("client", SPAN_ID_COLUMN)), + qcol("server", TRACE_TIMESTAMP_COLUMN) + .gt_eq(qcol("client", TRACE_TIMESTAMP_COLUMN) - interval(CHILD_SPAN_EARLY_NANOS)), + qcol("server", TRACE_TIMESTAMP_COLUMN) + .lt_eq(qcol("client", TRACE_TIMESTAMP_COLUMN) + interval(CHILD_SPAN_LATE_NANOS)), + ]; + + client + .join_on(server, JoinType::Inner, join_conditions)? + .select(vec![ + bin_ms(qcol("client", TRACE_TIMESTAMP_COLUMN)).alias("observed_at"), + // The cast inside entity_id_expr also normalizes tag columns, which + // come out of the storage engine dictionary-encoded. + entity_id_expr(&service.id_columns, &|c| qcol("client", c)).alias("src_id"), + entity_id_expr(&service.id_columns, &|c| qcol("server", c)).alias("dst_id"), + qcol("server", SPAN_STATUS_CODE_COLUMN).alias("status_code"), + qcol("server", DURATION_NANO_COLUMN).alias("duration_nano"), + ])? + // Exclude self-calls on the composed identity: an edge needs two + // distinct service entities. + .filter(ident("src_id").not_eq(ident("dst_id"))) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_catalog::consts::{SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SERVICE_NAME_COLUMN}; + use datafusion::arrow::array::{ + ArrayRef, BinaryArray, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, + TimestampNanosecondArray, UInt64Array, + }; + use datafusion::arrow::datatypes::{Field, Schema, TimeUnit}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::datasource::MemTable; + use datafusion::prelude::SessionContext; + + use super::super::test_util::{collect, json_texts, strings, ts_values}; + use super::*; + + fn test_window() -> GraphQueryWindow { + GraphQueryWindow::from_observed(0, 10 * 60 * 1000) + } + + /// A trace-like table in the fixed `greptime_trace_v1` shape (ns timestamps). + fn trace_table_ctx() -> SessionContext { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TRACE_TIMESTAMP_COLUMN, + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ), + Field::new(TRACE_ID_COLUMN, DataType::Utf8, false), + Field::new(SPAN_ID_COLUMN, DataType::Utf8, false), + Field::new(PARENT_SPAN_ID_COLUMN, DataType::Utf8, true), + Field::new(SPAN_KIND_COLUMN, DataType::Utf8, false), + Field::new(SPAN_STATUS_CODE_COLUMN, DataType::Utf8, false), + Field::new(SERVICE_NAME_COLUMN, DataType::Utf8, false), + Field::new("service_namespace", DataType::Utf8, false), + Field::new(DURATION_NANO_COLUMN, DataType::UInt64, false), + ])); + const MS: i64 = 1_000_000; + // Two client->server pairs frontend->cart (one errored), one pair + // cart->cart (self-call, excluded), one unmatched client span. + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampNanosecondArray::from(vec![ + 1_000 * MS, // client frontend->cart + 1_010 * MS, // server cart + 2_000 * MS, // client frontend->cart (error) + 2_010 * MS, // server cart (error) + 3_000 * MS, // client cart->cart (self-call) + 3_010 * MS, // server cart + 4_000 * MS, // client with no matching server + ])) as ArrayRef, + Arc::new(StringArray::from(vec![ + "t1", "t1", "t2", "t2", "t3", "t3", "t4", + ])), + Arc::new(StringArray::from(vec![ + "c1", "s1", "c2", "s2", "c3", "s3", "c4", + ])), + Arc::new(StringArray::from(vec![ + None, + Some("c1"), + None, + Some("c2"), + None, + Some("c3"), + None, + ])), + Arc::new(StringArray::from(vec![ + "SPAN_KIND_CLIENT", + "SPAN_KIND_SERVER", + "SPAN_KIND_CLIENT", + "SPAN_KIND_SERVER", + "SPAN_KIND_CLIENT", + "SPAN_KIND_SERVER", + "SPAN_KIND_CLIENT", + ])), + Arc::new(StringArray::from(vec![ + "STATUS_CODE_UNSET", + "STATUS_CODE_UNSET", + "STATUS_CODE_UNSET", + "STATUS_CODE_ERROR", + "STATUS_CODE_UNSET", + "STATUS_CODE_UNSET", + "STATUS_CODE_UNSET", + ])), + Arc::new(StringArray::from(vec![ + "frontend", "cart", "frontend", "cart", "cart", "cart", "frontend", + ])), + Arc::new(StringArray::from(vec!["ns1"; 7])), + Arc::new(UInt64Array::from(vec![ + 0, + 500_000_000, // 0.5s + 0, + 1_500_000_000, // 1.5s + 0, + 100, + 0, + ])), + ], + ) + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_table( + "opentelemetry_traces", + Arc::new(MemTable::try_new(schema.clone(), vec![vec![batch.clone()]]).unwrap()), + ) + .unwrap(); + ctx.register_table( + "opentelemetry_traces_2", + Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), + ) + .unwrap(); + ctx + } + + fn trace_service_decl(id_columns: &[&str]) -> EntityDeclaration { + EntityDeclaration { + schema: "public".to_string(), + table: "opentelemetry_traces".to_string(), + time_index: TRACE_TIMESTAMP_COLUMN.to_string(), + entity_type: "service".to_string(), + id_columns: id_columns.iter().map(|s| s.to_string()).collect(), + descriptive_columns: vec![], + scope_columns: vec![], + } + } + + #[tokio::test] + async fn calls_plan_aggregates_red_metrics() { + let ctx = trace_table_ctx(); + let trace = ctx.table("opentelemetry_traces").await.unwrap(); + let plan = build_relationships_plan_for_test( + vec![CallsSource { + service: trace_service_decl(&[SERVICE_NAME_COLUMN]), + scan: trace, + }], + &test_window(), + ) + .unwrap() + .unwrap(); + + let names = plan + .schema() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>(); + assert_eq!(names, RELATIONSHIP_COLUMNS); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + // frontend->cart only: the self-call and the unmatched client drop out; + // both pairs land in the same 60s bin. + assert_eq!(total, 1); + let batch = &batches[0]; + assert_eq!(strings(batch, 5), vec!["frontend"]); + assert_eq!(strings(batch, 7), vec!["cart"]); + assert_eq!(strings(batch, 8), vec!["calls"]); + assert_eq!(strings(batch, 9), vec!["trace"]); + + let request_count = batch + .column(11) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(request_count.value(0), 2); + let error_count = batch + .column(12) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(error_count.value(0), 1); + let duration_sum = batch + .column(13) + .as_any() + .downcast_ref::() + .unwrap(); + assert!((duration_sum.value(0) - 2.0).abs() < 1e-9); + let duration_count = batch + .column(14) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(duration_count.value(0), 2); + // Derived calls edges carry no attributes: a typed-JSON NULL. + assert!(json_texts(batch, 15)[0].is_none()); + } + + #[tokio::test] + async fn calls_plan_merges_edges_across_trace_tables() { + let ctx = trace_table_ctx(); + let t1 = ctx.table("opentelemetry_traces").await.unwrap(); + let t2 = ctx.table("opentelemetry_traces_2").await.unwrap(); + let plan = build_relationships_plan_for_test( + vec![ + CallsSource { + service: trace_service_decl(&[SERVICE_NAME_COLUMN]), + scan: t1, + }, + CallsSource { + service: trace_service_decl(&[SERVICE_NAME_COLUMN]), + scan: t2, + }, + ], + &test_window(), + ) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + // The same frontend->cart edge from both tables folds into one row with + // summed RED metrics. + assert_eq!(total, 1); + let batch = &batches[0]; + let request_count = batch + .column(11) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(request_count.value(0), 4); + let error_count = batch + .column(12) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(error_count.value(0), 2); + let duration_sum = batch + .column(13) + .as_any() + .downcast_ref::() + .unwrap(); + assert!((duration_sum.value(0) - 4.0).abs() < 1e-9); + } + + #[tokio::test] + async fn calls_endpoints_follow_service_declaration() { + let ctx = trace_table_ctx(); + let trace = ctx.table("opentelemetry_traces").await.unwrap(); + let plan = build_relationships_plan_for_test( + vec![CallsSource { + service: trace_service_decl(&[SERVICE_NAME_COLUMN, "service_namespace"]), + scan: trace, + }], + &test_window(), + ) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, 1); + let batch = &batches[0]; + // Composite service identity renders the same sorted `k=v` form the + // registry emits, so edges land on registry entity ids. + assert_eq!( + strings(batch, 5), + vec!["service_name=frontend,service_namespace=ns1"] + ); + assert_eq!( + strings(batch, 7), + vec!["service_name=cart,service_namespace=ns1"] + ); + } + + fn build_relationships_plan_for_test( + traces: Vec, + window: &GraphQueryWindow, + ) -> DfResult> { + build_relationships_plan(traces, None, window) + } + + /// Registers a declared-edge table holding: + /// - frontend->db: two revisions (0.5 then 1.0), the newer retiring the edge at 300s; + /// - api->cache: expired at 100s; + /// - agent-1->search: open-ended `provenance = 'agent'` with JSON attributes; + /// - webapp->auth: two revisions, the newer declared at 900s; + /// - batch->queue: declared at 100s but only valid from 700s; + /// - gateway->redis: asserted twice at the same instant as generations g1/g2. + fn declared_table(ctx: &SessionContext) { + let ts_field = |name: &str| { + Field::new( + name, + DataType::Timestamp(TimeUnit::Millisecond, None), + name != OBSERVED_AT_COLUMN, + ) + }; + let schema = Arc::new(Schema::new(vec![ + ts_field(OBSERVED_AT_COLUMN), + ts_field("window_start"), + ts_field("window_end"), + ts_field("fresh_until"), + ts_field("valid_from"), + ts_field("valid_until"), + Field::new("src_type", DataType::Utf8, false), + Field::new("src_id", DataType::Utf8, false), + Field::new("rel_type", DataType::Utf8, false), + Field::new("dst_type", DataType::Utf8, false), + Field::new("dst_id", DataType::Utf8, false), + Field::new("provenance", DataType::Utf8, false), + Field::new("scope", DataType::Utf8, false), + Field::new("generation_id", DataType::Utf8, false), + Field::new("confidence", DataType::Float64, true), + Field::new("request_count", DataType::Int64, true), + Field::new("error_count", DataType::Int64, true), + Field::new("duration_sum", DataType::Float64, true), + Field::new("duration_count", DataType::Int64, true), + Field::new("attributes", DataType::Binary, true), + ])); + let attrs = jsonb::parse_value(br#"{"connection_type":"virtual_node"}"#) + .unwrap() + .to_vec(); + let no_ts = || Arc::new(TimestampMillisecondArray::from(vec![None::; 9])) as ArrayRef; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![ + 100_000, 200_000, 50_000, 90_000, 100_000, 900_000, 100_000, 100_000, 100_000, + ])) as ArrayRef, + no_ts(), // window_start + no_ts(), // window_end + no_ts(), // fresh_until + Arc::new(TimestampMillisecondArray::from(vec![ + None, + None, + None, + None, + None, + None, + Some(700_000), + None, + None, + ])), + Arc::new(TimestampMillisecondArray::from(vec![ + None, + Some(300_000), + Some(100_000), + None, + None, + None, + None, + None, + None, + ])), + Arc::new(StringArray::from(vec![ + "service", "service", "service", "agent", "service", "service", "service", + "service", "service", + ])), + Arc::new(StringArray::from(vec![ + "frontend", "frontend", "api", "agent-1", "webapp", "webapp", "batch", + "gateway", "gateway", + ])), + Arc::new(StringArray::from(vec![ + "depends_on", + "depends_on", + "depends_on", + "uses", + "depends_on", + "depends_on", + "depends_on", + "depends_on", + "depends_on", + ])), + Arc::new(StringArray::from(vec![ + "service", "service", "service", "tool", "service", "service", "service", + "service", "service", + ])), + Arc::new(StringArray::from(vec![ + "db", "db", "cache", "search", "auth", "auth", "queue", "redis", "redis", + ])), + Arc::new(StringArray::from(vec![ + "declared", "declared", "declared", "agent", "declared", "declared", + "declared", "declared", "declared", + ])), + Arc::new(StringArray::from(vec![""; 9])), + Arc::new(StringArray::from(vec![ + "", "", "", "", "", "", "", "g1", "g2", + ])), + Arc::new(Float64Array::from(vec![ + Some(0.5), + Some(1.0), + Some(1.0), + Some(0.8), + Some(0.5), + Some(1.0), + Some(1.0), + Some(0.5), + Some(1.0), + ])), + Arc::new(Int64Array::from(vec![None::; 9])), + Arc::new(Int64Array::from(vec![None::; 9])), + Arc::new(Float64Array::from(vec![None::; 9])), + Arc::new(Int64Array::from(vec![None::; 9])), + Arc::new(BinaryArray::from(vec![ + None, + None, + None, + Some(attrs.as_slice()), + None, + None, + None, + None, + None, + ])), + ], + ) + .unwrap(); + ctx.register_table( + SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, + Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), + ) + .unwrap(); + } + + #[tokio::test] + async fn declared_edges_latest_revision_validity_and_synthetic_time() { + let ctx = SessionContext::new(); + declared_table(&ctx); + let scan = ctx + .table(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) + .await + .unwrap(); + + // Queried window [120s, 600s). + let window = GraphQueryWindow::from_observed(120_000, 600_000); + let plan = build_relationships_plan(vec![], Some(DeclaredSource { scan }), &window) + .unwrap() + .unwrap(); + let names = plan + .schema() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>(); + assert_eq!(names, RELATIONSHIP_COLUMNS); + + // (src_id, observed_at, window_start, window_end, fresh_until, + // provenance, attributes) + type DeclaredRow = (String, i64, i64, i64, i64, String, Option); + + let batches = collect(&ctx, plan).await; + let mut rows: Vec = vec![]; + for batch in &batches { + let observed = ts_values(batch, 0); + let window_start = ts_values(batch, 1); + let window_end = ts_values(batch, 2); + let fresh_until = ts_values(batch, 3); + let src = strings(batch, 5); + let provenance = strings(batch, 9); + let attributes = json_texts(batch, 15); + for i in 0..batch.num_rows() { + rows.push(( + src[i].clone(), + observed[i], + window_start[i], + window_end[i], + fresh_until[i], + provenance[i].clone(), + attributes[i].clone(), + )); + } + } + rows.sort(); + + // api->cache expired before the window; batch->queue is not yet valid; + // frontend->db keeps only the latest revision; webapp->auth keeps the + // first revision (the second postdates the window); gateway->redis + // collapses to one row across generations. + assert_eq!(rows.len(), 4); + + // agent->tool: open-ended validity declared at 90s. The synthesized + // observed_at is clamped into the queried window; window_end and + // fresh_until take the window's upper bound. + assert_eq!( + rows[0], + ( + "agent-1".to_string(), + 120_000, + 90_000, + 600_000, + 600_000, + "agent".to_string(), + Some(r#"{"connection_type":"virtual_node"}"#.to_string()), + ) + ); + + // frontend->db: the latest revision (declared at 200s, retired at 300s) + // supersedes the open-ended first revision. + assert_eq!( + rows[1], + ( + "frontend".to_string(), + 200_000, + 200_000, + 300_000, + 300_000, + "declared".to_string(), + None, + ) + ); + + assert_eq!( + rows[2], + ( + "gateway".to_string(), + 120_000, + 100_000, + 600_000, + 600_000, + "declared".to_string(), + None, + ) + ); + + assert_eq!( + rows[3], + ( + "webapp".to_string(), + 120_000, + 100_000, + 600_000, + 600_000, + "declared".to_string(), + None, + ) + ); + } + + #[tokio::test] + async fn declared_edges_pick_the_revision_as_of_the_window() { + let ctx = SessionContext::new(); + declared_table(&ctx); + let scan = ctx + .table(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) + .await + .unwrap(); + + // A historical window that predates frontend->db's second revision: + // the first revision was the edge's state then and must not be hidden + // by the newer one. + let window = GraphQueryWindow::from_observed(0, 150_000); + let plan = build_relationships_plan(vec![], Some(DeclaredSource { scan }), &window) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let mut rows: Vec<(String, f64)> = vec![]; + for batch in &batches { + let src = strings(batch, 5); + let confidence = batch + .column(10) + .as_any() + .downcast_ref::() + .unwrap(); + for (i, src) in src.into_iter().enumerate() { + rows.push((src, confidence.value(i))); + } + } + rows.sort_by(|a, b| a.0.cmp(&b.0)); + + // api->cache is valid inside this window; batch->queue is not yet. + assert_eq!( + rows, + vec![ + ("agent-1".to_string(), 0.8), + ("api".to_string(), 1.0), + ("frontend".to_string(), 0.5), + ("gateway".to_string(), 1.0), + ("webapp".to_string(), 0.5), + ] + ); + } + + #[tokio::test] + async fn declared_edges_union_calls_branch() { + let ctx = trace_table_ctx(); + declared_table(&ctx); + let trace = ctx.table("opentelemetry_traces").await.unwrap(); + let declared = ctx + .table(SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME) + .await + .unwrap(); + + let plan = build_relationships_plan( + vec![CallsSource { + service: trace_service_decl(&[SERVICE_NAME_COLUMN]), + scan: trace, + }], + Some(DeclaredSource { scan: declared }), + &test_window(), + ) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let mut provenance: Vec = batches.iter().flat_map(|b| strings(b, 9)).collect(); + provenance.sort(); + // One trace-derived calls edge plus the declared edges valid somewhere + // inside [0s, 600s): frontend, api, webapp, gateway, and the + // agent-asserted one. + assert_eq!( + provenance, + vec![ + "agent", "declared", "declared", "declared", "declared", "trace" + ] + ); + } +} From 3510ef7d4c13f9afa3ad8fdafa8842f98dfd5cb3 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Tue, 11 Aug 2026 16:55:27 +0800 Subject: [PATCH 26/92] feat!: update native histogram unsigned int types (#8824) * feat(native-histogram): store counts and span lengths as signed integers Native histograms are unreleased, so the on-disk integer payload columns are switched from unsigned to signed types without backward-compat: - count_u64 / zero_count_u64: uint64 -> int64 - positive_span_lengths / negative_span_lengths: list(uint32) -> list(int32) - Span.length (query-time model): u32 -> i32 The Prometheus remote-write v2 source carries these as uint64/uint32, so the unsigned->signed conversion at the ingestion boundary is overflow checked: an integer count >= 2^63 or a span length >= 2^31 is rejected with an explicit error rather than silently wrapping to a negative value. read_spans additionally rejects negative stored lengths to keep the non-negative invariant sound for downstream `as usize` casts. The UDAF accumulator's own observation counter (transient aggregation state, not part of the persisted histogram value) is intentionally left as uint64. Signed-off-by: Ning Sun * refactor(native-histogram): rename count/zero_count fields to _i64 Now that the integer payload columns are stored as int64, rename the field constants and persisted names to match: COUNT_U64_FIELD ("count_u64") -> COUNT_I64_FIELD ("count_i64") ZERO_COUNT_U64_FIELD ("zero_count_u64") -> ZERO_COUNT_I64_FIELD ("zero_count_i64") The local builder variables and the docs/JSON snapshot are updated to match. No backward-compat (unreleased feature). Signed-off-by: Ning Sun * test(native-histogram): refresh planner plan snapshot for signed types The mixed native-histogram range test embeds the full histogram Struct type in its expected plan string, which still carried the pre-rename unsigned fields. Update the snapshot to match the signed schema: positive/negative_span_lengths: List(UInt32) -> List(Int32) count_u64/zero_count_u64: UInt64 -> count_i64/zero_count_i64: Int64 Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- src/common/query/src/native_histogram.rs | 67 ++++++++++--------- src/promql/src/functions/native_histogram.rs | 2 +- src/query/src/promql/planner.rs | 26 +++---- src/servers/src/prom_remote_write/README.md | 2 +- src/servers/src/prom_remote_write/v2.rs | 50 +++++++++++--- .../tests/prom_remote_write_v2_test.rs | 6 +- tests-integration/tests/http.rs | 2 +- 7 files changed, 95 insertions(+), 60 deletions(-) diff --git a/src/common/query/src/native_histogram.rs b/src/common/query/src/native_histogram.rs index 931fb452ede..2db873487bd 100644 --- a/src/common/query/src/native_histogram.rs +++ b/src/common/query/src/native_histogram.rs @@ -23,13 +23,13 @@ use std::collections::BTreeMap; use std::sync::Arc; use datafusion::arrow::array::{ - Array, ArrayRef, Float64Array, Int32Array, ListArray, PrimitiveArray, StructArray, - TimestampMillisecondArray, UInt64Array, + Array, ArrayRef, Float64Array, Int32Array, Int64Array, ListArray, PrimitiveArray, StructArray, + TimestampMillisecondArray, }; use datafusion::arrow::buffer::NullBuffer; use datafusion::arrow::datatypes::{ ArrowPrimitiveType, DataType as ArrowDataType, Field, Float64Type, Int32Type, Int64Type, - TimestampMillisecondType, UInt32Type, UInt64Type, + TimestampMillisecondType, }; use datafusion_common::{DataFusionError, Result as DfResult}; use datatypes::data_type::{ConcreteDataType, DataType}; @@ -50,8 +50,8 @@ pub const POSITIVE_SPAN_OFFSETS_FIELD: &str = "positive_span_offsets"; pub const POSITIVE_SPAN_LENGTHS_FIELD: &str = "positive_span_lengths"; pub const NEGATIVE_SPAN_OFFSETS_FIELD: &str = "negative_span_offsets"; pub const NEGATIVE_SPAN_LENGTHS_FIELD: &str = "negative_span_lengths"; -pub const COUNT_U64_FIELD: &str = "count_u64"; -pub const ZERO_COUNT_U64_FIELD: &str = "zero_count_u64"; +pub const COUNT_I64_FIELD: &str = "count_i64"; +pub const ZERO_COUNT_I64_FIELD: &str = "zero_count_i64"; pub const POSITIVE_BUCKETS_I64_FIELD: &str = "positive_buckets_i64"; pub const NEGATIVE_BUCKETS_I64_FIELD: &str = "negative_buckets_i64"; pub const COUNT_F64_FIELD: &str = "count_f64"; @@ -72,8 +72,8 @@ pub const NATIVE_HISTOGRAM_FIELD_NAMES: &[&str] = &[ POSITIVE_SPAN_LENGTHS_FIELD, NEGATIVE_SPAN_OFFSETS_FIELD, NEGATIVE_SPAN_LENGTHS_FIELD, - COUNT_U64_FIELD, - ZERO_COUNT_U64_FIELD, + COUNT_I64_FIELD, + ZERO_COUNT_I64_FIELD, POSITIVE_BUCKETS_I64_FIELD, NEGATIVE_BUCKETS_I64_FIELD, COUNT_F64_FIELD, @@ -108,9 +108,9 @@ pub fn native_histogram_field_type(name: &str) -> Option { ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::int32_datatype())), ), POSITIVE_SPAN_LENGTHS_FIELD | NEGATIVE_SPAN_LENGTHS_FIELD => Some( - ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::uint32_datatype())), + ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::int32_datatype())), ), - COUNT_U64_FIELD | ZERO_COUNT_U64_FIELD => Some(ConcreteDataType::uint64_datatype()), + COUNT_I64_FIELD | ZERO_COUNT_I64_FIELD => Some(ConcreteDataType::int64_datatype()), POSITIVE_BUCKETS_I64_FIELD | NEGATIVE_BUCKETS_I64_FIELD => Some( ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::int64_datatype())), ), @@ -195,7 +195,7 @@ pub struct Span { /// The first bucket index, or the gap after the preceding span. pub offset: i32, /// Number of consecutive buckets in the span. - pub length: u32, + pub length: i32, } /// Inclusion rules for a materialized bucket's lower and upper bounds. @@ -1313,7 +1313,7 @@ where .collect() } -fn read_spans(offsets: Vec, lengths: Vec, name: &str) -> DfResult> { +fn read_spans(offsets: Vec, lengths: Vec, name: &str) -> DfResult> { if offsets.len() != lengths.len() { return Err(DataFusionError::Execution(format!( "native histogram {name} span offsets and lengths mismatch: {} vs {}", @@ -1321,11 +1321,18 @@ fn read_spans(offsets: Vec, lengths: Vec, name: &str) -> DfResult DfResult<()> { @@ -1352,12 +1359,12 @@ pub fn read_histogram(array: &StructArray, row: usize) -> DfResult(array, SCHEMA_FIELD, row)?; let positive_spans = read_spans( list_values::(array, POSITIVE_SPAN_OFFSETS_FIELD, row)?, - list_values::(array, POSITIVE_SPAN_LENGTHS_FIELD, row)?, + list_values::(array, POSITIVE_SPAN_LENGTHS_FIELD, row)?, "positive", )?; let negative_spans = read_spans( list_values::(array, NEGATIVE_SPAN_OFFSETS_FIELD, row)?, - list_values::(array, NEGATIVE_SPAN_LENGTHS_FIELD, row)?, + list_values::(array, NEGATIVE_SPAN_LENGTHS_FIELD, row)?, "negative", )?; @@ -1372,8 +1379,8 @@ pub fn read_histogram(array: &StructArray, row: usize) -> DfResult(array, COUNT_U64_FIELD, row)? as f64, - optional_primitive::(array, ZERO_COUNT_U64_FIELD, row)? + required_primitive::(array, COUNT_I64_FIELD, row)? as f64, + optional_primitive::(array, ZERO_COUNT_I64_FIELD, row)? .unwrap_or_default() as f64, list_values::(array, POSITIVE_BUCKETS_I64_FIELD, row)? .into_iter() @@ -1429,8 +1436,8 @@ pub fn build_histogram_array(values: &[Option]) -> ArrayRef { let mut positive_span_lengths = Vec::with_capacity(values.len()); let mut negative_span_offsets = Vec::with_capacity(values.len()); let mut negative_span_lengths = Vec::with_capacity(values.len()); - let mut count_u64 = Vec::>::with_capacity(values.len()); - let mut zero_count_u64 = Vec::>::with_capacity(values.len()); + let mut count_i64 = Vec::>::with_capacity(values.len()); + let mut zero_count_i64 = Vec::>::with_capacity(values.len()); let mut positive_buckets_i64 = Vec::with_capacity(values.len()); let mut negative_buckets_i64 = Vec::with_capacity(values.len()); let mut count_f64 = Vec::with_capacity(values.len()); @@ -1476,8 +1483,8 @@ pub fn build_histogram_array(values: &[Option]) -> ArrayRef { .map(|span| span.length) .collect(), )); - count_u64.push(None); - zero_count_u64.push(None); + count_i64.push(None); + zero_count_i64.push(None); positive_buckets_i64.push(list_opt(Vec::::new())); negative_buckets_i64.push(list_opt(Vec::::new())); count_f64.push(Some(histogram.count)); @@ -1495,8 +1502,8 @@ pub fn build_histogram_array(values: &[Option]) -> ArrayRef { positive_span_lengths.push(None); negative_span_offsets.push(None); negative_span_lengths.push(None); - count_u64.push(None); - zero_count_u64.push(None); + count_i64.push(None); + zero_count_i64.push(None); positive_buckets_i64.push(None); negative_buckets_i64.push(None); count_f64.push(None); @@ -1532,7 +1539,7 @@ pub fn build_histogram_array(values: &[Option]) -> ArrayRef { ), ( POSITIVE_SPAN_LENGTHS_FIELD, - Arc::new(ListArray::from_iter_primitive::( + Arc::new(ListArray::from_iter_primitive::( positive_span_lengths, )), ), @@ -1544,14 +1551,14 @@ pub fn build_histogram_array(values: &[Option]) -> ArrayRef { ), ( NEGATIVE_SPAN_LENGTHS_FIELD, - Arc::new(ListArray::from_iter_primitive::( + Arc::new(ListArray::from_iter_primitive::( negative_span_lengths, )), ), - (COUNT_U64_FIELD, Arc::new(UInt64Array::from(count_u64))), + (COUNT_I64_FIELD, Arc::new(Int64Array::from(count_i64))), ( - ZERO_COUNT_U64_FIELD, - Arc::new(UInt64Array::from(zero_count_u64)), + ZERO_COUNT_I64_FIELD, + Arc::new(Int64Array::from(zero_count_i64)), ), ( POSITIVE_BUCKETS_I64_FIELD, @@ -1662,7 +1669,7 @@ mod tests { assert_eq!(read_histogram(array, 0).unwrap(), Some(expected)); assert!( - primitive_child::(array, COUNT_U64_FIELD) + primitive_child::(array, COUNT_I64_FIELD) .unwrap() .is_null(0) ); diff --git a/src/promql/src/functions/native_histogram.rs b/src/promql/src/functions/native_histogram.rs index e6c9e84d59f..5b633defa96 100644 --- a/src/promql/src/functions/native_histogram.rs +++ b/src/promql/src/functions/native_histogram.rs @@ -2751,7 +2751,7 @@ mod tests { custom_values: Vec::new(), positive_spans: vec![Span { offset: 0, - length: positive_buckets.len() as u32, + length: positive_buckets.len() as i32, }], negative_spans: Vec::new(), count, diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 10462ea1828..0933bbcaca4 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -10026,19 +10026,19 @@ mod test { .unwrap() .display_indent_schema() .to_string(); - let expected = r#"Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_value, prom_mixed_range_histogram(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=["greptime_value", "greptime_native_histogram"] [timestamp:Timestamp(ms), greptime_value:Dictionary(Int64, Float64);N, greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, tag_0:Utf8, timestamp_range:Dictionary(Int64, Timestamp(ms))] - PromSeriesDivide: tags=["tag_0"] [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_value, prom_mixed_range_histogram(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] - PromRangeManipulate: req range=[-540000..100000000], interval=[60000], eval range=[300000], time index=[timestamp], values=["greptime_native_histogram", "greptime_value"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, greptime_value:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))] - PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - Filter: some_metric.timestamp >= TimestampMillisecond(-839999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] - TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(UInt32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(UInt32), "count_u64": UInt64, "zero_count_u64": UInt64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]"#; + let expected = r#"Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_value, prom_mixed_range_histogram(Utf8("sum_over_time"), timestamp_range, greptime_value, greptime_native_histogram) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[600000], time index=[timestamp], values=["greptime_value", "greptime_native_histogram"] [timestamp:Timestamp(ms), greptime_value:Dictionary(Int64, Float64);N, greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, tag_0:Utf8, timestamp_range:Dictionary(Int64, Timestamp(ms))] + PromSeriesDivide: tags=["tag_0"] [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Filter: greptime_value IS NOT NULL OR greptime_native_histogram IS NOT NULL [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + Projection: some_metric.timestamp, prom_mixed_range_float(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_value, prom_mixed_range_histogram(Utf8("rate"), timestamp_range, greptime_value, greptime_native_histogram, some_metric.timestamp, Int64(300000)) AS greptime_native_histogram, some_metric.tag_0 [timestamp:Timestamp(ms), greptime_value:Float64;N, greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, tag_0:Utf8] + PromRangeManipulate: req range=[-540000..100000000], interval=[60000], eval range=[300000], time index=[timestamp], values=["greptime_native_histogram", "greptime_value"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Dictionary(Int64, Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64)));N, greptime_value:Dictionary(Int64, Float64);N, timestamp_range:Dictionary(Int64, Timestamp(ms))] + PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + PromSeriesDivide: tags=["tag_0"] [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + Filter: some_metric.timestamp >= TimestampMillisecond(-839999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N] + TableScan: some_metric [tag_0:Utf8, timestamp:Timestamp(ms), greptime_native_histogram:Struct("schema": Int32, "zero_threshold": Float64, "sum": Float64, "reset_hint": Int32, "start_timestamp": Timestamp(ms), "custom_values": List(Float64), "positive_span_offsets": List(Int32), "positive_span_lengths": List(Int32), "negative_span_offsets": List(Int32), "negative_span_lengths": List(Int32), "count_i64": Int64, "zero_count_i64": Int64, "positive_buckets_i64": List(Int64), "negative_buckets_i64": List(Int64), "count_f64": Float64, "zero_count_f64": Float64, "positive_buckets_f64": List(Float64), "negative_buckets_f64": List(Float64));N, greptime_value:Float64;N]"#; assert_eq!(plan, expected); } diff --git a/src/servers/src/prom_remote_write/README.md b/src/servers/src/prom_remote_write/README.md index 1f15f8288a7..d21eead1901 100644 --- a/src/servers/src/prom_remote_write/README.md +++ b/src/servers/src/prom_remote_write/README.md @@ -47,7 +47,7 @@ Each histogram row stores one Struct field named - common scalar children: `schema`, `zero_threshold`, `sum`, `reset_hint`, `start_timestamp`; -- count children: `count_u64` / `zero_count_u64` or `count_f64` / `zero_count_f64`; +- count children: `count_i64` / `zero_count_i64` or `count_f64` / `zero_count_f64`; - list children for custom values, spans, and positive/negative buckets; - original Prometheus labels as Greptime tags. diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 9940ad2e7c2..5a25feb566a 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -371,6 +371,8 @@ fn native_histogram_struct_value(histogram: &Histogram) -> Result { validate_native_histogram(histogram, uses_float_counts)?; let mut items = Vec::with_capacity(NATIVE_HISTOGRAM_FIELD_NAMES.len()); + let positive_span_lengths = i32_span_lengths("positive", &histogram.positive_spans)?; + let negative_span_lengths = i32_span_lengths("negative", &histogram.negative_spans)?; items.extend([ pb_value(ValueData::I32Value(histogram.schema)), pb_value(ValueData::F64Value(histogram.zero_threshold)), @@ -381,9 +383,9 @@ fn native_histogram_struct_value(histogram: &Histogram) -> Result { )), f64_list_value(histogram.custom_values.iter().copied()), i32_list_value(histogram.positive_spans.iter().map(|span| span.offset)), - u32_list_value(histogram.positive_spans.iter().map(|span| span.length)), + i32_list_value(positive_span_lengths.iter().copied()), i32_list_value(histogram.negative_spans.iter().map(|span| span.offset)), - u32_list_value(histogram.negative_spans.iter().map(|span| span.length)), + i32_list_value(negative_span_lengths.iter().copied()), ]); if uses_float_counts { @@ -418,9 +420,23 @@ fn native_histogram_struct_value(histogram: &Histogram) -> Result { let positive_buckets = bucket_counts_from_deltas(&histogram.positive_deltas)?; let negative_buckets = bucket_counts_from_deltas(&histogram.negative_deltas)?; validate_integer_native_histogram_counts(histogram, &positive_buckets, &negative_buckets)?; + let count = i64::try_from(count) + .ok() + .context(error::InvalidPromRemoteRequestSnafu { + msg: format!( + "remote write v2 native histogram integer count {count} overflows i64" + ), + })?; + let zero_count = i64::try_from(zero_count).ok().context( + error::InvalidPromRemoteRequestSnafu { + msg: format!( + "remote write v2 native histogram integer zero_count {zero_count} overflows i64" + ), + }, + )?; items.extend([ - pb_value(ValueData::U64Value(count)), - pb_value(ValueData::U64Value(zero_count)), + pb_value(ValueData::I64Value(count)), + pb_value(ValueData::I64Value(zero_count)), i64_list_value(positive_buckets.iter().copied()), i64_list_value(negative_buckets.iter().copied()), null_pb_value(), @@ -779,8 +795,20 @@ fn i32_list_value(values: impl IntoIterator) -> Value { list_value(values.into_iter().map(ValueData::I32Value)) } -fn u32_list_value(values: impl IntoIterator) -> Value { - list_value(values.into_iter().map(ValueData::U32Value)) +fn i32_span_lengths(name: &str, spans: &[BucketSpan]) -> Result> { + spans + .iter() + .map(|span| { + i32::try_from(span.length) + .ok() + .context(error::InvalidPromRemoteRequestSnafu { + msg: format!( + "remote write v2 native histogram {name} span length {} overflows i32", + span.length + ), + }) + }) + .collect() } fn i64_list_value(values: impl IntoIterator) -> Value { @@ -1708,8 +1736,8 @@ mod tests { Some(ValueData::I32Value(0)) ); assert_eq!( - histogram_field_value(&rows, 0, COUNT_U64_FIELD), - Some(ValueData::U64Value(0)) + histogram_field_value(&rows, 0, COUNT_I64_FIELD), + Some(ValueData::I64Value(0)) ); assert_eq!(histogram_field_value(&rows, 0, COUNT_F64_FIELD), None); } @@ -1825,8 +1853,8 @@ mod tests { ); assert_eq!( - histogram_field_value(&rows, 0, COUNT_U64_FIELD), - Some(ValueData::U64Value(0)) + histogram_field_value(&rows, 0, COUNT_I64_FIELD), + Some(ValueData::I64Value(0)) ); assert_eq!(histogram_field_value(&rows, 0, COUNT_F64_FIELD), None); assert!(matches!( @@ -1839,7 +1867,7 @@ mod tests { POSITIVE_BUCKETS_F64_FIELD ))); - assert_eq!(histogram_field_value(&rows, 1, COUNT_U64_FIELD), None); + assert_eq!(histogram_field_value(&rows, 1, COUNT_I64_FIELD), None); assert_eq!( histogram_field_value(&rows, 1, COUNT_F64_FIELD), Some(ValueData::F64Value(6.0)) diff --git a/src/servers/tests/prom_remote_write_v2_test.rs b/src/servers/tests/prom_remote_write_v2_test.rs index a44583dccd1..94d260d3fac 100644 --- a/src/servers/tests/prom_remote_write_v2_test.rs +++ b/src/servers/tests/prom_remote_write_v2_test.rs @@ -18,7 +18,7 @@ use api::v1::value::ValueData; use api::v1::{ColumnSchema, Rows}; use bytes::Bytes; use common_query::native_histogram::{ - COUNT_U64_FIELD, NATIVE_HISTOGRAM_FIELD_NAMES, POSITIVE_BUCKETS_F64_FIELD, + COUNT_I64_FIELD, NATIVE_HISTOGRAM_FIELD_NAMES, POSITIVE_BUCKETS_F64_FIELD, POSITIVE_BUCKETS_I64_FIELD, POSITIVE_SPAN_OFFSETS_FIELD, SCHEMA_FIELD, }; use common_query::prelude::greptime_native_histogram; @@ -100,8 +100,8 @@ fn test_decode_remote_write_v2_native_histogram_dump() { Some(ValueData::I32Value(3)) ); assert_eq!( - histogram_field_value(rows, 0, COUNT_U64_FIELD), - Some(ValueData::U64Value(24)) + histogram_field_value(rows, 0, COUNT_I64_FIELD), + Some(ValueData::I64Value(24)) ); assert_eq!( list_i32_values(histogram_field_value(rows, 0, POSITIVE_SPAN_OFFSETS_FIELD)), diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 9d2f1eb8cb5..cd8af8768a4 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2878,7 +2878,7 @@ pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: Storag "prometheus_remote_write_v2_native_histogram_rows", &client, "select greptime_timestamp, greptime_native_histogram, job, instance from remote_write_v2_latency_seconds order by greptime_timestamp;", - "[[3000,{\"count_f64\":null,\"count_u64\":8,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[1],\"negative_span_lengths\":[1],\"negative_span_offsets\":[-2],\"positive_buckets_f64\":[],\"positive_buckets_i64\":[1,3,2],\"positive_span_lengths\":[3],\"positive_span_offsets\":[0],\"reset_hint\":2,\"schema\":1,\"start_timestamp\":1500,\"sum\":10.0,\"zero_count_f64\":null,\"zero_count_u64\":1,\"zero_threshold\":0.001},\"api\",\"localhost:9090\"],[4000,{\"count_f64\":6.0,\"count_u64\":null,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[],\"negative_span_lengths\":[],\"negative_span_offsets\":[],\"positive_buckets_f64\":[2.0,3.5],\"positive_buckets_i64\":[],\"positive_span_lengths\":[2],\"positive_span_offsets\":[3],\"reset_hint\":3,\"schema\":2,\"start_timestamp\":2500,\"sum\":20.0,\"zero_count_f64\":0.5,\"zero_count_u64\":null,\"zero_threshold\":0.002},\"api\",\"localhost:9090\"]]", + "[[3000,{\"count_f64\":null,\"count_i64\":8,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[1],\"negative_span_lengths\":[1],\"negative_span_offsets\":[-2],\"positive_buckets_f64\":[],\"positive_buckets_i64\":[1,3,2],\"positive_span_lengths\":[3],\"positive_span_offsets\":[0],\"reset_hint\":2,\"schema\":1,\"start_timestamp\":1500,\"sum\":10.0,\"zero_count_f64\":null,\"zero_count_i64\":1,\"zero_threshold\":0.001},\"api\",\"localhost:9090\"],[4000,{\"count_f64\":6.0,\"count_i64\":null,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[],\"negative_span_lengths\":[],\"negative_span_offsets\":[],\"positive_buckets_f64\":[2.0,3.5],\"positive_buckets_i64\":[],\"positive_span_lengths\":[2],\"positive_span_offsets\":[3],\"reset_hint\":3,\"schema\":2,\"start_timestamp\":2500,\"sum\":20.0,\"zero_count_f64\":0.5,\"zero_count_i64\":null,\"zero_threshold\":0.002},\"api\",\"localhost:9090\"]]", ) .await; From 72f6cf09bf7f7d78c84514163554ad2a29fc7f05 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 11 Aug 2026 17:05:26 +0800 Subject: [PATCH 27/92] refactor(procedure): centralize event context handling (#8834) * refactor(procedure): centralize event context handling Signed-off-by: WenyXu * refactor(meta): simplify migration trigger reason handling Signed-off-by: WenyXu * refactor(meta): avoid cloning event context Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- Cargo.lock | 1 + src/common/event-recorder/Cargo.toml | 1 + src/common/event-recorder/src/context.rs | 104 ++++++++ src/common/event-recorder/src/lib.rs | 2 + src/common/meta/src/ddl/alter_database.rs | 23 +- .../meta/src/ddl/alter_logical_tables.rs | 14 +- src/common/meta/src/ddl/alter_table.rs | 33 +-- .../src/ddl/alter_table/region_request.rs | 8 +- src/common/meta/src/ddl/create_database.rs | 19 +- src/common/meta/src/ddl/create_flow.rs | 22 +- .../meta/src/ddl/create_logical_tables.rs | 13 +- src/common/meta/src/ddl/create_table.rs | 34 +-- src/common/meta/src/ddl/create_view.rs | 14 +- src/common/meta/src/ddl/drop_database.rs | 25 +- src/common/meta/src/ddl/drop_flow.rs | 15 +- src/common/meta/src/ddl/drop_table.rs | 23 +- src/common/meta/src/ddl/drop_view.rs | 12 +- src/common/meta/src/ddl/event/database.rs | 25 +- src/common/meta/src/ddl/event/flow.rs | 28 +-- src/common/meta/src/ddl/event/table.rs | 57 +---- src/common/meta/src/ddl/event/view.rs | 25 +- .../meta/src/ddl/purge_dropped_table.rs | 36 +-- src/common/meta/src/ddl/test_util.rs | 9 +- .../src/ddl/tests/alter_logical_tables.rs | 52 +--- src/common/meta/src/ddl/tests/alter_table.rs | 73 ++---- src/common/meta/src/ddl/tests/create_flow.rs | 87 ++----- .../src/ddl/tests/create_logical_tables.rs | 39 +-- src/common/meta/src/ddl/tests/create_table.rs | 78 ++---- src/common/meta/src/ddl/tests/create_view.rs | 34 ++- .../meta/src/ddl/tests/drop_database.rs | 6 +- src/common/meta/src/ddl/tests/drop_flow.rs | 14 +- src/common/meta/src/ddl/tests/drop_table.rs | 225 +++++------------- src/common/meta/src/ddl/tests/drop_view.rs | 14 +- .../meta/src/ddl/tests/event/database.rs | 48 +--- src/common/meta/src/ddl/tests/event/flow.rs | 25 +- src/common/meta/src/ddl/tests/event/table.rs | 166 ++++--------- .../meta/src/ddl/tests/event/test_util.rs | 21 +- src/common/meta/src/ddl/tests/event/view.rs | 67 ++---- src/common/meta/src/ddl/truncate_table.rs | 24 +- src/common/meta/src/ddl/undrop_table.rs | 50 +--- src/common/meta/src/ddl_manager.rs | 210 ++++++++-------- src/common/meta/src/procedure_executor.rs | 5 +- src/common/meta/src/rpc/ddl.rs | 92 ++----- src/common/procedure-test/src/lib.rs | 4 + src/common/procedure/src/event.rs | 82 ++++++- src/common/procedure/src/lib.rs | 4 +- src/common/procedure/src/local.rs | 80 ++++++- src/common/procedure/src/local/runner.rs | 32 ++- src/common/procedure/src/procedure.rs | 44 +++- src/common/procedure/src/store.rs | 46 +++- src/common/procedure/src/watcher.rs | 1 + src/meta-srv/src/event/gc.rs | 86 +++---- src/meta-srv/src/event/region_migration.rs | 22 +- src/meta-srv/src/event/repartition.rs | 46 +--- src/meta-srv/src/gc/ctx.rs | 12 +- src/meta-srv/src/gc/handler.rs | 11 +- src/meta-srv/src/gc/mock.rs | 6 +- src/meta-srv/src/gc/mock/basic.rs | 10 +- src/meta-srv/src/gc/mock/concurrent.rs | 12 +- src/meta-srv/src/gc/mock/full_list.rs | 12 +- src/meta-srv/src/gc/mock/misc.rs | 6 +- src/meta-srv/src/gc/procedure.rs | 36 +-- src/meta-srv/src/gc/scheduler.rs | 16 +- .../src/procedure/region_migration.rs | 152 +++++++----- .../close_downgraded_region.rs | 4 +- .../downgrade_leader_region.rs | 7 +- .../src/procedure/region_migration/manager.rs | 14 +- .../region_migration/migration_abort.rs | 4 +- .../region_migration/open_candidate_region.rs | 30 ++- .../procedure/region_migration/test_util.rs | 8 +- .../upgrade_candidate_region.rs | 7 +- src/meta-srv/src/procedure/repartition.rs | 47 ++-- .../procedure/repartition/allocate_region.rs | 1 + .../src/procedure/repartition/collect.rs | 4 +- .../procedure/repartition/gc_requirement.rs | 1 + .../src/procedure/repartition/group.rs | 22 +- .../src/procedure/repartition/test_util.rs | 4 +- src/meta-srv/src/procedure/tests.rs | 14 +- src/meta-srv/src/procedure/wal_prune.rs | 12 +- src/meta-srv/src/service/procedure.rs | 1 + src/standalone/src/procedure.rs | 2 - tests-integration/src/tests/gc.rs | 2 - tests-integration/tests/gc_event.rs | 6 +- tests-integration/tests/repartition.rs | 2 - tests-integration/tests/table_ddl_event.rs | 4 +- tests/conf/frontend-test.toml.template | 3 + 86 files changed, 1142 insertions(+), 1650 deletions(-) create mode 100644 src/common/event-recorder/src/context.rs diff --git a/Cargo.lock b/Cargo.lock index 2e6fe8ae92c..18248bd442d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2409,6 +2409,7 @@ dependencies = [ "serde_json", "snafu 0.8.6", "store-api", + "strum 0.27.1", "tokio", "tokio-util", "toml 0.8.23", diff --git a/src/common/event-recorder/Cargo.toml b/src/common/event-recorder/Cargo.toml index 7990cd3d183..9a033d7009d 100644 --- a/src/common/event-recorder/Cargo.toml +++ b/src/common/event-recorder/Cargo.toml @@ -23,6 +23,7 @@ serde.workspace = true serde_json.workspace = true snafu.workspace = true store-api.workspace = true +strum.workspace = true tokio.workspace = true tokio-util.workspace = true diff --git a/src/common/event-recorder/src/context.rs b/src/common/event-recorder/src/context.rs new file mode 100644 index 00000000000..601276967e3 --- /dev/null +++ b/src/common/event-recorder/src/context.rs @@ -0,0 +1,104 @@ +// Copyright 2023 Greptime Team +// +// 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. + +use serde::{Deserialize, Serialize}; +use strum::{AsRefStr, EnumString}; + +/// Stable context recorded for a procedure event. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct PersistentEventContext { + pub reason: TriggerReason, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub protocol: Option, + #[serde(default, skip_serializing_if = "serde_json::Map::is_empty")] + pub extensions: serde_json::Map, +} + +impl PersistentEventContext { + /// Creates an event context with no additional extensions. + pub fn new(reason: TriggerReason) -> Self { + Self { + reason, + protocol: None, + extensions: Default::default(), + } + } + + /// Adds the protocol that originated the operation. + pub fn with_protocol(mut self, protocol: impl Into) -> Self { + self.protocol = Some(protocol.into()); + self + } +} + +impl Default for PersistentEventContext { + fn default() -> Self { + Self::new(TriggerReason::default()) + } +} + +/// Stable classification of a procedure trigger. +#[derive( + Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, AsRefStr, EnumString, +)] +#[serde(rename_all = "snake_case")] +#[strum(serialize_all = "snake_case")] +pub enum TriggerReason { + Manual, + AutoCreate, + AutoAlter, + AutoRepartition, + AutoRebalance, + RegionFailover, + ScheduledGc, + #[default] + #[serde(other)] + Unknown, +} + +impl TriggerReason { + pub fn from_extension(value: &str) -> Self { + value.parse().unwrap_or_default() + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn test_event_context_serialization() { + let context = PersistentEventContext::new(TriggerReason::Manual).with_protocol("mysql"); + + assert_eq!( + json!({ + "reason": "manual", + "protocol": "mysql", + }), + serde_json::to_value(context).unwrap() + ); + assert_eq!( + json!({ "reason": "manual" }), + serde_json::to_value(PersistentEventContext::new(TriggerReason::Manual)).unwrap() + ); + } + + #[test] + fn test_trigger_reason_deserializes_unknown_value() { + let reason: TriggerReason = serde_json::from_str("\"future_reason\"").unwrap(); + assert_eq!(TriggerReason::Unknown, reason); + } +} diff --git a/src/common/event-recorder/src/lib.rs b/src/common/event-recorder/src/lib.rs index 3e624824980..5b9cb4a7b97 100644 --- a/src/common/event-recorder/src/lib.rs +++ b/src/common/event-recorder/src/lib.rs @@ -14,6 +14,7 @@ #![feature(duration_constructors)] +pub mod context; pub mod error; pub mod event_table; pub mod recorder; @@ -21,4 +22,5 @@ pub mod recorder; #[cfg(any(test, feature = "testing"))] pub mod testing; +pub use context::{PersistentEventContext, TriggerReason}; pub use recorder::*; diff --git a/src/common/meta/src/ddl/alter_database.rs b/src/common/meta/src/ddl/alter_database.rs index f9344213025..ac24ff77e11 100644 --- a/src/common/meta/src/ddl/alter_database.rs +++ b/src/common/meta/src/ddl/alter_database.rs @@ -15,7 +15,7 @@ use async_trait::async_trait; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, Status, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, Status, }; use common_telemetry::tracing::info; use serde::{Deserialize, Serialize}; @@ -32,7 +32,7 @@ use crate::key::DeserializedValueWithBytes; use crate::key::schema_name::{SchemaName, SchemaNameKey, SchemaNameValue}; use crate::lock_key::{CatalogLock, SchemaLock}; use crate::rpc::ddl::UnsetDatabaseOption::{self}; -use crate::rpc::ddl::{AlterDatabaseKind, AlterDatabaseTask, EventContext, SetDatabaseOption}; +use crate::rpc::ddl::{AlterDatabaseKind, AlterDatabaseTask, SetDatabaseOption}; pub struct AlterDatabaseProcedure { pub context: DdlContext, @@ -73,14 +73,10 @@ fn build_new_schema_value( impl AlterDatabaseProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::AlterDatabase"; - pub fn new( - task: AlterDatabaseTask, - event_context: EventContext, - context: DdlContext, - ) -> Result { + pub fn new(task: AlterDatabaseTask, context: DdlContext) -> Result { Ok(Self { context, - data: AlterDatabaseData::new(task, event_context)?, + data: AlterDatabaseData::new(task)?, }) } @@ -180,10 +176,7 @@ impl Procedure for AlterDatabaseProcedure { LockKey::new(lock_key) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(ALTER_DATABASE_EVENT_TYPE) { return None; } @@ -193,7 +186,6 @@ impl Procedure for AlterDatabaseProcedure { self.data.catalog(), self.data.schema(), &self.data.kind, - self.data.event_context.clone(), ) } else { DatabaseDdlEvent::alter_lifecycle(self.data.catalog(), self.data.schema()) @@ -217,19 +209,16 @@ pub struct AlterDatabaseData { catalog_name: String, schema_name: String, schema_value: Option>, - #[serde(default)] - event_context: EventContext, } impl AlterDatabaseData { - pub fn new(task: AlterDatabaseTask, event_context: EventContext) -> Result { + pub fn new(task: AlterDatabaseTask) -> Result { Ok(Self { state: AlterDatabaseState::Prepare, kind: AlterDatabaseKind::try_from(task.alter_expr.kind.unwrap())?, catalog_name: task.alter_expr.catalog_name, schema_name: task.alter_expr.schema_name, schema_value: None, - event_context, }) } diff --git a/src/common/meta/src/ddl/alter_logical_tables.rs b/src/common/meta/src/ddl/alter_logical_tables.rs index f2f15430933..9207f910577 100644 --- a/src/common/meta/src/ddl/alter_logical_tables.rs +++ b/src/common/meta/src/ddl/alter_logical_tables.rs @@ -20,7 +20,7 @@ use api::region::RegionResponse; use async_trait::async_trait; use common_catalog::format_full_table_name; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; -use common_procedure::{Context, EventRuntimeContext, EventTrigger, LockKey, Procedure, Status}; +use common_procedure::{Context, EventContext, EventTrigger, LockKey, Procedure, Status}; use common_telemetry::{debug, error, info, warn}; pub use executor::make_alter_region_request; use serde::{Deserialize, Serialize}; @@ -47,7 +47,7 @@ use crate::key::table_info::TableInfoValue; use crate::key::table_route::PhysicalTableRouteValue; use crate::lock_key::{CatalogLock, SchemaLock, TableLock}; use crate::metrics; -use crate::rpc::ddl::{AlterTableTask, EventContext}; +use crate::rpc::ddl::AlterTableTask; use crate::rpc::router::RegionRoute; pub struct AlterLogicalTablesProcedure { @@ -88,7 +88,6 @@ impl AlterLogicalTablesProcedure { pub fn new( tasks: Vec, physical_table_id: TableId, - event_context: EventContext, context: DdlContext, ) -> Self { Self { @@ -101,7 +100,6 @@ impl AlterLogicalTablesProcedure { physical_table_info: None, physical_columns: vec![], table_cache_keys_to_invalidate: vec![], - event_context, }, physical_table_route: None, } @@ -322,10 +320,7 @@ impl Procedure for AlterLogicalTablesProcedure { LockKey::new(lock_key) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::AlterLogicalTables.as_str()) @@ -358,7 +353,6 @@ impl Procedure for AlterLogicalTablesProcedure { locators, self.data.tasks.len(), kinds, - self.data.event_context.clone(), ))) } } @@ -375,8 +369,6 @@ pub struct AlterTablesData { physical_table_info: Option>, physical_columns: Vec, table_cache_keys_to_invalidate: Vec, - #[serde(default)] - event_context: EventContext, } impl AlterTablesData { diff --git a/src/common/meta/src/ddl/alter_table.rs b/src/common/meta/src/ddl/alter_table.rs index 21975fab714..63342a7bbb6 100644 --- a/src/common/meta/src/ddl/alter_table.rs +++ b/src/common/meta/src/ddl/alter_table.rs @@ -27,7 +27,7 @@ use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE}; use common_error::ext::BoxedError; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, ContextProvider, Error as ProcedureError, EventRuntimeContext, + Context as ProcedureContext, ContextProvider, Error as ProcedureError, EventContext, EventTrigger, LockKey, PoisonKey, PoisonKeys, Procedure, ProcedureId, Status, StringKey, }; use common_telemetry::{error, info, warn}; @@ -58,7 +58,7 @@ use crate::key::{DeserializedValueWithBytes, RegionDistribution}; use crate::lock_key::{CatalogLock, RegionLock, SchemaLock, TableLock, TableNameLock}; use crate::metrics; use crate::poison_key::table_poison_key; -use crate::rpc::ddl::{AlterTableTask, EventContext}; +use crate::rpc::ddl::AlterTableTask; use crate::rpc::router::{RegionRoute, find_leaders, region_distribution}; /// The alter table procedure @@ -97,24 +97,18 @@ fn build_executor_from_alter_expr(alter_data: &AlterTableData) -> AlterTableExec impl AlterTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::AlterTable"; - pub fn new( - table_id: TableId, - task: AlterTableTask, - event_context: EventContext, - context: DdlContext, - ) -> Result { - Self::new_with_region_locks(table_id, task, event_context, vec![], context) + pub fn new(table_id: TableId, task: AlterTableTask, context: DdlContext) -> Result { + Self::new_with_region_locks(table_id, task, vec![], context) } pub(crate) fn new_with_region_locks( table_id: TableId, task: AlterTableTask, - event_context: EventContext, region_locks: Vec, context: DdlContext, ) -> Result { task.validate()?; - let data = AlterTableData::new(task, table_id, event_context, region_locks); + let data = AlterTableData::new(task, table_id, region_locks); let executor = build_executor_from_alter_expr(&data); Ok(Self { context, @@ -569,10 +563,7 @@ impl Procedure for AlterTableProcedure { PoisonKeys::new(vec![self.table_poison_key()]) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::AlterTable.as_str()) @@ -591,7 +582,7 @@ impl Procedure for AlterTableProcedure { .kind .as_ref() .and_then(alter_table_kind_name); - TableDdlEvent::alter_table_submitted(locator, kind, self.data.event_context.clone()) + TableDdlEvent::alter_table_submitted(locator, kind) } _ => TableDdlEvent::lifecycle(TableDdlEventType::AlterTable, [locator]), }; @@ -619,8 +610,6 @@ pub struct AlterTableData { task: AlterTableTask, table_id: TableId, #[serde(default)] - event_context: EventContext, - #[serde(default)] column_metadatas: Vec, /// Table info value before alteration. table_info_value: Option>, @@ -632,17 +621,11 @@ pub struct AlterTableData { } impl AlterTableData { - pub fn new( - task: AlterTableTask, - table_id: TableId, - event_context: EventContext, - region_locks: Vec, - ) -> Self { + pub fn new(task: AlterTableTask, table_id: TableId, region_locks: Vec) -> Self { Self { state: AlterTableState::Prepare, task, table_id, - event_context, column_metadatas: vec![], table_info_value: None, region_distribution: None, diff --git a/src/common/meta/src/ddl/alter_table/region_request.rs b/src/common/meta/src/ddl/alter_table/region_request.rs index ef8d1487c70..73edd4c3615 100644 --- a/src/common/meta/src/ddl/alter_table/region_request.rs +++ b/src/common/meta/src/ddl/alter_table/region_request.rs @@ -144,7 +144,7 @@ mod tests { }; use crate::key::table_route::TableRouteValue; use crate::peer::Peer; - use crate::rpc::ddl::{AlterTableTask, EventContext}; + use crate::rpc::ddl::AlterTableTask; use crate::rpc::router::{Region, RegionRoute}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -240,8 +240,7 @@ mod tests { }, }; - let mut procedure = - AlterTableProcedure::new(table_id, task, EventContext::default(), ddl_context).unwrap(); + let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context).unwrap(); procedure.on_prepare().await.unwrap(); let alter_kind = procedure.make_region_alter_kind().unwrap(); let Some(Body::Alter(alter_region_request)) = @@ -297,8 +296,7 @@ mod tests { }, }; - let mut procedure = - AlterTableProcedure::new(table_id, task, EventContext::default(), ddl_context).unwrap(); + let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context).unwrap(); procedure.on_prepare().await.unwrap(); let alter_kind = procedure.make_region_alter_kind().unwrap(); let Some(Body::Alter(alter_region_request)) = diff --git a/src/common/meta/src/ddl/create_database.rs b/src/common/meta/src/ddl/create_database.rs index b48e700a519..82d467888a5 100644 --- a/src/common/meta/src/ddl/create_database.rs +++ b/src/common/meta/src/ddl/create_database.rs @@ -19,8 +19,8 @@ use async_trait::async_trait; use common_error::ext::{BoxedError, ErrorExt}; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, - ProcedureId, Status, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, ProcedureId, + Status, }; use serde::{Deserialize, Serialize}; use serde_with::{DefaultOnNull, serde_as}; @@ -34,7 +34,7 @@ use crate::error::{self, Result}; use crate::instruction::{CacheIdent, UserCacheIdent}; use crate::key::schema_name::{SchemaNameKey, SchemaNameValue}; use crate::lock_key::{CatalogLock, SchemaLock}; -use crate::rpc::ddl::{CreatorGrantIntent, EventContext}; +use crate::rpc::ddl::CreatorGrantIntent; /// Describes the creator-access result of an atomic create. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -93,7 +93,6 @@ impl CreateDatabaseProcedure { create_if_not_exists: bool, options: HashMap, creator: Option, - event_context: EventContext, context: DdlContext, ) -> Self { Self { @@ -105,7 +104,6 @@ impl CreateDatabaseProcedure { create_if_not_exists, options, creator, - event_context, }, } } @@ -280,10 +278,7 @@ impl Procedure for CreateDatabaseProcedure { LockKey::new(lock_key) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(CREATE_DATABASE_EVENT_TYPE) { return None; } @@ -294,7 +289,6 @@ impl Procedure for CreateDatabaseProcedure { &self.data.schema, self.data.create_if_not_exists, &self.data.options, - self.data.event_context.clone(), ) } else { DatabaseDdlEvent::create_lifecycle(&self.data.catalog, &self.data.schema) @@ -327,8 +321,6 @@ pub struct CreateDatabaseData { /// Authenticated creator whose access is ensured, absent for legacy schema-only requests. #[serde(default)] pub creator: Option, - #[serde(default)] - pub event_context: EventContext, } #[cfg(test)] @@ -423,7 +415,6 @@ mod tests { username: "alice".to_string(), created_at_ns: 1, }), - EventContext::default(), context, ); (procedure, committer) @@ -583,7 +574,6 @@ mod tests { false, HashMap::new(), None, - EventContext::default(), context, ); procedure.data.state = CreateDatabaseState::CreateMetadata; @@ -619,7 +609,6 @@ mod tests { username: "alice".to_string(), created_at_ns: 1, }), - EventContext::default(), context, ); procedure.data.state = CreateDatabaseState::CreateMetadata; diff --git a/src/common/meta/src/ddl/create_flow.rs b/src/common/meta/src/ddl/create_flow.rs index 3ef4dadbc63..a2f306306d8 100644 --- a/src/common/meta/src/ddl/create_flow.rs +++ b/src/common/meta/src/ddl/create_flow.rs @@ -24,8 +24,8 @@ use async_trait::async_trait; use common_catalog::format_full_flow_name; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, - ProcedureState, Result as ProcedureResult, Status, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, ProcedureState, + Result as ProcedureResult, Status, }; use common_telemetry::info; use common_telemetry::tracing_context::TracingContext; @@ -50,7 +50,7 @@ use crate::key::{DeserializedValueWithBytes, FlowId, FlowPartitionId}; use crate::lock_key::{CatalogLock, FlowNameLock}; use crate::metrics; use crate::peer::Peer; -use crate::rpc::ddl::{CreateFlowTask, EventContext, FlowQueryContext, QueryContext}; +use crate::rpc::ddl::{CreateFlowTask, FlowQueryContext, QueryContext}; /// The procedure of flow creation. pub struct CreateFlowProcedure { @@ -62,12 +62,7 @@ impl CreateFlowProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::CreateFlow"; /// Returns a new [CreateFlowProcedure]. - pub fn new( - task: CreateFlowTask, - query_context: QueryContext, - event_context: EventContext, - context: DdlContext, - ) -> Self { + pub fn new(task: CreateFlowTask, query_context: QueryContext, context: DdlContext) -> Self { Self { context, data: CreateFlowData { @@ -81,7 +76,6 @@ impl CreateFlowProcedure { prev_flow_info_value: None, did_replace: false, flow_type: None, - event_context, }, } } @@ -410,10 +404,7 @@ impl Procedure for CreateFlowProcedure { ]) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(CREATE_FLOW_EVENT_TYPE) { return None; } @@ -428,7 +419,6 @@ impl Procedure for CreateFlowProcedure { expire_after: self.data.task.expire_after, eval_interval_secs: self.data.task.eval_interval_secs, }, - self.data.event_context.clone(), ), EventTrigger::Succeeded => { let flow_id = match ctx.lifecycle_state { @@ -718,8 +708,6 @@ pub struct CreateFlowData { #[serde(default)] pub(crate) did_replace: bool, pub(crate) flow_type: Option, - #[serde(default)] - pub(crate) event_context: EventContext, } impl CreateFlowData { diff --git a/src/common/meta/src/ddl/create_logical_tables.rs b/src/common/meta/src/ddl/create_logical_tables.rs index 1c9efb865e4..9b231f3c538 100644 --- a/src/common/meta/src/ddl/create_logical_tables.rs +++ b/src/common/meta/src/ddl/create_logical_tables.rs @@ -24,8 +24,8 @@ use common_catalog::consts::METRIC_ENGINE; use common_event_recorder::Event; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, - ProcedureState, Status, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, ProcedureState, + Status, }; use common_telemetry::{debug, error, warn}; use futures::future; @@ -48,7 +48,7 @@ use crate::error::Result; use crate::key::table_route::TableRouteValue; use crate::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; use crate::metrics; -use crate::rpc::ddl::{CreateTableTask, EventContext}; +use crate::rpc::ddl::CreateTableTask; use crate::rpc::router::{RegionRoute, find_leaders}; pub struct CreateLogicalTablesProcedure { @@ -62,7 +62,6 @@ impl CreateLogicalTablesProcedure { pub fn new( tasks: Vec, physical_table_id: TableId, - event_context: EventContext, context: DdlContext, ) -> Self { Self { @@ -75,7 +74,6 @@ impl CreateLogicalTablesProcedure { physical_region_numbers: vec![], physical_columns: vec![], physical_partition_columns: vec![], - event_context, }, } } @@ -273,7 +271,7 @@ impl Procedure for CreateLogicalTablesProcedure { LockKey::new(lock_key) } - fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::CreateLogicalTables.as_str()) @@ -284,7 +282,6 @@ impl Procedure for CreateLogicalTablesProcedure { EventTrigger::Submitted => TableDdlEvent::create_logical_tables_submitted( self.event_locators(), self.data.tasks.len(), - self.data.event_context.clone(), ), EventTrigger::Succeeded => match ctx.lifecycle_state { ProcedureState::Done { @@ -329,8 +326,6 @@ pub struct CreateTablesData { physical_region_numbers: Vec, physical_columns: Vec, physical_partition_columns: Vec, - #[serde(default)] - event_context: EventContext, } impl CreateTablesData { diff --git a/src/common/meta/src/ddl/create_table.rs b/src/common/meta/src/ddl/create_table.rs index c63d7f7a34d..ce3e1aa060a 100644 --- a/src/common/meta/src/ddl/create_table.rs +++ b/src/common/meta/src/ddl/create_table.rs @@ -23,8 +23,8 @@ use common_procedure::error::{ }; use common_procedure::local::DynamicKeyLockGuard; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, - ProcedureId, ProcedureState, Status, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, ProcedureId, + ProcedureState, Status, }; use common_telemetry::info; use serde::{Deserialize, Serialize}; @@ -47,7 +47,7 @@ use crate::lock_key::{CatalogLock, SchemaLock, TableNameLock}; use crate::metrics; use crate::peer::PeerAllocContext; use crate::region_keeper::OperatingRegionGuard; -use crate::rpc::ddl::{CreateTableTask, EventContext, QueryContext}; +use crate::rpc::ddl::{CreateTableTask, QueryContext}; use crate::rpc::router::{RegionRoute, operating_leader_region_roles}; use crate::wal_provider::{ RegionWalOptions, acquire_remote_wal_read_locks, optional_region_wal_options_serde, @@ -84,17 +84,20 @@ fn build_executor_from_create_table_data( impl CreateTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::CreateTable"; - pub fn new( + pub fn new(task: CreateTableTask, context: DdlContext) -> Result { + Self::new_with_query_context(task, QueryContext::default(), context) + } + + pub fn new_with_query_context( task: CreateTableTask, query_context: QueryContext, - event_context: EventContext, context: DdlContext, ) -> Result { let executor = build_executor_from_create_table_data(&task.create_table)?; Ok(Self { context, - data: CreateTableData::new_with_query_context(task, event_context, query_context), + data: CreateTableData::new(task, query_context), opening_regions: vec![], executor, remote_wal_lock_guards: vec![], @@ -396,10 +399,7 @@ impl Procedure for CreateTableProcedure { ]) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::CreateTable.as_str()) @@ -415,7 +415,6 @@ impl Procedure for CreateTableProcedure { locator, create_table.create_if_not_exists, &create_table.engine, - self.data.event_context.clone(), ) } EventTrigger::Succeeded => match ctx.lifecycle_state { @@ -456,8 +455,6 @@ pub struct CreateTableData { #[serde(default)] pub query_context: QueryContext, #[serde(default)] - pub event_context: EventContext, - #[serde(default)] pub column_metadatas: Vec, /// None stands for not allocated yet. pub(crate) table_route: Option, @@ -468,21 +465,12 @@ pub struct CreateTableData { } impl CreateTableData { - pub fn new(task: CreateTableTask, event_context: EventContext) -> Self { - Self::new_with_query_context(task, event_context, QueryContext::default()) - } - - pub fn new_with_query_context( - task: CreateTableTask, - event_context: EventContext, - query_context: QueryContext, - ) -> Self { + pub fn new(task: CreateTableTask, query_context: QueryContext) -> Self { CreateTableData { state: CreateTableState::Prepare, column_metadatas: vec![], task, query_context, - event_context, table_route: None, region_wal_options: None, } diff --git a/src/common/meta/src/ddl/create_view.rs b/src/common/meta/src/ddl/create_view.rs index a7175854b76..666ef19efdf 100644 --- a/src/common/meta/src/ddl/create_view.rs +++ b/src/common/meta/src/ddl/create_view.rs @@ -16,8 +16,8 @@ use async_trait::async_trait; use common_event_recorder::Event; use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, - ProcedureState, Status, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, ProcedureState, + Status, }; use common_telemetry::info; use serde::{Deserialize, Serialize}; @@ -35,7 +35,7 @@ use crate::instruction::CacheIdent; use crate::key::table_name::TableNameKey; use crate::lock_key::{CatalogLock, SchemaLock, TableNameLock}; use crate::metrics; -use crate::rpc::ddl::{CreateViewTask, EventContext}; +use crate::rpc::ddl::CreateViewTask; // The procedure to execute `[CreateViewTask]`. pub struct CreateViewProcedure { @@ -46,14 +46,13 @@ pub struct CreateViewProcedure { impl CreateViewProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::CreateView"; - pub fn new(task: CreateViewTask, event_context: EventContext, context: DdlContext) -> Self { + pub fn new(task: CreateViewTask, context: DdlContext) -> Self { Self { context, data: CreateViewData { state: CreateViewState::Prepare, task, need_update: false, - event_context, }, } } @@ -272,7 +271,7 @@ impl Procedure for CreateViewProcedure { ]) } - fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(CREATE_VIEW_EVENT_TYPE) { return None; } @@ -290,7 +289,6 @@ impl Procedure for CreateViewProcedure { referenced_table_count: self.data.task.table_names().len(), column_count: self.data.task.columns().len(), }, - self.data.event_context.clone(), ) } EventTrigger::Succeeded => match ctx.lifecycle_state { @@ -344,8 +342,6 @@ pub struct CreateViewData { pub task: CreateViewTask, /// Whether to update the view info. pub need_update: bool, - #[serde(default)] - pub event_context: EventContext, } impl CreateViewData { diff --git a/src/common/meta/src/ddl/drop_database.rs b/src/common/meta/src/ddl/drop_database.rs index 7772e966cb2..5bbd229ec5b 100644 --- a/src/common/meta/src/ddl/drop_database.rs +++ b/src/common/meta/src/ddl/drop_database.rs @@ -23,7 +23,7 @@ use std::fmt::Debug; use common_error::ext::BoxedError; use common_procedure::error::{ExternalSnafu, FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, Result as ProcedureResult, Status, }; use futures::stream::BoxStream; @@ -38,13 +38,11 @@ use crate::ddl::utils::map_to_procedure_error; use crate::error::Result; use crate::key::table_name::TableNameValue; use crate::lock_key::{CatalogLock, SchemaLock}; -use crate::rpc::ddl::EventContext; pub struct DropDatabaseProcedure { /// The context of procedure runtime. runtime_context: DdlContext, context: DropDatabaseContext, - event_context: EventContext, state: Box, } @@ -87,16 +85,9 @@ pub(crate) trait State: Send + Debug { impl DropDatabaseProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::DropDatabase"; - pub fn new( - catalog: String, - schema: String, - drop_if_exists: bool, - event_context: EventContext, - context: DdlContext, - ) -> Self { + pub fn new(catalog: String, schema: String, drop_if_exists: bool, context: DdlContext) -> Self { Self { runtime_context: context, - event_context, context: DropDatabaseContext { catalog, schema, @@ -113,13 +104,11 @@ impl DropDatabaseProcedure { catalog, schema, drop_if_exists, - event_context, state, } = serde_json::from_str(json).context(FromJsonSnafu)?; Ok(Self { runtime_context, - event_context, context: DropDatabaseContext { catalog, schema, @@ -170,7 +159,6 @@ impl Procedure for DropDatabaseProcedure { catalog: &self.context.catalog, schema: &self.context.schema, drop_if_exists: self.context.drop_if_exists, - event_context: &self.event_context, state: self.state.as_ref(), }; @@ -186,10 +174,7 @@ impl Procedure for DropDatabaseProcedure { LockKey::new(lock_key) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(DROP_DATABASE_EVENT_TYPE) { return None; } @@ -199,7 +184,6 @@ impl Procedure for DropDatabaseProcedure { &self.context.catalog, &self.context.schema, self.context.drop_if_exists, - self.event_context.clone(), ) } else { DatabaseDdlEvent::drop_lifecycle(&self.context.catalog, &self.context.schema) @@ -215,7 +199,6 @@ struct DropDatabaseData<'a> { // The schema name schema: &'a str, drop_if_exists: bool, - event_context: &'a EventContext, state: &'a dyn State, } @@ -226,7 +209,5 @@ struct DropDatabaseOwnedData { // The schema name schema: String, drop_if_exists: bool, - #[serde(default)] - event_context: EventContext, state: Box, } diff --git a/src/common/meta/src/ddl/drop_flow.rs b/src/common/meta/src/ddl/drop_flow.rs index 99c18ebc13b..783fb7796e0 100644 --- a/src/common/meta/src/ddl/drop_flow.rs +++ b/src/common/meta/src/ddl/drop_flow.rs @@ -21,7 +21,7 @@ use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, Result as ProcedureResult, Status, }; use common_telemetry::info; @@ -41,7 +41,7 @@ use crate::key::flow::flow_info::FlowInfoValue; use crate::key::flow::flow_route::FlowRouteValue; use crate::lock_key::{CatalogLock, FlowLock}; use crate::metrics; -use crate::rpc::ddl::{DropFlowTask, EventContext}; +use crate::rpc::ddl::DropFlowTask; /// The procedure for dropping a flow. pub struct DropFlowProcedure { @@ -54,7 +54,7 @@ pub struct DropFlowProcedure { impl DropFlowProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::DropFlow"; - pub fn new(task: DropFlowTask, event_context: EventContext, context: DdlContext) -> Self { + pub fn new(task: DropFlowTask, context: DdlContext) -> Self { Self { context, data: DropFlowData { @@ -62,7 +62,6 @@ impl DropFlowProcedure { task, flow_info_value: None, flow_route_values: vec![], - event_context, }, } } @@ -223,10 +222,7 @@ impl Procedure for DropFlowProcedure { LockKey::new(lock_key) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(DROP_FLOW_EVENT_TYPE) { return None; } @@ -237,7 +233,6 @@ impl Procedure for DropFlowProcedure { &self.data.task.flow_name, self.data.task.flow_id, self.data.task.drop_if_exists, - self.data.event_context.clone(), ), _ => FlowDdlEvent::drop_lifecycle( &self.data.task.catalog_name, @@ -257,8 +252,6 @@ pub(crate) struct DropFlowData { task: DropFlowTask, pub(crate) flow_info_value: Option, pub(crate) flow_route_values: Vec, - #[serde(default)] - pub(crate) event_context: EventContext, } /// The state of drop flow diff --git a/src/common/meta/src/ddl/drop_table.rs b/src/common/meta/src/ddl/drop_table.rs index 7c8dec55e49..b73bd95e1bc 100644 --- a/src/common/meta/src/ddl/drop_table.rs +++ b/src/common/meta/src/ddl/drop_table.rs @@ -22,8 +22,8 @@ use common_error::ext::BoxedError; use common_event_recorder::Event; use common_procedure::error::{ExternalSnafu, FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, Error as ProcedureError, EventRuntimeContext, EventTrigger, - LockKey, Procedure, Result as ProcedureResult, Status, + Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, + Procedure, Result as ProcedureResult, Status, }; use common_telemetry::info; use common_telemetry::tracing::warn; @@ -48,7 +48,7 @@ use crate::key::table_route::TableRouteValue; use crate::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; use crate::metrics; use crate::region_keeper::OperatingRegionGuard; -use crate::rpc::ddl::{DropTableTask, EventContext}; +use crate::rpc::ddl::DropTableTask; use crate::rpc::router::{RegionRoute, operating_leader_region_roles}; #[cfg(feature = "enterprise")] @@ -74,12 +74,11 @@ pub struct DropTableProcedure { impl DropTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::DropTable"; - pub fn new(task: DropTableTask, context: DdlContext, event_context: EventContext) -> Self { + pub fn new(task: DropTableTask, context: DdlContext) -> Self { let data = DropTableData::new( task, cfg!(feature = "enterprise") && context.soft_drop_enabled, context.soft_drop_retention, - event_context, ); let executor = data.build_executor(); Self { @@ -414,7 +413,7 @@ impl Procedure for DropTableProcedure { LockKey::new(lock_key) } - fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::DropTable.as_str()) @@ -425,11 +424,9 @@ impl Procedure for DropTableProcedure { let locator = TableDdlLocator::new(&task.catalog, &task.schema, &task.table) .with_table_id(task.table_id); let event = match &ctx.trigger { - EventTrigger::Submitted => TableDdlEvent::drop_table_submitted( - locator, - task.drop_if_exists, - self.data.event_context.clone(), - ), + EventTrigger::Submitted => { + TableDdlEvent::drop_table_submitted(locator, task.drop_if_exists) + } _ => TableDdlEvent::lifecycle(TableDdlEventType::DropTable, [locator]), }; @@ -486,8 +483,6 @@ pub struct DropTableData { pub soft_drop_retention_millis: Option, #[serde(default)] pub drop_generation: Option, - #[serde(default)] - pub event_context: EventContext, } impl DropTableData { @@ -495,7 +490,6 @@ impl DropTableData { task: DropTableTask, soft_drop_enabled: bool, soft_drop_retention: Option, - event_context: EventContext, ) -> Self { Self { state: DropTableState::Prepare, @@ -509,7 +503,6 @@ impl DropTableData { retention_expires_at: None, soft_drop_retention_millis: soft_drop_retention_millis(soft_drop_retention), drop_generation: None, - event_context, } } diff --git a/src/common/meta/src/ddl/drop_view.rs b/src/common/meta/src/ddl/drop_view.rs index ece1d2b06b7..32d843fbd48 100644 --- a/src/common/meta/src/ddl/drop_view.rs +++ b/src/common/meta/src/ddl/drop_view.rs @@ -16,7 +16,7 @@ use async_trait::async_trait; use common_event_recorder::Event; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, Result as ProcedureResult, Status, }; use common_telemetry::info; @@ -35,7 +35,7 @@ use crate::instruction::CacheIdent; use crate::key::table_name::TableNameKey; use crate::lock_key::{CatalogLock, SchemaLock, TableLock}; use crate::metrics; -use crate::rpc::ddl::{DropViewTask, EventContext}; +use crate::rpc::ddl::DropViewTask; /// The procedure for dropping a view. pub struct DropViewProcedure { @@ -48,13 +48,12 @@ pub struct DropViewProcedure { impl DropViewProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::DropView"; - pub fn new(task: DropViewTask, event_context: EventContext, context: DdlContext) -> Self { + pub fn new(task: DropViewTask, context: DdlContext) -> Self { Self { context, data: DropViewData { state: DropViewState::Prepare, task, - event_context, }, } } @@ -214,7 +213,7 @@ impl Procedure for DropViewProcedure { LockKey::new(lock_key) } - fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(DROP_VIEW_EVENT_TYPE) { return None; } @@ -228,7 +227,6 @@ impl Procedure for DropViewProcedure { table_ref.table, self.data.view_id(), self.data.task.drop_if_exists, - self.data.event_context.clone(), ) } _ => { @@ -251,8 +249,6 @@ impl Procedure for DropViewProcedure { pub(crate) struct DropViewData { state: DropViewState, task: DropViewTask, - #[serde(default)] - event_context: EventContext, } impl DropViewData { diff --git a/src/common/meta/src/ddl/event/database.rs b/src/common/meta/src/ddl/event/database.rs index b24611799c6..cab42020f9f 100644 --- a/src/common/meta/src/ddl/event/database.rs +++ b/src/common/meta/src/ddl/event/database.rs @@ -19,14 +19,13 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN as EVENT_TABLE_CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, - SCHEMA_NAME_COLUMN as EVENT_TABLE_SCHEMA_NAME_COLUMN, column_schemas, nullable_json, - nullable_string, + CATALOG_NAME_COLUMN as EVENT_TABLE_CATALOG_NAME_COLUMN, + SCHEMA_NAME_COLUMN as EVENT_TABLE_SCHEMA_NAME_COLUMN, column_schemas, nullable_string, }; use serde::Serialize; use snafu::ResultExt; -use crate::rpc::ddl::{AlterDatabaseKind, EventContext, SetDatabaseOption, UnsetDatabaseOption}; +use crate::rpc::ddl::{AlterDatabaseKind, SetDatabaseOption, UnsetDatabaseOption}; pub(crate) const CREATE_DATABASE_EVENT_TYPE: &str = "create_database"; pub(crate) const ALTER_DATABASE_EVENT_TYPE: &str = "alter_database"; @@ -40,7 +39,6 @@ pub(crate) struct DatabaseDdlEvent { catalog_name: Option, schema_name: Option, payload: Option, - event_context: Option, } #[derive(Debug, Serialize)] @@ -90,7 +88,6 @@ impl DatabaseDdlEvent { schema_name: &str, create_if_not_exists: bool, options: &HashMap, - event_context: EventContext, ) -> Self { let mut options = options.iter().collect::>(); options.sort_unstable_by_key(|(left, _)| *left); @@ -110,7 +107,6 @@ impl DatabaseDdlEvent { create_if_not_exists, options, }), - event_context, ) } @@ -118,7 +114,6 @@ impl DatabaseDdlEvent { catalog_name: &str, schema_name: &str, kind: &AlterDatabaseKind, - event_context: EventContext, ) -> Self { let intent = match kind { AlterDatabaseKind::SetDatabaseOptions(options) => AlterDatabaseIntent::Set { @@ -157,7 +152,6 @@ impl DatabaseDdlEvent { version: PAYLOAD_VERSION, intent, }), - event_context, ) } @@ -165,7 +159,6 @@ impl DatabaseDdlEvent { catalog_name: &str, schema_name: &str, drop_if_exists: bool, - event_context: EventContext, ) -> Self { Self::submitted( DROP_DATABASE_EVENT_TYPE, @@ -175,7 +168,6 @@ impl DatabaseDdlEvent { version: PAYLOAD_VERSION, drop_if_exists, }), - event_context, ) } @@ -196,14 +188,12 @@ impl DatabaseDdlEvent { catalog_name: &str, schema_name: &str, payload: DatabaseDdlPayload, - event_context: EventContext, ) -> Self { Self { event_type, catalog_name: Some(catalog_name.to_string()), schema_name: Some(schema_name.to_string()), payload: Some(payload), - event_context: Some(event_context), } } @@ -213,7 +203,6 @@ impl DatabaseDdlEvent { catalog_name: Some(catalog_name.to_string()), schema_name: Some(schema_name.to_string()), payload: None, - event_context: None, } } } @@ -234,22 +223,14 @@ impl Event for DatabaseDdlEvent { column_schemas([ &EVENT_TABLE_CATALOG_NAME_COLUMN, &EVENT_TABLE_SCHEMA_NAME_COLUMN, - &EVENT_CONTEXT_COLUMN, ]) } fn extra_rows(&self) -> Result> { - let event_context = self - .event_context - .as_ref() - .map(serde_json::to_value) - .transpose() - .context(SerializeEventSnafu)?; Ok(vec![Row { values: vec![ nullable_string(self.catalog_name.as_deref()), nullable_string(self.schema_name.as_deref()), - nullable_json(event_context.as_ref()), ], }]) } diff --git a/src/common/meta/src/ddl/event/flow.rs b/src/common/meta/src/ddl/event/flow.rs index 7d0d6cb539b..0a9ec2dd244 100644 --- a/src/common/meta/src/ddl/event/flow.rs +++ b/src/common/meta/src/ddl/event/flow.rs @@ -19,14 +19,12 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, column_schemas, - nullable_json, nullable_string, nullable_value, + CATALOG_NAME_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, column_schemas, nullable_string, + nullable_value, }; use serde::Serialize; use snafu::ResultExt; -use crate::rpc::ddl::EventContext; - pub(crate) const CREATE_FLOW_EVENT_TYPE: &str = "create_flow"; pub(crate) const DROP_FLOW_EVENT_TYPE: &str = "drop_flow"; @@ -71,15 +69,14 @@ pub(crate) struct FlowDdlEvent { flow_name: Option, flow_id: Option, payload: Option, - event_context: Option, } impl FlowDdlEvent { + /// Builds the bounded event emitted when creating a Flow is submitted. pub(crate) fn create_submitted( catalog_name: &str, flow_name: &str, intent: CreateFlowEventIntent, - event_context: EventContext, ) -> Self { Self { event_type: CREATE_FLOW_EVENT_TYPE, @@ -93,16 +90,15 @@ impl FlowDdlEvent { expire_after: intent.expire_after, eval_interval_secs: intent.eval_interval_secs, })), - event_context: Some(event_context), } } + /// Builds the bounded event emitted when dropping a Flow is submitted. pub(crate) fn drop_submitted( catalog_name: &str, flow_name: &str, flow_id: u32, drop_if_exists: bool, - event_context: EventContext, ) -> Self { Self { event_type: DROP_FLOW_EVENT_TYPE, @@ -113,7 +109,6 @@ impl FlowDdlEvent { version: PAYLOAD_VERSION, drop_if_exists, })), - event_context: Some(event_context), } } @@ -149,7 +144,6 @@ impl FlowDdlEvent { flow_name: Some(flow_name.to_string()), flow_id: None, payload: None, - event_context: None, } } } @@ -167,27 +161,15 @@ impl Event for FlowDdlEvent { } fn extra_schema(&self) -> Vec { - column_schemas([ - &CATALOG_NAME_COLUMN, - &FLOW_NAME_COLUMN, - &FLOW_ID_COLUMN, - &EVENT_CONTEXT_COLUMN, - ]) + column_schemas([&CATALOG_NAME_COLUMN, &FLOW_NAME_COLUMN, &FLOW_ID_COLUMN]) } fn extra_rows(&self) -> Result> { - let event_context = self - .event_context - .as_ref() - .map(serde_json::to_value) - .transpose() - .context(SerializeEventSnafu)?; Ok(vec![Row { values: vec![ nullable_string(self.catalog_name.as_deref()), nullable_string(self.flow_name.as_deref()), nullable_value(self.flow_id.map(ValueData::U32Value)), - nullable_json(event_context.as_ref()), ], }]) } diff --git a/src/common/meta/src/ddl/event/table.rs b/src/common/meta/src/ddl/event/table.rs index 9ae6f530761..670f3c7c1dc 100644 --- a/src/common/meta/src/ddl/event/table.rs +++ b/src/common/meta/src/ddl/event/table.rs @@ -21,17 +21,14 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, PHYSICAL_TABLE_ID_COLUMN, SCHEMA_NAME_COLUMN, - TABLE_ID_COLUMN, TABLE_NAME_COLUMN, column_schemas, nullable_json, nullable_string, - nullable_value, + CATALOG_NAME_COLUMN, PHYSICAL_TABLE_ID_COLUMN, SCHEMA_NAME_COLUMN, TABLE_ID_COLUMN, + TABLE_NAME_COLUMN, column_schemas, nullable_string, nullable_value, }; use serde::Serialize; use serde_json::Value as JsonValue; use snafu::ResultExt; use store_api::storage::TableId; -use crate::rpc::ddl::EventContext; - /// Current version of table DDL event payloads. pub(crate) const TABLE_DDL_PAYLOAD_VERSION: u8 = 1; @@ -215,7 +212,6 @@ pub(crate) struct TableDdlEvent { event_type: TableDdlEventType, locators: Vec, payload: Option, - event_context: Option, } impl TableDdlEvent { @@ -224,7 +220,6 @@ impl TableDdlEvent { locator: TableDdlLocator, create_if_not_exists: bool, engine: &str, - event_context: EventContext, ) -> Self { Self::submitted( TableDdlEventType::CreateTable, @@ -234,7 +229,6 @@ impl TableDdlEvent { create_if_not_exists, engine: engine.to_string(), }), - Some(event_context), ) } @@ -242,7 +236,6 @@ impl TableDdlEvent { pub(crate) fn create_logical_tables_submitted( locators: impl IntoIterator, table_count: usize, - event_context: EventContext, ) -> Self { Self::submitted( TableDdlEventType::CreateLogicalTables, @@ -251,7 +244,6 @@ impl TableDdlEvent { version: TABLE_DDL_PAYLOAD_VERSION, table_count, }), - Some(event_context), ) } @@ -259,7 +251,6 @@ impl TableDdlEvent { pub(crate) fn alter_table_submitted( locator: TableDdlLocator, kind: Option<&'static str>, - event_context: EventContext, ) -> Self { Self::submitted( TableDdlEventType::AlterTable, @@ -268,7 +259,6 @@ impl TableDdlEvent { version: TABLE_DDL_PAYLOAD_VERSION, kind, }), - Some(event_context), ) } @@ -277,7 +267,6 @@ impl TableDdlEvent { locators: impl IntoIterator, table_count: usize, kinds: impl IntoIterator, - event_context: EventContext, ) -> Self { let kinds = kinds .into_iter() @@ -292,16 +281,11 @@ impl TableDdlEvent { table_count, kinds, }), - Some(event_context), ) } /// Builds the bounded event emitted when dropping a table is submitted. - pub(crate) fn drop_table_submitted( - locator: TableDdlLocator, - drop_if_exists: bool, - event_context: EventContext, - ) -> Self { + pub(crate) fn drop_table_submitted(locator: TableDdlLocator, drop_if_exists: bool) -> Self { Self::submitted( TableDdlEventType::DropTable, [locator], @@ -309,39 +293,30 @@ impl TableDdlEvent { version: TABLE_DDL_PAYLOAD_VERSION, drop_if_exists, }), - Some(event_context), ) } /// Builds the bounded event emitted when restoring a dropped table is submitted. #[cfg(feature = "enterprise")] - pub(crate) fn undrop_table_submitted( - locator: TableDdlLocator, - event_context: EventContext, - ) -> Self { + pub(crate) fn undrop_table_submitted(locator: TableDdlLocator) -> Self { Self::submitted( TableDdlEventType::UndropTable, [locator], TableDdlPayload::UndropTable(UndropTablePayload { version: TABLE_DDL_PAYLOAD_VERSION, }), - Some(event_context), ) } /// Builds the bounded event emitted when purging a dropped table is submitted. #[cfg(feature = "enterprise")] - pub(crate) fn purge_dropped_table_submitted( - locator: TableDdlLocator, - event_context: EventContext, - ) -> Self { + pub(crate) fn purge_dropped_table_submitted(locator: TableDdlLocator) -> Self { Self::submitted( TableDdlEventType::PurgeDroppedTable, [locator], TableDdlPayload::PurgeDroppedTable(PurgeDroppedTablePayload { version: TABLE_DDL_PAYLOAD_VERSION, }), - Some(event_context), ) } @@ -349,7 +324,6 @@ impl TableDdlEvent { pub(crate) fn truncate_table_submitted( locator: TableDdlLocator, time_range_count: usize, - event_context: EventContext, ) -> Self { Self::submitted( TableDdlEventType::TruncateTable, @@ -358,7 +332,6 @@ impl TableDdlEvent { version: TABLE_DDL_PAYLOAD_VERSION, time_range_count, }), - Some(event_context), ) } @@ -371,7 +344,6 @@ impl TableDdlEvent { event_type, locators: locators.into_iter().collect(), payload: None, - event_context: None, } } @@ -394,13 +366,11 @@ impl TableDdlEvent { event_type: TableDdlEventType, locators: impl IntoIterator, payload: TableDdlPayload, - event_context: Option, ) -> Self { Self { event_type, locators: locators.into_iter().collect(), payload: Some(payload), - event_context, } } @@ -413,7 +383,7 @@ impl TableDdlEvent { ]) } - fn locator_row(&self, locator: &TableDdlLocator) -> Result { + fn locator_row(&self, locator: &TableDdlLocator) -> Row { let mut values = vec![ nullable_string(locator.catalog_name.as_deref()), nullable_string(locator.schema_name.as_deref()), @@ -423,14 +393,7 @@ impl TableDdlEvent { if self.event_type.has_physical_table_id() { values.push(nullable_table_id(locator.physical_table_id)); } - let event_context = self - .event_context - .as_ref() - .map(serde_json::to_value) - .transpose() - .context(SerializeEventSnafu)?; - values.push(nullable_json(event_context.as_ref())); - Ok(Row { values }) + Row { values } } } @@ -451,15 +414,15 @@ impl Event for TableDdlEvent { if self.event_type.has_physical_table_id() { schema.push(PHYSICAL_TABLE_ID_COLUMN.column_schema()); } - schema.push(EVENT_CONTEXT_COLUMN.column_schema()); schema } fn extra_rows(&self) -> Result> { - self.locators + Ok(self + .locators .iter() .map(|locator| self.locator_row(locator)) - .collect() + .collect()) } fn as_any(&self) -> &dyn Any { diff --git a/src/common/meta/src/ddl/event/view.rs b/src/common/meta/src/ddl/event/view.rs index c1d5846fcf2..cdce5954086 100644 --- a/src/common/meta/src/ddl/event/view.rs +++ b/src/common/meta/src/ddl/event/view.rs @@ -19,14 +19,12 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, - VIEW_NAME_COLUMN, column_schemas, nullable_json, nullable_string, nullable_value, + CATALOG_NAME_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, VIEW_NAME_COLUMN, column_schemas, + nullable_string, nullable_value, }; use serde::Serialize; use snafu::ResultExt; -use crate::rpc::ddl::EventContext; - pub(crate) const CREATE_VIEW_EVENT_TYPE: &str = "create_view"; pub(crate) const DROP_VIEW_EVENT_TYPE: &str = "drop_view"; @@ -64,7 +62,6 @@ pub(crate) struct ViewDdlEvent { view_name: Option, view_id: Option, payload: Option, - event_context: Option, } #[derive(Debug, Serialize)] @@ -75,12 +72,12 @@ enum ViewDdlPayload { } impl ViewDdlEvent { + /// Builds the bounded event emitted when creating a View is submitted. pub(crate) fn create_submitted( catalog_name: &str, schema_name: &str, view_name: &str, intent: CreateViewEventIntent, - event_context: EventContext, ) -> Self { Self::submitted( CREATE_VIEW_EVENT_TYPE, @@ -95,17 +92,16 @@ impl ViewDdlEvent { referenced_table_count: intent.referenced_table_count, column_count: intent.column_count, }), - event_context, ) } + /// Builds the bounded event emitted when dropping a View is submitted. pub(crate) fn drop_submitted( catalog_name: &str, schema_name: &str, view_name: &str, view_id: u32, drop_if_exists: bool, - event_context: EventContext, ) -> Self { Self::submitted( DROP_VIEW_EVENT_TYPE, @@ -117,7 +113,6 @@ impl ViewDdlEvent { version: PAYLOAD_VERSION, drop_if_exists, }), - event_context, ) } @@ -162,7 +157,6 @@ impl ViewDdlEvent { view_name: &str, view_id: Option, payload: ViewDdlPayload, - event_context: EventContext, ) -> Self { Self { event_type, @@ -171,7 +165,6 @@ impl ViewDdlEvent { view_name: Some(view_name.to_string()), view_id, payload: Some(payload), - event_context: Some(event_context), } } @@ -188,7 +181,6 @@ impl ViewDdlEvent { view_name: Some(view_name.to_string()), view_id: None, payload: None, - event_context: None, } } @@ -206,7 +198,6 @@ impl ViewDdlEvent { view_name: Some(view_name.to_string()), view_id: Some(view_id), payload: None, - event_context: None, } } } @@ -229,24 +220,16 @@ impl Event for ViewDdlEvent { &SCHEMA_NAME_COLUMN, &VIEW_NAME_COLUMN, &VIEW_ID_COLUMN, - &EVENT_CONTEXT_COLUMN, ]) } fn extra_rows(&self) -> Result> { - let event_context = self - .event_context - .as_ref() - .map(serde_json::to_value) - .transpose() - .context(SerializeEventSnafu)?; Ok(vec![Row { values: vec![ nullable_string(self.catalog_name.as_deref()), nullable_string(self.schema_name.as_deref()), nullable_string(self.view_name.as_deref()), nullable_value(self.view_id.map(ValueData::U32Value)), - nullable_json(event_context.as_ref()), ], }]) } diff --git a/src/common/meta/src/ddl/purge_dropped_table.rs b/src/common/meta/src/ddl/purge_dropped_table.rs index c81aa043dc2..cdafe40de81 100644 --- a/src/common/meta/src/ddl/purge_dropped_table.rs +++ b/src/common/meta/src/ddl/purge_dropped_table.rs @@ -14,7 +14,7 @@ use std::collections::HashMap; use async_trait::async_trait; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, Result as ProcedureResult, Status, }; use common_telemetry::info; @@ -38,7 +38,7 @@ use crate::error::{self, Result}; use crate::key::DroppedTableMetadata; use crate::key::table_route::TableRouteValue; use crate::lock_key::TableLock; -use crate::rpc::ddl::{EventContext, PurgeDroppedTableTask}; +use crate::rpc::ddl::PurgeDroppedTableTask; use crate::rpc::router::RegionRoute; pub struct PurgeDroppedTableProcedure { @@ -50,25 +50,17 @@ impl PurgeDroppedTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::PurgeDroppedTable"; pub const EXPIRED_TYPE_NAME: &'static str = "metasrv-procedure::PurgeExpiredDroppedTable"; - pub fn new( - task: PurgeDroppedTableTask, - context: DdlContext, - event_context: EventContext, - ) -> Self { + pub fn new(task: PurgeDroppedTableTask, context: DdlContext) -> Self { Self { context, - data: PurgeDroppedTableData::new(task, event_context), + data: PurgeDroppedTableData::new(task), } } - pub fn new_if_expired( - task: PurgeDroppedTableTask, - context: DdlContext, - event_context: EventContext, - ) -> Self { + pub fn new_if_expired(task: PurgeDroppedTableTask, context: DdlContext) -> Self { Self { context, - data: PurgeDroppedTableData::new_if_expired(task, event_context), + data: PurgeDroppedTableData::new_if_expired(task), } } @@ -278,10 +270,7 @@ impl Procedure for PurgeDroppedTableProcedure { LockKey::new(vec![TableLock::Write(self.data.task.table_id).into()]) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::PurgeDroppedTable.as_str()) @@ -290,7 +279,7 @@ impl Procedure for PurgeDroppedTableProcedure { } let locator = self.event_locator(); let event = if ctx.trigger == EventTrigger::Submitted { - TableDdlEvent::purge_dropped_table_submitted(locator, self.data.event_context.clone()) + TableDdlEvent::purge_dropped_table_submitted(locator) } else { TableDdlEvent::lifecycle(TableDdlEventType::PurgeDroppedTable, [locator]) }; @@ -316,12 +305,10 @@ pub struct PurgeDroppedTableData { drop_generation_loaded: bool, #[serde(default)] purging_claimed: bool, - #[serde(default)] - event_context: EventContext, } impl PurgeDroppedTableData { - fn new(task: PurgeDroppedTableTask, event_context: EventContext) -> Self { + fn new(task: PurgeDroppedTableTask) -> Self { Self { state: PurgeDroppedTableState::Prepare, task, @@ -334,14 +321,13 @@ impl PurgeDroppedTableData { drop_generation: None, drop_generation_loaded: false, purging_claimed: false, - event_context, } } - fn new_if_expired(task: PurgeDroppedTableTask, event_context: EventContext) -> Self { + fn new_if_expired(task: PurgeDroppedTableTask) -> Self { Self { check_expired: true, - ..Self::new(task, event_context) + ..Self::new(task) } } diff --git a/src/common/meta/src/ddl/test_util.rs b/src/common/meta/src/ddl/test_util.rs index 3686ceb33d2..7d8cdd895b2 100644 --- a/src/common/meta/src/ddl/test_util.rs +++ b/src/common/meta/src/ddl/test_util.rs @@ -45,7 +45,7 @@ use crate::key::node_address::{NodeAddressKey, NodeAddressValue}; use crate::key::table_route::TableRouteValue; use crate::key::{MetadataKey, MetadataValue}; use crate::peer::Peer; -use crate::rpc::ddl::{CreateTableTask, EventContext}; +use crate::rpc::ddl::CreateTableTask; use crate::rpc::store::PutRequest; pub async fn create_physical_table_metadata( @@ -104,12 +104,7 @@ pub async fn create_logical_table( table_name: &str, ) -> TableId { let tasks = vec![test_create_logical_table_task(table_name)]; - let mut procedure = CreateLogicalTablesProcedure::new( - tasks, - physical_table_id, - EventContext::default(), - ddl_context, - ); + let mut procedure = CreateLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); let status = procedure.on_prepare().await.unwrap(); assert_matches!( status, diff --git a/src/common/meta/src/ddl/tests/alter_logical_tables.rs b/src/common/meta/src/ddl/tests/alter_logical_tables.rs index 09662199389..351c0fd96ec 100644 --- a/src/common/meta/src/ddl/tests/alter_logical_tables.rs +++ b/src/common/meta/src/ddl/tests/alter_logical_tables.rs @@ -43,7 +43,7 @@ use crate::error::Error::{AlterLogicalTablesInvalidArguments, TableNotFound}; use crate::error::Result; use crate::key::table_name::TableNameKey; use crate::key::table_route::{PhysicalTableRouteValue, TableRouteValue}; -use crate::rpc::ddl::{AlterTableTask, EventContext}; +use crate::rpc::ddl::AlterTableTask; use crate::rpc::router::{Region, RegionRoute}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -162,12 +162,7 @@ async fn test_on_prepare_check_schema() { ), ]; let physical_table_id = 1024u32; - let mut procedure = AlterLogicalTablesProcedure::new( - tasks, - physical_table_id, - EventContext::default(), - ddl_context, - ); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, AlterLogicalTablesInvalidArguments { .. }); } @@ -182,12 +177,7 @@ async fn test_on_prepare_check_alter_kind() { "new_table1", )]; let physical_table_id = 1024u32; - let mut procedure = AlterLogicalTablesProcedure::new( - tasks, - physical_table_id, - EventContext::default(), - ddl_context, - ); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, AlterLogicalTablesInvalidArguments { .. }); } @@ -207,8 +197,7 @@ async fn test_on_prepare_different_physical_table() { make_alter_logical_table_add_column_task(None, "table2", vec!["column2".to_string()]), ]; - let mut procedure = - AlterLogicalTablesProcedure::new(tasks, phy1_id, EventContext::default(), ddl_context); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy1_id, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, AlterLogicalTablesInvalidArguments { .. }); } @@ -229,8 +218,7 @@ async fn test_on_prepare_logical_table_not_exists() { make_alter_logical_table_add_column_task(None, "table2", vec!["column2".to_string()]), ]; - let mut procedure = - AlterLogicalTablesProcedure::new(tasks, phy_id, EventContext::default(), ddl_context); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, TableNotFound { .. }); } @@ -253,8 +241,7 @@ async fn test_on_prepare() { make_alter_logical_table_add_column_task(None, "table3", vec!["column3".to_string()]), ]; - let mut procedure = - AlterLogicalTablesProcedure::new(tasks, phy_id, EventContext::default(), ddl_context); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context); let result = procedure.on_prepare().await; assert_matches!( result, @@ -290,12 +277,7 @@ async fn test_on_update_metadata() { make_alter_logical_table_add_column_task(None, "table3", vec!["new_col".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new( - tasks, - phy_id, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context.clone()); let mut status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -308,6 +290,7 @@ async fn test_on_update_metadata() { let ctx = common_procedure::Context { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; // on_submit_alter_region_requests status = procedure.execute(&ctx).await.unwrap(); @@ -379,12 +362,7 @@ async fn test_on_part_duplicate_alter_request() { make_alter_logical_table_add_column_task(None, "table2", vec!["col_0".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new( - tasks, - phy_id, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context.clone()); let mut status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -397,6 +375,7 @@ async fn test_on_part_duplicate_alter_request() { let ctx = common_procedure::Context { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; // on_submit_alter_region_requests status = procedure.execute(&ctx).await.unwrap(); @@ -469,12 +448,7 @@ async fn test_on_part_duplicate_alter_request() { ), ]; - let mut procedure = AlterLogicalTablesProcedure::new( - tasks, - phy_id, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context.clone()); let mut status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -487,6 +461,7 @@ async fn test_on_part_duplicate_alter_request() { let ctx = common_procedure::Context { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; // on_submit_alter_region_requests status = procedure.execute(&ctx).await.unwrap(); @@ -645,8 +620,7 @@ async fn test_on_submit_alter_region_request() { make_alter_logical_table_add_column_task(None, "table2", vec!["mew_col".to_string()]), ]; - let mut procedure = - AlterLogicalTablesProcedure::new(tasks, phy_id, EventContext::default(), ddl_context); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context); procedure.on_prepare().await.unwrap(); procedure.on_submit_alter_region_requests().await.unwrap(); let mut results = Vec::new(); diff --git a/src/common/meta/src/ddl/tests/alter_table.rs b/src/common/meta/src/ddl/tests/alter_table.rs index d68d1c96dfa..b35666e52e1 100644 --- a/src/common/meta/src/ddl/tests/alter_table.rs +++ b/src/common/meta/src/ddl/tests/alter_table.rs @@ -59,7 +59,7 @@ use crate::lock_key::RegionLock; use crate::node_manager::NodeManagerRef; use crate::peer::Peer; use crate::poison_key::table_poison_key; -use crate::rpc::ddl::{AlterTableTask, EventContext}; +use crate::rpc::ddl::AlterTableTask; use crate::rpc::router::{Region, RegionRoute}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -121,8 +121,7 @@ async fn test_on_prepare_table_exists_err() { .unwrap(); let task = test_rename_alter_table_task("non-exists", "foo"); - let mut procedure = - AlterTableProcedure::new(1024, task, EventContext::default(), ddl_context).unwrap(); + let mut procedure = AlterTableProcedure::new(1024, task, ddl_context).unwrap(); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err.status_code(), StatusCode::TableAlreadyExists); } @@ -132,8 +131,7 @@ async fn test_on_prepare_table_not_exists_err() { let node_manager = Arc::new(MockDatanodeManager::new(())); let ddl_context = new_ddl_context(node_manager); let task = test_rename_alter_table_task("non-exists", "foo"); - let mut procedure = - AlterTableProcedure::new(1024, task, EventContext::default(), ddl_context).unwrap(); + let mut procedure = AlterTableProcedure::new(1024, task, ddl_context).unwrap(); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err.status_code(), StatusCode::TableNotFound); } @@ -247,13 +245,8 @@ async fn test_on_submit_alter_request() { let alter_table_task = test_alter_table_task(table_name); let procedure_id = ProcedureId::random(); let provider = Arc::new(MockContextProvider::default()); - let mut procedure = AlterTableProcedure::new( - table_id, - alter_table_task, - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut procedure = + AlterTableProcedure::new(table_id, alter_table_task, ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); procedure .submit_alter_region_requests(procedure_id, provider.as_ref()) @@ -306,13 +299,8 @@ async fn test_on_submit_alter_request_without_sync_request() { let alter_table_task = test_alter_table_task(table_name); let procedure_id = ProcedureId::random(); let provider = Arc::new(MockContextProvider::default()); - let mut procedure = AlterTableProcedure::new( - table_id, - alter_table_task, - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut procedure = + AlterTableProcedure::new(table_id, alter_table_task, ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); procedure .submit_alter_region_requests(procedure_id, provider.as_ref()) @@ -369,13 +357,7 @@ async fn test_on_submit_alter_request_with_outdated_request() { }; let procedure_id = ProcedureId::random(); let provider = Arc::new(MockContextProvider::default()); - let mut procedure = AlterTableProcedure::new( - table_id, - alter_table_task, - EventContext::default(), - ddl_context, - ) - .unwrap(); + let mut procedure = AlterTableProcedure::new(table_id, alter_table_task, ddl_context).unwrap(); procedure.on_prepare().await.unwrap(); let err = procedure .submit_alter_region_requests(procedure_id, provider.as_ref()) @@ -404,9 +386,7 @@ async fn test_on_update_metadata_rename() { .unwrap(); let task = test_rename_alter_table_task(table_name, new_table_name); - let mut procedure = - AlterTableProcedure::new(table_id, task, EventContext::default(), ddl_context.clone()) - .unwrap(); + let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); procedure.on_update_metadata().await.unwrap(); @@ -486,9 +466,7 @@ async fn test_on_update_metadata_add_columns() { }; let procedure_id = ProcedureId::random(); let provider = Arc::new(MockContextProvider::default()); - let mut procedure = - AlterTableProcedure::new(table_id, task, EventContext::default(), ddl_context.clone()) - .unwrap(); + let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); procedure .submit_alter_region_requests(procedure_id, provider.as_ref()) @@ -585,9 +563,7 @@ async fn test_on_update_table_options() { }; let procedure_id = ProcedureId::random(); let provider = Arc::new(MockContextProvider::default()); - let mut procedure = - AlterTableProcedure::new(table_id, task, EventContext::default(), ddl_context.clone()) - .unwrap(); + let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); procedure .submit_alter_region_requests(procedure_id, provider.as_ref()) @@ -658,7 +634,6 @@ async fn test_skip_wal_rejects_mixed_table_options() { })), }, }, - EventContext::default(), ddl_context, ) .unwrap(); @@ -687,13 +662,8 @@ async fn test_skip_wal_rejects_logical_table() { }, }; - let mut procedure = AlterTableProcedure::new( - logical_table_id, - task, - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut procedure = + AlterTableProcedure::new(logical_table_id, task, ddl_context.clone()).unwrap(); let error = procedure.on_prepare().await.unwrap_err(); assert_matches!(error, Error::Unsupported { .. }); @@ -739,9 +709,7 @@ async fn test_skip_wal_rejects_file_engine_table() { }, }; - let mut procedure = - AlterTableProcedure::new(table_id, task, EventContext::default(), ddl_context.clone()) - .unwrap(); + let mut procedure = AlterTableProcedure::new(table_id, task, ddl_context.clone()).unwrap(); let error = procedure.on_prepare().await.unwrap_err(); assert_matches!(error, Error::Unsupported { .. }); @@ -778,7 +746,6 @@ fn test_skip_wal_holds_region_locks() { let procedure = AlterTableProcedure::new_with_region_locks( table_id, task, - EventContext::default(), region_ids.clone(), context.clone(), ) @@ -830,7 +797,6 @@ async fn test_skip_wal_detects_region_route_change() { let mut procedure = AlterTableProcedure::new_with_region_locks( table_id, alter_task, - EventContext::default(), stale_region_locks, ddl_context.clone(), ) @@ -890,7 +856,6 @@ async fn test_skip_wal_rejects_no_leader_before_updating_metadata() { let mut procedure = AlterTableProcedure::new_with_region_locks( table_id, alter_task, - EventContext::default(), region_locks, ddl_context.clone(), ) @@ -951,7 +916,6 @@ async fn test_skip_wal_updates_metadata_before_all_replicas() { let mut procedure = AlterTableProcedure::new_with_region_locks( table_id, alter_task, - EventContext::default(), region_locks, ddl_context.clone(), ) @@ -1033,13 +997,8 @@ async fn prepare_alter_table_procedure( }, }; let procedure_id = ProcedureId::random(); - let mut procedure = AlterTableProcedure::new( - table_id, - alter_table_task, - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut procedure = + AlterTableProcedure::new(table_id, alter_table_task, ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); (procedure, procedure_id) } diff --git a/src/common/meta/src/ddl/tests/create_flow.rs b/src/common/meta/src/ddl/tests/create_flow.rs index 280a2633e32..b526856bdac 100644 --- a/src/common/meta/src/ddl/tests/create_flow.rs +++ b/src/common/meta/src/ddl/tests/create_flow.rs @@ -37,7 +37,7 @@ use crate::key::flow::flow_info::{ FlowInfoValue, FlowMissedTickPolicy, FlowScheduleConfig, FlowStatus, }; use crate::key::table_route::TableRouteValue; -use crate::rpc::ddl::{CreateFlowTask, EventContext, FlowQueryContext, QueryContext}; +use crate::rpc::ddl::{CreateFlowTask, FlowQueryContext, QueryContext}; use crate::test_util::{MockFlownodeManager, new_ddl_context}; pub(crate) fn test_query_context() -> QueryContext { @@ -120,8 +120,7 @@ async fn test_create_flow_source_table_not_found() { let node_manager = Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler)); let ddl_context = new_ddl_context(node_manager); let query_ctx = test_query_context(); - let mut procedure = - CreateFlowProcedure::new(task, query_ctx, EventContext::default(), ddl_context); + let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::Unsupported { .. }); assert!( @@ -144,12 +143,7 @@ async fn test_create_pending_flow_source_table_not_found_with_defer() { let node_manager = Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler)); let ddl_context = new_ddl_context(node_manager); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new( - task, - query_ctx, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context.clone()); let status = procedure.on_prepare().await.unwrap(); assert_matches!(status, Status::Executing { persist: true, .. }); assert_eq!(procedure.data.unresolved_source_table_names.len(), 1); @@ -189,8 +183,7 @@ async fn test_create_pending_flow_source_table_not_found_with_defer_false() { let node_manager = Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler)); let ddl_context = new_ddl_context(node_manager); let query_ctx = test_query_context(); - let mut procedure = - CreateFlowProcedure::new(task, query_ctx, EventContext::default(), ddl_context); + let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::Unsupported { .. }); assert!( @@ -240,12 +233,7 @@ async fn test_create_pending_flow_records_partial_source_resolution() { ); enable_defer_on_missing_source(&mut task); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new( - task, - query_ctx, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context.clone()); let status = procedure.on_prepare().await.unwrap(); assert_matches!(status, Status::Executing { persist: true, .. }); assert_eq!(procedure.data.source_table_ids, vec![existing_table_id]); @@ -440,7 +428,6 @@ fn test_resolved_schedule_defaults_in_create_request() { // Construct CreateFlowData and verify CreateRequest carries internal key. let data = CreateFlowData { - event_context: EventContext::default(), state: CreateFlowState::CreateFlows, task, flow_id: Some(1024), @@ -454,7 +441,6 @@ fn test_resolved_schedule_defaults_in_create_request() { }; let data2 = CreateFlowData { - event_context: EventContext::default(), state: CreateFlowState::CreateMetadata, task: data.task.clone(), flow_id: Some(1024), @@ -594,8 +580,7 @@ async fn test_create_flow_rejects_unknown_option_in_meta_task() { let node_manager = Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler)); let ddl_context = new_ddl_context(node_manager); let query_ctx = test_query_context(); - let mut procedure = - CreateFlowProcedure::new(task, query_ctx, EventContext::default(), ddl_context); + let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::Unexpected { .. }); @@ -616,7 +601,6 @@ fn test_create_request_strips_defer_on_missing_source_runtime_option() { enable_defer_on_missing_source(&mut task); let data = CreateFlowData { - event_context: EventContext::default(), state: CreateFlowState::CreateFlows, task, flow_id: Some(1024), @@ -666,12 +650,7 @@ pub(crate) async fn create_test_flow( false, ); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new( - task.clone(), - query_ctx, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = CreateFlowProcedure::new(task.clone(), query_ctx, ddl_context.clone()); let output = execute_procedure_until_done(&mut procedure).await.unwrap(); let flow_id = output.downcast_ref::().unwrap(); @@ -692,12 +671,7 @@ pub(crate) async fn create_test_pending_flow( ); enable_defer_on_missing_source(&mut task); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new( - task, - query_ctx, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context.clone()); let output = execute_procedure_until_done(&mut procedure).await.unwrap(); let flow_id = output.downcast_ref::().unwrap(); @@ -744,12 +718,7 @@ async fn test_create_flow() { true, ); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new( - task.clone(), - query_ctx, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = CreateFlowProcedure::new(task.clone(), query_ctx, ddl_context.clone()); let output = execute_procedure_until_done(&mut procedure).await.unwrap(); let flow_id = output.downcast_ref::().unwrap(); assert_eq!(*flow_id, 1024); @@ -757,12 +726,7 @@ async fn test_create_flow() { // Creates again let task = test_create_flow_task("my_flow", source_table_names, sink_table_name, false); let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new( - task.clone(), - query_ctx, - EventContext::default(), - ddl_context, - ); + let mut procedure = CreateFlowProcedure::new(task.clone(), query_ctx, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::FlowAlreadyExists { .. }); } @@ -819,12 +783,7 @@ async fn test_replace_pending_flow_with_active_flow_is_unsupported() { ); replace_task.or_replace = true; let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new( - replace_task, - query_ctx, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = CreateFlowProcedure::new(replace_task, query_ctx, ddl_context.clone()); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::Unsupported { .. }); assert!( @@ -882,12 +841,7 @@ async fn test_replace_active_flow_with_pending_flow_is_unsupported() { enable_defer_on_missing_source(&mut replace_task); replace_task.or_replace = true; let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new( - replace_task, - query_ctx, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = CreateFlowProcedure::new(replace_task, query_ctx, ddl_context.clone()); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::Unsupported { .. }); assert!( @@ -947,12 +901,7 @@ async fn test_replace_pending_flow_with_pending_flow_updates_metadata() { enable_defer_on_missing_source(&mut replace_task); replace_task.or_replace = true; let query_ctx = test_query_context(); - let mut procedure = CreateFlowProcedure::new( - replace_task, - query_ctx, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = CreateFlowProcedure::new(replace_task, query_ctx, ddl_context.clone()); let output = execute_procedure_until_done(&mut procedure).await.unwrap(); let replaced_flow_id = *output.downcast_ref::().unwrap(); assert_eq!(replaced_flow_id, original_flow_id); @@ -1004,8 +953,7 @@ async fn test_create_flow_same_source_and_sink_table() { // Try to create a flow with same source and sink table - should fail let task = test_create_flow_task("my_flow", source_table_names, sink_table_name, false); let query_ctx = test_query_context(); - let mut procedure = - CreateFlowProcedure::new(task, query_ctx, EventContext::default(), ddl_context); + let mut procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::Unsupported { .. }); @@ -1038,7 +986,7 @@ fn test_create_flow_lock_key_does_not_lock_sink_table_name() { let task = create_test_flow_task_for_serialization(); let query_ctx = test_query_context(); let ddl_context = new_ddl_context(Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler))); - let procedure = CreateFlowProcedure::new(task, query_ctx, EventContext::default(), ddl_context); + let procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); let lock_keys = procedure.lock_key().get_keys(); @@ -1113,7 +1061,6 @@ fn test_create_flow_data_new_format_serialization() { }; let data = CreateFlowData { - event_context: EventContext::default(), state: CreateFlowState::Prepare, task: create_test_flow_task_for_serialization(), flow_id: None, @@ -1182,7 +1129,7 @@ fn test_create_flow_procedure_strips_scheduled_time_extension() { .extensions .insert("flow.other".to_string(), "kept".to_string()); - let procedure = CreateFlowProcedure::new(task, query_ctx, EventContext::default(), ddl_context); + let procedure = CreateFlowProcedure::new(task, query_ctx, ddl_context); assert!( !procedure @@ -1215,7 +1162,6 @@ fn test_flow_info_conversion_with_flow_context() { }; let data = CreateFlowData { - event_context: EventContext::default(), state: CreateFlowState::CreateMetadata, task: create_test_flow_task_for_serialization(), flow_id: Some(123), @@ -1258,7 +1204,6 @@ fn test_flow_info_conversion_strips_scheduled_time_extension() { }; let data = CreateFlowData { - event_context: EventContext::default(), state: CreateFlowState::CreateMetadata, task: create_test_flow_task_for_serialization(), flow_id: Some(123), diff --git a/src/common/meta/src/ddl/tests/create_logical_tables.rs b/src/common/meta/src/ddl/tests/create_logical_tables.rs index d1cc2420e04..89ec486ee03 100644 --- a/src/common/meta/src/ddl/tests/create_logical_tables.rs +++ b/src/common/meta/src/ddl/tests/create_logical_tables.rs @@ -39,7 +39,6 @@ use crate::ddl::test_util::{ }; use crate::error::{Error, Result}; use crate::key::table_route::{PhysicalTableRouteValue, TableRouteValue}; -use crate::rpc::ddl::EventContext; use crate::rpc::router::{Region, RegionRoute}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -97,12 +96,7 @@ async fn test_on_prepare_physical_table_not_found() { let ddl_context = new_ddl_context(node_manager); let tasks = vec![test_create_logical_table_task("foo")]; let physical_table_id = 1024u32; - let mut procedure = CreateLogicalTablesProcedure::new( - tasks, - physical_table_id, - EventContext::default(), - ddl_context, - ); + let mut procedure = CreateLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::TableRouteNotFound { .. }); } @@ -132,12 +126,7 @@ async fn test_on_prepare() { // The create logical table procedure. let tasks = vec![test_create_logical_table_task("foo")]; let physical_table_id = table_id; - let mut procedure = CreateLogicalTablesProcedure::new( - tasks, - physical_table_id, - EventContext::default(), - ddl_context, - ); + let mut procedure = CreateLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); let status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -183,12 +172,8 @@ async fn test_on_prepare_logical_table_exists_err() { .unwrap(); // The create logical table procedure. let physical_table_id = table_id; - let mut procedure = CreateLogicalTablesProcedure::new( - vec![task], - physical_table_id, - EventContext::default(), - ddl_context, - ); + let mut procedure = + CreateLogicalTablesProcedure::new(vec![task], physical_table_id, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::TableAlreadyExists { .. }); assert_eq!(err.status_code(), StatusCode::TableAlreadyExists); @@ -231,12 +216,8 @@ async fn test_on_prepare_with_create_if_table_exists() { let physical_table_id = table_id; // Sets `create_if_not_exists` task.create_table.create_if_not_exists = true; - let mut procedure = CreateLogicalTablesProcedure::new( - vec![task], - physical_table_id, - EventContext::default(), - ddl_context, - ); + let mut procedure = + CreateLogicalTablesProcedure::new(vec![task], physical_table_id, ddl_context); let status = procedure.on_prepare().await.unwrap(); let output = status.downcast_output_ref::>().unwrap(); assert_eq!(*output, vec![8192]); @@ -283,7 +264,6 @@ async fn test_on_prepare_part_logical_tables_exist() { let mut procedure = CreateLogicalTablesProcedure::new( vec![task, non_exist_task], physical_table_id, - EventContext::default(), ddl_context, ); let status = procedure.on_prepare().await.unwrap(); @@ -331,7 +311,6 @@ async fn test_on_create_metadata() { let mut procedure = CreateLogicalTablesProcedure::new( vec![task, yet_another_task], physical_table_id, - EventContext::default(), ddl_context.clone(), ); let status = procedure.on_prepare().await.unwrap(); @@ -345,6 +324,7 @@ async fn test_on_create_metadata() { let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; procedure.execute(&ctx).await.unwrap(); // Triggers procedure to create table metadata @@ -425,7 +405,6 @@ async fn test_on_create_metadata_part_logical_tables_exist() { let mut procedure = CreateLogicalTablesProcedure::new( vec![task, non_exist_task], physical_table_id, - EventContext::default(), ddl_context.clone(), ); let status = procedure.on_prepare().await.unwrap(); @@ -439,6 +418,7 @@ async fn test_on_create_metadata_part_logical_tables_exist() { let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; procedure.execute(&ctx).await.unwrap(); // Triggers procedure to create table metadata @@ -498,7 +478,6 @@ async fn test_on_create_metadata_err() { let mut procedure = CreateLogicalTablesProcedure::new( vec![task.clone(), yet_another_task], physical_table_id, - EventContext::default(), ddl_context.clone(), ); let status = procedure.on_prepare().await.unwrap(); @@ -512,6 +491,7 @@ async fn test_on_create_metadata_err() { let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; procedure.execute(&ctx).await.unwrap(); // Creates logical table metadata(different with the task) @@ -559,7 +539,6 @@ async fn test_on_submit_create_request() { let mut procedure = CreateLogicalTablesProcedure::new( vec![task, yet_another_task], physical_table_id, - EventContext::default(), ddl_context, ); procedure.on_prepare().await.unwrap(); diff --git a/src/common/meta/src/ddl/tests/create_table.rs b/src/common/meta/src/ddl/tests/create_table.rs index 2eee317dc14..c97e61926c0 100644 --- a/src/common/meta/src/ddl/tests/create_table.rs +++ b/src/common/meta/src/ddl/tests/create_table.rs @@ -48,7 +48,7 @@ use crate::ddl::test_util::{assert_column_name, get_raw_table_info, put_datanode use crate::error::{Error, Result}; use crate::key::table_route::TableRouteValue; use crate::kv_backend::memory::MemoryKvBackend; -use crate::rpc::ddl::{CreateTableTask, EventContext, QueryContext}; +use crate::rpc::ddl::CreateTableTask; use crate::test_util::{MockDatanodeManager, new_ddl_context, new_ddl_context_with_kv_backend}; fn create_request_handler(_peer: Peer, request: RegionRequest) -> Result { @@ -189,13 +189,7 @@ async fn test_on_prepare_table_exists_err() { ) .await .unwrap(); - let mut procedure = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - ddl_context, - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task, ddl_context).unwrap(); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::TableAlreadyExists { .. }); assert_eq!(err.status_code(), StatusCode::TableAlreadyExists); @@ -218,13 +212,7 @@ async fn test_on_prepare_with_create_if_table_exists() { ) .await .unwrap(); - let mut procedure = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - ddl_context, - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task, ddl_context).unwrap(); let status = procedure.on_prepare().await.unwrap(); assert_matches!(status, Status::Done { output: Some(..) }); let table_id = *status.downcast_output_ref::().unwrap(); @@ -237,13 +225,7 @@ async fn test_on_prepare_without_create_if_table_exists() { let ddl_context = new_ddl_context(node_manager); let mut task = test_create_table_task("foo"); task.create_table.create_if_not_exists = true; - let mut procedure = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - ddl_context, - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task, ddl_context).unwrap(); let status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -262,17 +244,12 @@ async fn test_on_datanode_create_regions_should_retry() { let ddl_context = new_ddl_context(node_manager); let task = test_create_table_task("foo"); assert!(!task.create_table.create_if_not_exists); - let mut procedure = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - ddl_context, - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task, ddl_context).unwrap(); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; let error = procedure.execute(&ctx).await.unwrap_err(); assert!(error.is_retry_later()); @@ -285,17 +262,12 @@ async fn test_on_datanode_create_regions_should_not_retry() { let ddl_context = new_ddl_context(node_manager); let task = test_create_table_task("foo"); assert!(!task.create_table.create_if_not_exists); - let mut procedure = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - ddl_context, - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task, ddl_context).unwrap(); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; let error = procedure.execute(&ctx).await.unwrap_err(); assert!(!error.is_retry_later()); @@ -308,13 +280,7 @@ async fn test_on_datanode_create_regions_remaps_addresses_when_retrying() { let node_manager = Arc::new(MockDatanodeManager::new(datanode_handler)); let ddl_context = new_ddl_context(node_manager); let task = test_create_table_task("foo"); - let mut procedure = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task, ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); let table_route = procedure.data.table_route.as_mut().unwrap(); @@ -335,17 +301,12 @@ async fn test_on_create_metadata_error() { let ddl_context = new_ddl_context(node_manager); let task = test_create_table_task("foo"); assert!(!task.create_table.create_if_not_exists); - let mut procedure = CreateTableProcedure::new( - task.clone(), - QueryContext::default(), - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task.clone(), ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; procedure.execute(&ctx).await.unwrap(); let mut task = task; @@ -374,17 +335,12 @@ async fn test_on_create_metadata() { let ddl_context = new_ddl_context(node_manager); let task = test_create_table_task("foo"); assert!(!task.create_table.create_if_not_exists); - let mut procedure = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task, ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; procedure.execute(&ctx).await.unwrap(); // Triggers procedure to create table metadata @@ -408,13 +364,7 @@ async fn test_memory_region_keeper_guard_dropped_on_procedure_done() { let ddl_context = new_ddl_context_with_kv_backend(node_manager, kv_backend); let task = test_create_table_task("foo"); - let mut procedure = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task, ddl_context.clone()).unwrap(); execute_procedure_until(&mut procedure, |p| { p.data.state == CreateTableState::CreateMetadata diff --git a/src/common/meta/src/ddl/tests/create_view.rs b/src/common/meta/src/ddl/tests/create_view.rs index 618249fd1cd..b73977c79f2 100644 --- a/src/common/meta/src/ddl/tests/create_view.rs +++ b/src/common/meta/src/ddl/tests/create_view.rs @@ -29,7 +29,7 @@ use crate::ddl::create_view::CreateViewProcedure; use crate::ddl::test_util::datanode_handler::NaiveDatanodeHandler; use crate::ddl::tests::create_table::test_create_table_task; use crate::error::Error; -use crate::rpc::ddl::{CreateViewTask, EventContext, QueryContext}; +use crate::rpc::ddl::CreateViewTask; use crate::test_util::{MockDatanodeManager, new_ddl_context}; pub(crate) fn test_table_names() -> HashSet { @@ -112,7 +112,7 @@ async fn test_on_prepare_view_exists_err() { ) .await .unwrap(); - let mut procedure = CreateViewProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = CreateViewProcedure::new(task, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::ViewAlreadyExists { .. }); assert_eq!(err.status_code(), StatusCode::TableAlreadyExists); @@ -138,7 +138,7 @@ async fn test_on_prepare_with_create_if_view_exists() { ) .await .unwrap(); - let mut procedure = CreateViewProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = CreateViewProcedure::new(task, ddl_context); let status = procedure.on_prepare().await.unwrap(); assert_matches!(status, Status::Done { output: Some(..) }); let table_id = *status.downcast_output_ref::().unwrap(); @@ -151,7 +151,7 @@ async fn test_on_prepare_without_create_if_table_exists() { let ddl_context = new_ddl_context(node_manager); let mut task = test_create_view_task("foo"); task.create_view.create_if_not_exists = true; - let mut procedure = CreateViewProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = CreateViewProcedure::new(task, ddl_context); let status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -170,11 +170,12 @@ async fn test_on_create_metadata() { let ddl_context = new_ddl_context(node_manager); let task = test_create_view_task("foo"); assert!(!task.create_view.create_if_not_exists); - let mut procedure = CreateViewProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = CreateViewProcedure::new(task, ddl_context); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; // Triggers procedure to create view metadata let status = procedure.execute(&ctx).await.unwrap(); @@ -189,12 +190,12 @@ async fn test_replace_view_metadata() { let ddl_context = new_ddl_context(node_manager.clone()); let task = test_create_view_task("foo"); assert!(!task.create_view.create_if_not_exists); - let mut procedure = - CreateViewProcedure::new(task.clone(), EventContext::default(), ddl_context.clone()); + let mut procedure = CreateViewProcedure::new(task.clone(), ddl_context.clone()); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; // Triggers procedure to create view metadata let status = procedure.execute(&ctx).await.unwrap(); @@ -215,8 +216,7 @@ async fn test_replace_view_metadata() { let mut task = test_create_view_task("foo"); // The view already exists, prepare should fail { - let mut procedure = - CreateViewProcedure::new(task.clone(), EventContext::default(), ddl_context.clone()); + let mut procedure = CreateViewProcedure::new(task.clone(), ddl_context.clone()); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::ViewAlreadyExists { .. }); assert_eq!(err.status_code(), StatusCode::TableAlreadyExists); @@ -227,12 +227,12 @@ async fn test_replace_view_metadata() { task.create_view.logical_plan = vec![4, 5, 6]; task.create_view.definition = "new_definition".to_string(); - let mut procedure = - CreateViewProcedure::new(task, EventContext::default(), ddl_context.clone()); + let mut procedure = CreateViewProcedure::new(task, ddl_context.clone()); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; // Triggers procedure to replace view metadata, but the view_id is unchanged. let status = procedure.execute(&ctx).await.unwrap(); @@ -262,17 +262,12 @@ async fn test_replace_table() { { // Create a `foo` table. let task = test_create_table_task("foo"); - let mut procedure = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut procedure = CreateTableProcedure::new(task, ddl_context.clone()).unwrap(); procedure.on_prepare().await.unwrap(); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; procedure.execute(&ctx).await.unwrap(); procedure.execute(&ctx).await.unwrap(); @@ -281,8 +276,7 @@ async fn test_replace_table() { // Try to replace a view named `foo` too. let mut task = test_create_view_task("foo"); task.create_view.or_replace = true; - let mut procedure = - CreateViewProcedure::new(task.clone(), EventContext::default(), ddl_context.clone()); + let mut procedure = CreateViewProcedure::new(task.clone(), ddl_context.clone()); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::TableAlreadyExists { .. }); assert_eq!(err.status_code(), StatusCode::TableAlreadyExists); diff --git a/src/common/meta/src/ddl/tests/drop_database.rs b/src/common/meta/src/ddl/tests/drop_database.rs index 955625e4bb3..05a70962bf0 100644 --- a/src/common/meta/src/ddl/tests/drop_database.rs +++ b/src/common/meta/src/ddl/tests/drop_database.rs @@ -26,7 +26,6 @@ use crate::ddl::drop_database::executor::DropDatabaseExecutor; use crate::ddl::test_util::datanode_handler::{NaiveDatanodeHandler, RetryErrorDatanodeHandler}; use crate::ddl::test_util::{create_logical_table, create_physical_table}; use crate::key::schema_name::SchemaNameKey; -use crate::rpc::ddl::EventContext; use crate::test_util::{MockDatanodeManager, new_ddl_context}; #[tokio::test] @@ -55,13 +54,13 @@ async fn test_drop_database_with_logical_tables() { DEFAULT_CATALOG_NAME.to_string(), DEFAULT_SCHEMA_NAME.to_string(), false, - EventContext::default(), ddl_context.clone(), ); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; while !procedure.execute(&ctx).await.unwrap().is_done() { @@ -104,13 +103,13 @@ async fn test_drop_database_retryable_error() { DEFAULT_CATALOG_NAME.to_string(), DEFAULT_SCHEMA_NAME.to_string(), false, - EventContext::default(), ddl_context.clone(), ); let ctx = ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; loop { @@ -149,7 +148,6 @@ async fn test_drop_database_recover() { DEFAULT_CATALOG_NAME.to_string(), DEFAULT_SCHEMA_NAME.to_string(), false, - EventContext::default(), ddl_context.clone(), ); let num_operating_regions = 1; diff --git a/src/common/meta/src/ddl/tests/drop_flow.rs b/src/common/meta/src/ddl/tests/drop_flow.rs index 0976aa572f7..6d36fc93ba3 100644 --- a/src/common/meta/src/ddl/tests/drop_flow.rs +++ b/src/common/meta/src/ddl/tests/drop_flow.rs @@ -26,7 +26,7 @@ use crate::ddl::test_util::flownode_handler::NaiveFlownodeHandler; use crate::ddl::tests::create_flow::{create_test_flow, create_test_pending_flow}; use crate::error; use crate::key::table_route::TableRouteValue; -use crate::rpc::ddl::{DropFlowTask, EventContext}; +use crate::rpc::ddl::DropFlowTask; use crate::test_util::{MockFlownodeManager, new_ddl_context}; pub(crate) fn test_drop_flow_task( @@ -48,7 +48,7 @@ async fn test_drop_flow_not_found() { let node_manager = Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler)); let ddl_context = new_ddl_context(node_manager); let task = test_drop_flow_task("my_flow", flow_id, false); - let mut procedure = DropFlowProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = DropFlowProcedure::new(task, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::FlowNotFound { .. }); } @@ -81,17 +81,17 @@ async fn test_drop_flow() { create_test_flow(&ddl_context, "my_flow", source_table_names, sink_table_name).await; // Drops the flows let task = test_drop_flow_task("my_flow", flow_id, false); - let mut procedure = DropFlowProcedure::new(task, EventContext::default(), ddl_context.clone()); + let mut procedure = DropFlowProcedure::new(task, ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; // Drops if not exists let task = test_drop_flow_task("my_flow", flow_id, true); - let mut procedure = DropFlowProcedure::new(task, EventContext::default(), ddl_context.clone()); + let mut procedure = DropFlowProcedure::new(task, ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; // Drops again let task = test_drop_flow_task("my_flow", flow_id, false); - let mut procedure = DropFlowProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = DropFlowProcedure::new(task, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::FlowNotFound { .. }); } @@ -129,11 +129,11 @@ async fn test_drop_pending_flow_without_routes() { assert!(flow_info.flownode_ids().is_empty()); let task = test_drop_flow_task("drop_pending_flow", flow_id, false); - let mut procedure = DropFlowProcedure::new(task, EventContext::default(), ddl_context.clone()); + let mut procedure = DropFlowProcedure::new(task, ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; let task = test_drop_flow_task("drop_pending_flow", flow_id, false); - let mut procedure = DropFlowProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = DropFlowProcedure::new(task, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, error::Error::FlowNotFound { .. }); } diff --git a/src/common/meta/src/ddl/tests/drop_table.rs b/src/common/meta/src/ddl/tests/drop_table.rs index 176f847c428..a278b951dd5 100644 --- a/src/common/meta/src/ddl/tests/drop_table.rs +++ b/src/common/meta/src/ddl/tests/drop_table.rs @@ -65,9 +65,9 @@ use crate::key::table_route::TableRouteValue; use crate::kv_backend::KvBackend; use crate::kv_backend::memory::MemoryKvBackend; use crate::peer::Peer; -use crate::rpc::ddl::{DropTableTask, EventContext}; +use crate::rpc::ddl::DropTableTask; #[cfg(feature = "enterprise")] -use crate::rpc::ddl::{PurgeDroppedTableTask, QueryContext, UndropTableTask}; +use crate::rpc::ddl::{PurgeDroppedTableTask, UndropTableTask}; use crate::rpc::router::{Region, RegionRoute}; #[cfg(feature = "enterprise")] use crate::rpc::store::{BatchDeleteRequest, PutRequest}; @@ -128,11 +128,7 @@ async fn undrop_at_restore_metadata( table_id: TableId, ) -> UndropTableProcedure { create_dropped_table(context, table_id, Some(1), Some(i64::MAX)).await; - let mut undrop = UndropTableProcedure::new( - new_undrop_table_task(table_id), - context.clone(), - EventContext::default(), - ); + let mut undrop = UndropTableProcedure::new(new_undrop_table_task(table_id), context.clone()); let procedure_context = new_test_procedure_context(); undrop.execute(&procedure_context).await.unwrap(); undrop.execute(&procedure_context).await.unwrap(); @@ -161,7 +157,6 @@ fn test_old_drop_table_json_defaults_to_hard_drop() { let procedure = DropTableProcedure::new( new_drop_table_task("foo", 1024, false), runtime_context.clone(), - EventContext::default(), ); let mut old_data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); old_data @@ -169,7 +164,6 @@ fn test_old_drop_table_json_defaults_to_hard_drop() { .unwrap() .remove("soft_drop_enabled"); old_data.as_object_mut().unwrap().remove("dropped_at"); - old_data.as_object_mut().unwrap().remove("event_context"); let recovered = DropTableProcedure::from_json(&old_data.to_string(), runtime_context).unwrap(); let recovered_data: serde_json::Value = @@ -177,20 +171,12 @@ fn test_old_drop_table_json_defaults_to_hard_drop() { assert_eq!(recovered_data["soft_drop_enabled"], false); assert_eq!(recovered_data["dropped_at"], serde_json::Value::Null); - assert_eq!( - recovered_data["event_context"], - serde_json::json!({"reason": "unknown"}) - ); } #[test] fn test_drop_table_lock_key_includes_table_name() { let context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); - let procedure = DropTableProcedure::new( - new_drop_table_task("foo", 1024, false), - context, - EventContext::default(), - ); + let procedure = DropTableProcedure::new(new_drop_table_task("foo", 1024, false), context); let keys = procedure .lock_key() @@ -219,11 +205,7 @@ fn test_new_drop_table_procedure_disables_soft_drop_without_enterprise() { let mut context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); context.soft_drop_enabled = true; - let procedure = DropTableProcedure::new( - new_drop_table_task("foo", 1024, false), - context, - EventContext::default(), - ); + let procedure = DropTableProcedure::new(new_drop_table_task("foo", 1024, false), context); assert!(!procedure.data.soft_drop_enabled); } @@ -234,11 +216,7 @@ fn test_new_drop_table_procedure_preserves_soft_drop_with_enterprise() { let mut context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); context.soft_drop_enabled = true; - let procedure = DropTableProcedure::new( - new_drop_table_task("foo", 1024, false), - context, - EventContext::default(), - ); + let procedure = DropTableProcedure::new(new_drop_table_task("foo", 1024, false), context); assert!(procedure.data.soft_drop_enabled); } @@ -273,7 +251,6 @@ async fn test_recovered_soft_drop_preserves_persisted_mode() { let procedure = DropTableProcedure::new( new_drop_table_task("foo", table_id, false), soft_context.clone(), - EventContext::default(), ); let hard_context = new_ddl_context_with_kv_backend(node_manager, kv_backend); @@ -313,13 +290,9 @@ async fn test_legacy_soft_drop_prepare_preserves_persisted_mode() { .unwrap(); let mut persisted: serde_json::Value = serde_json::from_str( - &DropTableProcedure::new( - new_drop_table_task("foo", table_id, false), - context.clone(), - EventContext::default(), - ) - .dump() - .unwrap(), + &DropTableProcedure::new(new_drop_table_task("foo", table_id, false), context.clone()) + .dump() + .unwrap(), ) .unwrap(); persisted["soft_drop_enabled"] = true.into(); @@ -349,11 +322,7 @@ async fn test_disabled_soft_drop_creates_hard_drop_procedure() { context.soft_drop_enabled = false; context.soft_drop_retention = Some(Duration::from_millis(100)); - let procedure = DropTableProcedure::new( - new_drop_table_task("foo", 1024, false), - context, - EventContext::default(), - ); + let procedure = DropTableProcedure::new(new_drop_table_task("foo", 1024, false), context); assert!(!procedure.data.soft_drop_enabled); assert_eq!(None, procedure.data.dropped_at); @@ -387,7 +356,6 @@ async fn test_recovered_hard_drop_ignores_soft_runtime_context() { let procedure = DropTableProcedure::new( new_drop_table_task("foo", table_id, false), hard_context.clone(), - EventContext::default(), ); let mut soft_context = new_ddl_context_with_kv_backend(node_manager, kv_backend); @@ -429,7 +397,7 @@ async fn test_on_prepare_table_not_exists_err() { .unwrap(); let task = new_drop_table_task("bar", table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context, EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context); assert_eq!(procedure.data.dropped_at, None); assert_eq!(procedure.data.retention_expires_at, None); let err = procedure.on_prepare().await.unwrap_err(); @@ -459,7 +427,6 @@ async fn test_soft_drop_prepare_assigns_stable_deadline_and_recovers_it() { let mut procedure = DropTableProcedure::new( new_drop_table_task("foo", table_id, false), ddl_context.clone(), - EventContext::default(), ); procedure.on_prepare().await.unwrap(); @@ -496,7 +463,6 @@ async fn test_soft_drop_prepare_rejects_deadline_overflow_before_metadata_delete let mut procedure = DropTableProcedure::new( new_drop_table_task("foo", table_id, false), ddl_context.clone(), - EventContext::default(), ); let error = procedure.on_prepare().await.unwrap_err(); @@ -536,13 +502,13 @@ async fn test_on_prepare_table() { let task = new_drop_table_task("bar", table_id, true); // Drop if exists - let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); procedure.on_prepare().await.unwrap(); assert!(!procedure.rollback_supported()); let task = new_drop_table_task(table_name, table_id, false); // Drop table - let mut procedure = DropTableProcedure::new(task, ddl_context, EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context); procedure.on_prepare().await.unwrap(); } @@ -593,7 +559,7 @@ async fn test_on_datanode_drop_regions() { let task = new_drop_table_task(table_name, table_id, false); // Drop table - let mut procedure = DropTableProcedure::new(task, ddl_context, EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context); procedure.on_prepare().await.unwrap(); procedure.on_datanode_drop_regions(false).await.unwrap(); @@ -662,7 +628,7 @@ async fn test_on_datanode_drop_regions_remaps_addresses_when_retrying() { .unwrap(); let task = new_drop_table_task(table_name, table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); procedure.on_prepare().await.unwrap(); put_datanode_address(&ddl_context, 1, "new-leader").await; @@ -721,7 +687,7 @@ async fn test_soft_drop_closes_regions_and_keeps_tombstone() { .unwrap(); let task = new_drop_table_task(table_name, table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; @@ -835,7 +801,7 @@ async fn test_soft_drop_keeps_metadata_live_until_regions_close() { .unwrap(); let task = new_drop_table_task(table_name, table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); let ctx = new_test_procedure_context(); procedure.execute(&ctx).await.unwrap(); put_datanode_address(&ddl_context, 1, "new-leader").await; @@ -922,7 +888,6 @@ async fn test_soft_drop_timestamp_is_stable_across_retry_and_recovery() { let procedure = DropTableProcedure::new( new_drop_table_task("foo", table_id, false), ddl_context.clone(), - EventContext::default(), ); assert_eq!(procedure.data.dropped_at, None); let mut procedure = procedure; @@ -993,7 +958,7 @@ async fn test_hard_drop_keeps_delete_tombstone_flow() { .unwrap(); let task = new_drop_table_task(table_name, table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); execute_procedure_until(&mut procedure, |p| { p.data.state == DropTableState::DeleteTombstone @@ -1036,20 +1001,13 @@ async fn test_create_table_succeeds_while_tombstone_exists() { .unwrap(); let drop_task = new_drop_table_task(table_name, dropped_table_id, false); - let mut drop_procedure = - DropTableProcedure::new(drop_task, ddl_context.clone(), EventContext::default()); + let mut drop_procedure = DropTableProcedure::new(drop_task, ddl_context.clone()); execute_procedure_until_done(&mut drop_procedure).await; let mut create_task = test_create_table_task(table_name, 1025); create_task.create_table.table_id = None; create_task.table_info.ident.table_id = 0; - let mut create_procedure = CreateTableProcedure::new( - create_task, - QueryContext::default(), - EventContext::default(), - ddl_context.clone(), - ) - .unwrap(); + let mut create_procedure = CreateTableProcedure::new(create_task, ddl_context.clone()).unwrap(); execute_procedure_until_done(&mut create_procedure).await; let live_table = ddl_context @@ -1094,8 +1052,7 @@ async fn test_hard_drop_recreated_table_fails_when_soft_tombstone_exists() { .unwrap(); let drop_task = new_drop_table_task(table_name, original_table_id, false); - let mut drop_procedure = - DropTableProcedure::new(drop_task, ddl_context.clone(), EventContext::default()); + let mut drop_procedure = DropTableProcedure::new(drop_task, ddl_context.clone()); execute_procedure_until_done(&mut drop_procedure).await; ddl_context.soft_drop_enabled = false; @@ -1112,7 +1069,6 @@ async fn test_hard_drop_recreated_table_fails_when_soft_tombstone_exists() { let mut procedure = DropTableProcedure::new( new_drop_table_task(table_name, recreated_table_id, false), ddl_context, - EventContext::default(), ); let err = procedure.on_prepare().await.unwrap_err(); @@ -1142,8 +1098,7 @@ async fn test_hard_drop_recreated_table_ignores_previous_orphan_tombstone() { .unwrap(); let drop_task = new_drop_table_task(table_name, original_table_id, false); - let mut drop_procedure = - DropTableProcedure::new(drop_task, ddl_context.clone(), EventContext::default()); + let mut drop_procedure = DropTableProcedure::new(drop_task, ddl_context.clone()); execute_procedure_until(&mut drop_procedure, |p| { p.data.state == DropTableState::DeleteTombstone }) @@ -1162,7 +1117,6 @@ async fn test_hard_drop_recreated_table_ignores_previous_orphan_tombstone() { let mut procedure = DropTableProcedure::new( new_drop_table_task(table_name, recreated_table_id, false), ddl_context, - EventContext::default(), ); procedure.on_prepare().await.unwrap(); @@ -1202,7 +1156,6 @@ async fn test_undrop_table_restores_metadata_and_reopens_regions() { let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; assert!( @@ -1219,11 +1172,8 @@ async fn test_undrop_table_restores_metadata_and_reopens_regions() { while rx.try_recv().is_ok() {} - let mut procedure = UndropTableProcedure::new( - new_undrop_table_task(table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; let live_table = ddl_context @@ -1314,16 +1264,12 @@ async fn test_undrop_table_opens_regions_before_restoring_live_metadata() { let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; while rx.try_recv().is_ok() {} - let mut procedure = UndropTableProcedure::new( - new_undrop_table_task(table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); let ctx = new_test_procedure_context(); procedure.execute(&ctx).await.unwrap(); procedure.execute(&ctx).await.unwrap(); @@ -1387,17 +1333,13 @@ async fn test_undrop_logical_table_skips_datanode_open() { let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, logical_table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; while rx.try_recv().is_ok() {} - let mut procedure = UndropTableProcedure::new( - new_undrop_table_task(logical_table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(logical_table_id), ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; let live_table = ddl_context @@ -1428,7 +1370,6 @@ async fn test_soft_drop_metric_logical_table_falls_back_to_hard_drop() { let mut procedure = DropTableProcedure::new( new_drop_table_task("foo", logical_table_id, false), ddl_context.clone(), - EventContext::default(), ); procedure.on_prepare().await.unwrap(); assert!(!procedure.data.soft_drop_enabled); @@ -1472,7 +1413,6 @@ async fn test_soft_drop_metric_physical_table_remains_enabled() { let mut procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut procedure).await; @@ -1517,7 +1457,6 @@ async fn test_soft_drop_file_engine_table_falls_back_to_hard_drop() { let mut procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), - EventContext::default(), ); procedure.on_prepare().await.unwrap(); assert!(!procedure.data.soft_drop_enabled); @@ -1554,7 +1493,6 @@ async fn test_file_engine_fallback_is_resolved_before_tombstone_conflict() { let mut procedure = DropTableProcedure::new( new_drop_table_task(table_name, recreated_table_id, false), ddl_context, - EventContext::default(), ); procedure.on_prepare().await.unwrap(); @@ -1570,11 +1508,8 @@ async fn test_undrop_metric_logical_table_fails() { let logical_table_id = create_metric_logical_table_tombstone(&ddl_context, physical_table_id, "foo").await; - let mut procedure = UndropTableProcedure::new( - new_undrop_table_task(logical_table_id), - ddl_context, - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(logical_table_id), ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_eq!(err.status_code(), StatusCode::Unsupported); @@ -1592,7 +1527,6 @@ async fn test_purge_metric_logical_table_fails() { let mut procedure = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(logical_table_id), ddl_context, - EventContext::default(), ); let err = procedure .execute(&new_test_procedure_context()) @@ -1624,7 +1558,6 @@ async fn test_undrop_table_fails_when_live_name_exists() { let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, dropped_table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; ddl_context @@ -1637,11 +1570,8 @@ async fn test_undrop_table_fails_when_live_name_exists() { .await .unwrap(); - let mut procedure = UndropTableProcedure::new( - new_undrop_table_task(dropped_table_id), - ddl_context, - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(dropped_table_id), ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, Error::TableAlreadyExists { .. }); @@ -1669,15 +1599,11 @@ async fn test_undrop_table_fails_when_live_name_is_created_after_prepare() { let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, dropped_table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; - let mut procedure = UndropTableProcedure::new( - new_undrop_table_task(dropped_table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(dropped_table_id), ddl_context.clone()); procedure.on_prepare().await.unwrap(); ddl_context .table_metadata_manager @@ -1741,17 +1667,13 @@ async fn test_undrop_table_closes_opened_regions_when_restore_metadata_races_wit let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, dropped_table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; while rx.try_recv().is_ok() {} detector_controller.clear().await; - let mut procedure = UndropTableProcedure::new( - new_undrop_table_task(dropped_table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(dropped_table_id), ddl_context.clone()); let ctx = new_test_procedure_context(); procedure.execute(&ctx).await.unwrap(); procedure.execute(&ctx).await.unwrap(); @@ -1845,7 +1767,6 @@ async fn test_undrop_table_lock_key_includes_original_table_name_before_prepare( let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; @@ -1860,7 +1781,6 @@ async fn test_undrop_table_lock_key_includes_original_table_name_before_prepare( new_undrop_table_task(table_id), ddl_context, Some(original_table_name), - EventContext::default(), ); let keys = procedure @@ -1928,12 +1848,10 @@ async fn test_undrop_rejects_stale_id_after_name_tombstone_is_consumed() { new_undrop_table_task(original_table_id), ddl_context.clone(), Some(original_table_name.clone()), - EventContext::default(), ); let mut hard_drop = DropTableProcedure::new( new_drop_table_task(table_name, recreated_table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut hard_drop).await; assert!( @@ -1988,15 +1906,11 @@ async fn test_undrop_table_replayed_restore_metadata_is_idempotent() { let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; - let mut procedure = UndropTableProcedure::new( - new_undrop_table_task(table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); let ctx = new_test_procedure_context(); procedure.execute(&ctx).await.unwrap(); procedure.execute(&ctx).await.unwrap(); @@ -2100,7 +2014,6 @@ async fn test_purge_dropped_table_cleans_regions_offline_and_deletes_tombstone() let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; while rx.try_recv().is_ok() {} @@ -2109,7 +2022,6 @@ async fn test_purge_dropped_table_cleans_regions_offline_and_deletes_tombstone() let mut procedure = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(table_id), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut procedure).await; @@ -2188,7 +2100,6 @@ async fn test_automatic_purge_rechecks_unexpired_tombstone() { let procedure = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(table_id), ddl_context.clone(), - EventContext::default(), ); let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); data["check_expired"] = true.into(); @@ -2245,7 +2156,6 @@ async fn test_recovered_automatic_purge_rechecks_unexpired_tombstone() { let procedure = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(table_id), ddl_context.clone(), - EventContext::default(), ); let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); data["check_expired"] = true.into(); @@ -2323,7 +2233,6 @@ async fn test_recovered_automatic_purge_obeys_runtime_disable() { let procedure = PurgeDroppedTableProcedure::new_if_expired( new_purge_dropped_table_task(table_id), ddl_context.clone(), - EventContext::default(), ); let mut recovered = PurgeDroppedTableProcedure::from_json(&procedure.dump().unwrap(), ddl_context.clone()) @@ -2355,7 +2264,6 @@ async fn test_recovered_automatic_purge_finishes_after_cleanup_is_claimed() { let mut procedure = PurgeDroppedTableProcedure::new_if_expired( new_purge_dropped_table_task(table_id), enabled_context.clone(), - EventContext::default(), ); procedure .execute(&new_test_procedure_context()) @@ -2419,7 +2327,6 @@ async fn test_recovered_automatic_purge_ignores_foreign_generation_claim() { let mut procedure = PurgeDroppedTableProcedure::new_if_expired( new_purge_dropped_table_task(table_id), ddl_context.clone(), - EventContext::default(), ); procedure .execute(&new_test_procedure_context()) @@ -2462,11 +2369,8 @@ async fn test_undrop_rejects_table_claimed_for_purge() { .await .unwrap(); - let mut procedure = UndropTableProcedure::new( - new_undrop_table_task(table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); let error = procedure.on_prepare().await.unwrap_err(); assert_eq!(StatusCode::TableNotFound, error.status_code()); @@ -2488,11 +2392,8 @@ async fn test_recovered_undrop_rejects_tombstone_deleted_by_purge() { let table_id = 1024; create_dropped_table(&ddl_context, table_id, Some(1), Some(i64::MAX)).await; - let mut undrop = UndropTableProcedure::new( - new_undrop_table_task(table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut undrop = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); undrop.on_prepare().await.unwrap(); let open_regions_state = undrop.dump().unwrap(); undrop.execute(&new_test_procedure_context()).await.unwrap(); @@ -2501,7 +2402,6 @@ async fn test_recovered_undrop_rejects_tombstone_deleted_by_purge() { let mut purge = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(table_id), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut purge).await; @@ -2533,11 +2433,8 @@ async fn test_recovered_undrop_rejects_replacement_generation() { .await .unwrap(); - let mut undrop = UndropTableProcedure::new( - new_undrop_table_task(table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut undrop = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); undrop.on_prepare().await.unwrap(); let persisted = undrop.dump().unwrap(); kv_backend @@ -2603,11 +2500,8 @@ async fn test_legacy_unmarked_tombstone_can_be_undropped_and_purged() { assert_eq!(dropped_tables.len(), 1); assert_eq!(dropped_tables[0].dropped_at, None); - let mut undrop = UndropTableProcedure::new( - new_undrop_table_task(table_id), - ddl_context.clone(), - EventContext::default(), - ); + let mut undrop = + UndropTableProcedure::new(new_undrop_table_task(table_id), ddl_context.clone()); execute_procedure_until_done(&mut undrop).await; ddl_context @@ -2624,7 +2518,6 @@ async fn test_legacy_unmarked_tombstone_can_be_undropped_and_purged() { let mut purge = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(table_id), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut purge).await; assert!( @@ -2667,7 +2560,6 @@ async fn test_purge_dropped_table_by_id_selects_tombstone_when_live_table_exists let mut drop_procedure = DropTableProcedure::new( new_drop_table_task(table_name, dropped_table_id, false), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut drop_procedure).await; ddl_context @@ -2684,7 +2576,6 @@ async fn test_purge_dropped_table_by_id_selects_tombstone_when_live_table_exists let mut procedure = PurgeDroppedTableProcedure::new( new_purge_dropped_table_task(dropped_table_id), ddl_context.clone(), - EventContext::default(), ); execute_procedure_until_done(&mut procedure).await; @@ -2735,12 +2626,8 @@ async fn test_on_rollback() { let physical_table_id = table_id; // Creates the logical table metadata. let task = test_create_logical_table_task("foo"); - let mut procedure = CreateLogicalTablesProcedure::new( - vec![task], - physical_table_id, - EventContext::default(), - ddl_context.clone(), - ); + let mut procedure = + CreateLogicalTablesProcedure::new(vec![task], physical_table_id, ddl_context.clone()); procedure.on_prepare().await.unwrap(); let ctx = new_test_procedure_context(); procedure.execute(&ctx).await.unwrap(); @@ -2753,8 +2640,7 @@ async fn test_on_rollback() { // Drops the physical table { let task = new_drop_table_task("phy_table", physical_table_id, false); - let mut procedure = - DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); procedure.on_prepare().await.unwrap(); assert!(procedure.rollback_supported()); procedure.on_delete_metadata().await.unwrap(); @@ -2765,8 +2651,7 @@ async fn test_on_rollback() { // The physical table can be dropped again after the metadata rollback. let retry_task = new_drop_table_task("phy_table", physical_table_id, false); - let mut retry = - DropTableProcedure::new(retry_task, ddl_context.clone(), EventContext::default()); + let mut retry = DropTableProcedure::new(retry_task, ddl_context.clone()); retry.on_prepare().await.unwrap(); retry.on_delete_metadata().await.unwrap(); retry.rollback(&ctx).await.unwrap(); @@ -2783,7 +2668,7 @@ async fn test_on_rollback() { // Drops the logical table ddl_context.soft_drop_enabled = false; let task = new_drop_table_task("foo", table_ids[0], false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); procedure.on_prepare().await.unwrap(); assert!(!procedure.rollback_supported()); } @@ -2880,8 +2765,7 @@ async fn test_memory_region_keeper_guard_dropped_on_procedure_done() { let logical_table_id = create_logical_table(ddl_context.clone(), physical_table_id, "s").await; let inner_test = |task: DropTableTask| async { - let mut procedure = - DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); execute_procedure_until(&mut procedure, |p| { p.data.state == DropTableState::InvalidateTableCache }) @@ -2933,8 +2817,7 @@ async fn test_from_json() { let physical_table_id = create_physical_table(&ddl_context, "t").await; let task = new_drop_table_task("t", physical_table_id, false); - let mut procedure = - DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); execute_procedure_until(&mut procedure, |p| p.data.state == state).await; let data = procedure.dump().unwrap(); assert_eq!( @@ -2963,7 +2846,7 @@ async fn test_from_json() { let physical_table_id = create_physical_table(&ddl_context, "t").await; let task = new_drop_table_task("t", physical_table_id, false); - let mut procedure = DropTableProcedure::new(task, ddl_context.clone(), EventContext::default()); + let mut procedure = DropTableProcedure::new(task, ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; let data = procedure.dump().unwrap(); assert_eq!( diff --git a/src/common/meta/src/ddl/tests/drop_view.rs b/src/common/meta/src/ddl/tests/drop_view.rs index 09b2a450c52..412c5a78f0b 100644 --- a/src/common/meta/src/ddl/tests/drop_view.rs +++ b/src/common/meta/src/ddl/tests/drop_view.rs @@ -24,7 +24,7 @@ use crate::ddl::drop_view::{DropViewProcedure, DropViewState}; use crate::ddl::test_util::create_table::test_create_table_task; use crate::ddl::tests::create_view::{test_create_view_task, test_table_names}; use crate::key::table_route::TableRouteValue; -use crate::rpc::ddl::{DropViewTask, EventContext}; +use crate::rpc::ddl::DropViewTask; use crate::test_util::{MockDatanodeManager, new_ddl_context}; pub(crate) fn new_drop_view_task( @@ -63,7 +63,7 @@ async fn test_on_prepare_view_not_exists_err() { .unwrap(); let task = new_drop_view_task("bar", view_id, false); - let mut procedure = DropViewProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = DropViewProcedure::new(task, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_eq!(err.status_code(), StatusCode::TableNotFound); } @@ -87,7 +87,7 @@ async fn test_on_prepare_not_view_err() { .unwrap(); let task = new_drop_view_task(view_name, view_id, false); - let mut procedure = DropViewProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = DropViewProcedure::new(task, ddl_context); // It's not a view, expect error let err = procedure.on_prepare().await.unwrap_err(); assert_eq!(err.status_code(), StatusCode::InvalidArguments); @@ -117,12 +117,12 @@ async fn test_on_prepare_success() { let task = new_drop_view_task("bar", view_id, true); // Drop if exists - let mut procedure = DropViewProcedure::new(task, EventContext::default(), ddl_context.clone()); + let mut procedure = DropViewProcedure::new(task, ddl_context.clone()); procedure.on_prepare().await.unwrap(); let task = new_drop_view_task(view_name, view_id, false); // Prepare success - let mut procedure = DropViewProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = DropViewProcedure::new(task, ddl_context); procedure.on_prepare().await.unwrap(); assert_eq!(DropViewState::DeleteMetadata, procedure.state()); } @@ -161,7 +161,7 @@ async fn test_drop_view_success() { let task = new_drop_view_task(view_name, view_id, false); // Prepare success - let mut procedure = DropViewProcedure::new(task, EventContext::default(), ddl_context.clone()); + let mut procedure = DropViewProcedure::new(task, ddl_context.clone()); execute_procedure_until_done(&mut procedure).await; assert_eq!(DropViewState::InvalidateViewCache, procedure.state()); @@ -178,7 +178,7 @@ async fn test_drop_view_success() { // Drop again let task = new_drop_view_task(view_name, view_id, false); - let mut procedure = DropViewProcedure::new(task, EventContext::default(), ddl_context); + let mut procedure = DropViewProcedure::new(task, ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_eq!(err.status_code(), StatusCode::TableNotFound); } diff --git a/src/common/meta/src/ddl/tests/event/database.rs b/src/common/meta/src/ddl/tests/event/database.rs index 29cb48f7b80..8649982ed15 100644 --- a/src/common/meta/src/ddl/tests/event/database.rs +++ b/src/common/meta/src/ddl/tests/event/database.rs @@ -25,12 +25,12 @@ use common_event_recorder::event_table::{ PROCEDURE_ID_COLUMN as EVENT_TABLE_PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN as EVENT_TABLE_PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN as EVENT_TABLE_PROCEDURE_TRIGGER_COLUMN, - SCHEMA_NAME_COLUMN as EVENT_TABLE_SCHEMA_NAME_COLUMN, + SCHEMA_NAME_COLUMN as EVENT_TABLE_SCHEMA_NAME_COLUMN, jsonb_value, }; use common_event_recorder::testing::assert_event_contract; use common_procedure::{EventTrigger, ProcedureEvent, ProcedureId, ProcedureState}; -use super::test_util::{assert_event_filter, default_event_context_value, procedure_trigger_value}; +use super::test_util::assert_event_filter; use crate::ddl::alter_database::AlterDatabaseProcedure; use crate::ddl::create_database::CreateDatabaseProcedure; use crate::ddl::drop_database::DropDatabaseProcedure; @@ -39,7 +39,7 @@ use crate::ddl::event::database::{ DatabaseDdlEvent, }; use crate::rpc::ddl::{ - AlterDatabaseKind, EventContext, SetDatabaseOption, SetDatabaseOptions, UnsetDatabaseOption, + AlterDatabaseKind, SetDatabaseOption, SetDatabaseOptions, UnsetDatabaseOption, UnsetDatabaseOptions, }; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -50,13 +50,7 @@ fn test_create_database_submitted_event_contract() { ("password".to_string(), "do-not-record".to_string()), ("compaction.type".to_string(), "twcs".to_string()), ]); - let event = DatabaseDdlEvent::create_submitted( - "greptime", - "metrics", - true, - &options, - EventContext::default(), - ); + let event = DatabaseDdlEvent::create_submitted("greptime", "metrics", true, &options); assert_event_locator( &event, @@ -93,7 +87,6 @@ fn test_alter_database_set_and_unset_event_contracts() { SetDatabaseOption::Other("secret_token".to_string(), "hidden".to_string()), SetDatabaseOption::Ttl(std::time::Duration::from_secs(3600).into()), ])), - EventContext::default(), ); assert_event_locator( &set, @@ -120,7 +113,6 @@ fn test_alter_database_set_and_unset_event_contracts() { UnsetDatabaseOption::Ttl, UnsetDatabaseOption::Other("compaction.type".to_string()), ])), - EventContext::default(), ); assert_event_locator( &unset, @@ -140,8 +132,7 @@ fn test_alter_database_set_and_unset_event_contracts() { #[test] fn test_drop_database_submitted_event_contract() { - let event = - DatabaseDdlEvent::drop_submitted("greptime", "metrics", true, EventContext::default()); + let event = DatabaseDdlEvent::drop_submitted("greptime", "metrics", true); assert_event_locator( &event, @@ -177,14 +168,7 @@ fn test_database_lifecycle_events_preserve_locator_and_null_intent() { assert_eq!( DatabaseDdlEvent::create_lifecycle("greptime", "metrics").extra_schema(), - DatabaseDdlEvent::create_submitted( - "c", - "s", - false, - &HashMap::new(), - EventContext::default(), - ) - .extra_schema() + DatabaseDdlEvent::create_submitted("c", "s", false, &HashMap::new()).extra_schema() ); } @@ -198,7 +182,6 @@ fn test_database_events_preserve_procedure_envelope_contract() { "metrics", false, &HashMap::new(), - EventContext::default(), )), ProcedureState::Running, EventTrigger::Submitted, @@ -237,7 +220,6 @@ fn test_create_database_event_filter() { false, HashMap::new(), None, - EventContext::default(), new_ddl_context(Arc::new(MockDatanodeManager::new(()))), ); @@ -258,7 +240,6 @@ fn test_alter_database_event_filter() { )), }, }, - EventContext::default(), new_ddl_context(Arc::new(MockDatanodeManager::new(()))), ) .unwrap(); @@ -272,7 +253,6 @@ fn test_drop_database_event_filter() { "greptime".to_string(), "metrics".to_string(), false, - EventContext::default(), new_ddl_context(Arc::new(MockDatanodeManager::new(()))), ); @@ -285,14 +265,12 @@ fn assert_event_locator( catalog_name: Option<&str>, schema_name: Option<&str>, ) { - let has_event_context = !event.json_payload().unwrap().is_null(); assert_event_contract( event, event_type, &[ EVENT_TABLE_CATALOG_NAME_COLUMN.column_schema(), EVENT_TABLE_SCHEMA_NAME_COLUMN.column_schema(), - EVENT_CONTEXT_COLUMN.column_schema(), ], &[Row { values: vec![ @@ -302,11 +280,6 @@ fn assert_event_locator( Value { value_data: schema_name.map(|value| ValueData::StringValue(value.to_string())), }, - if has_event_context { - default_event_context_value() - } else { - Value { value_data: None } - }, ], }], ); @@ -320,7 +293,6 @@ fn assert_procedure_event_contract( catalog_name: Option<&str>, schema_name: Option<&str>, ) { - let has_event_context = procedure_trigger == "Submitted"; assert_event_contract( event, event_type, @@ -346,18 +318,14 @@ fn assert_procedure_event_contract( Value { value_data: Some(ValueData::StringValue(String::new())), }, - procedure_trigger_value(procedure_trigger), + jsonb_value(&serde_json::json!({"type": procedure_trigger})), Value { value_data: catalog_name.map(|value| ValueData::StringValue(value.to_string())), }, Value { value_data: schema_name.map(|value| ValueData::StringValue(value.to_string())), }, - if has_event_context { - default_event_context_value() - } else { - Value { value_data: None } - }, + Value { value_data: None }, ], }], ); diff --git a/src/common/meta/src/ddl/tests/event/flow.rs b/src/common/meta/src/ddl/tests/event/flow.rs index 39f24b74e51..aae3613a087 100644 --- a/src/common/meta/src/ddl/tests/event/flow.rs +++ b/src/common/meta/src/ddl/tests/event/flow.rs @@ -21,12 +21,13 @@ use common_event_recorder::Event; use common_event_recorder::event_table::{ CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, + jsonb_value, }; use common_event_recorder::testing::assert_event_contract; use common_procedure::{EventTrigger, ProcedureEvent, ProcedureId, ProcedureState}; use table::table_name::TableName; -use super::test_util::{assert_event_filter, default_event_context_value, procedure_trigger_value}; +use super::test_util::assert_event_filter; use crate::ddl::create_flow::CreateFlowProcedure; use crate::ddl::drop_flow::DropFlowProcedure; use crate::ddl::event::flow::{ @@ -35,7 +36,6 @@ use crate::ddl::event::flow::{ use crate::ddl::test_util::flownode_handler::NaiveFlownodeHandler; use crate::ddl::tests::create_flow::{test_create_flow_task, test_query_context}; use crate::ddl::tests::drop_flow::test_drop_flow_task; -use crate::rpc::ddl::EventContext; use crate::test_util::{MockFlownodeManager, new_ddl_context}; #[test] @@ -49,7 +49,6 @@ fn test_flow_submitted_event_contracts() { expire_after: Some(300), eval_interval_secs: Some(60), }, - EventContext::default(), ); assert_event_contract( &create, @@ -60,7 +59,6 @@ fn test_flow_submitted_event_contracts() { ValueData::StringValue("greptime".to_string()).into(), ValueData::StringValue("metrics".to_string()).into(), Value { value_data: None }, - default_event_context_value(), ], }], ); @@ -75,8 +73,7 @@ fn test_flow_submitted_event_contracts() { }) ); - let drop = - FlowDdlEvent::drop_submitted("greptime", "metrics", 42, true, EventContext::default()); + let drop = FlowDdlEvent::drop_submitted("greptime", "metrics", 42, true); assert_event_contract( &drop, DROP_FLOW_EVENT_TYPE, @@ -86,7 +83,6 @@ fn test_flow_submitted_event_contracts() { ValueData::StringValue("greptime".to_string()).into(), ValueData::StringValue("metrics".to_string()).into(), ValueData::U32Value(42).into(), - default_event_context_value(), ], }], ); @@ -125,7 +121,6 @@ fn test_flow_lifecycle_events_have_fixed_schema_and_null_intent() { } else { Value { value_data: None } }, - Value { value_data: None }, ], }], ); @@ -142,7 +137,6 @@ fn test_flow_lifecycle_events_have_fixed_schema_and_null_intent() { ValueData::StringValue("greptime".to_string()).into(), ValueData::StringValue("metrics".to_string()).into(), ValueData::U32Value(42).into(), - Value { value_data: None }, ], }], ); @@ -162,7 +156,6 @@ fn test_flow_events_preserve_procedure_envelope_contract() { expire_after: None, eval_interval_secs: None, }, - EventContext::default(), )), ProcedureState::Running, EventTrigger::Submitted, @@ -212,7 +205,6 @@ fn test_create_flow_event_filter() { false, ), test_query_context(), - EventContext::default(), new_ddl_context(Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler))), ); assert_event_filter(&procedure, CREATE_FLOW_EVENT_TYPE); @@ -222,7 +214,6 @@ fn test_create_flow_event_filter() { fn test_drop_flow_event_filter() { let procedure = DropFlowProcedure::new( test_drop_flow_task("flow", 42, false), - EventContext::default(), new_ddl_context(Arc::new(MockFlownodeManager::new(NaiveFlownodeHandler))), ); assert_event_filter(&procedure, DROP_FLOW_EVENT_TYPE); @@ -233,7 +224,6 @@ fn flow_schema() -> Vec { CATALOG_NAME_COLUMN.column_schema(), FLOW_NAME_COLUMN.column_schema(), FLOW_ID_COLUMN.column_schema(), - EVENT_CONTEXT_COLUMN.column_schema(), ] } @@ -257,6 +247,7 @@ fn assert_procedure_event_contract( PROCEDURE_TRIGGER_COLUMN.column_schema(), ]; schema.extend(flow_schema()); + schema.push(EVENT_CONTEXT_COLUMN.column_schema()); assert_event_contract( event, event_type, @@ -266,7 +257,7 @@ fn assert_procedure_event_contract( ValueData::StringValue(event.procedure_id.to_string()).into(), ValueData::StringValue(state.to_string()).into(), ValueData::StringValue(String::new()).into(), - procedure_trigger_value(trigger), + jsonb_value(&serde_json::json!({"type": trigger})), optional_string(locator.catalog_name), optional_string(locator.flow_name), locator @@ -274,11 +265,7 @@ fn assert_procedure_event_contract( .map(ValueData::U32Value) .map(Into::into) .unwrap_or(Value { value_data: None }), - if trigger == "Submitted" { - default_event_context_value() - } else { - Value { value_data: None } - }, + Value { value_data: None }, ], }], ); diff --git a/src/common/meta/src/ddl/tests/event/table.rs b/src/common/meta/src/ddl/tests/event/table.rs index 3326cb06808..24182fd19af 100644 --- a/src/common/meta/src/ddl/tests/event/table.rs +++ b/src/common/meta/src/ddl/tests/event/table.rs @@ -19,13 +19,13 @@ use api::v1::value::ValueData; use api::v1::{ColumnDataType, Repartition, SemanticType, Value}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, PHYSICAL_TABLE_ID_COLUMN, SCHEMA_NAME_COLUMN, - TABLE_ID_COLUMN, TABLE_NAME_COLUMN, jsonb_value, + CATALOG_NAME_COLUMN, PHYSICAL_TABLE_ID_COLUMN, SCHEMA_NAME_COLUMN, TABLE_ID_COLUMN, + TABLE_NAME_COLUMN, }; use common_event_recorder::{Event, EventTypeFilter}; use common_procedure::{ - ChildSubmissionOutcome, EventRuntimeContext, EventTrigger, Procedure, ProcedureId, - ProcedureState, RetryPhase, + ChildSubmissionOutcome, EventContext, EventTrigger, Procedure, ProcedureId, ProcedureState, + RetryPhase, }; use common_time::Timestamp; use serde_json::{Value as JsonValue, json}; @@ -52,9 +52,9 @@ use crate::ddl::truncate_table::TruncateTableProcedure; use crate::ddl::undrop_table::UndropTableProcedure; use crate::key::DeserializedValueWithBytes; use crate::key::table_info::TableInfoValue; -use crate::rpc::ddl::{DropTableTask, EventContext, QueryContext, TruncateTableTask}; +use crate::rpc::ddl::{DropTableTask, TruncateTableTask}; #[cfg(feature = "enterprise")] -use crate::rpc::ddl::{PurgeDroppedTableTask, TriggerReason, UndropTableTask}; +use crate::rpc::ddl::{PurgeDroppedTableTask, UndropTableTask}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; struct EventCase { @@ -175,11 +175,7 @@ fn later_lifecycle_events_are_uniform() { .extra_rows() .unwrap() .into_iter() - .map(|mut row| { - // Lifecycle events retain locators but not submitted event context. - *row.values.last_mut().unwrap() = Value::default(); - row.values - }) + .map(|row| row.values) .collect::>(); for trigger in &triggers { @@ -205,13 +201,7 @@ fn later_lifecycle_events_are_uniform() { fn create_success_events_keep_allocated_ids() { let mut task = test_create_table_task("create_success"); task.table_info.ident.table_id = 7; - let create_table = CreateTableProcedure::new( - task, - QueryContext::default(), - EventContext::default(), - test_context(), - ) - .unwrap(); + let create_table = CreateTableProcedure::new(task, test_context()).unwrap(); let state = ProcedureState::Done { output: Some(Arc::new(42_u32)), }; @@ -230,7 +220,6 @@ fn create_success_events_keep_allocated_ids() { test_create_logical_table_task("bar"), ], 1024, - EventContext::default(), test_context(), ); let state = ProcedureState::Done { @@ -269,14 +258,13 @@ fn event_cases() -> Vec { TableDdlLocator::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "create"), true, "mito2", - EventContext::default(), ), payload: json!({ "version": TABLE_DDL_PAYLOAD_VERSION, "create_if_not_exists": true, "engine": "mito2", }), - rows: vec![submitted_table_locator_values(Some("create"), None)], + rows: vec![table_locator_values(Some("create"), None)], }, EventCase { event_type: TableDdlEventType::CreateLogicalTables, @@ -288,15 +276,14 @@ fn event_cases() -> Vec { .with_physical_table_id(10), ], 2, - EventContext::default(), ), payload: json!({ "version": TABLE_DDL_PAYLOAD_VERSION, "table_count": 2, }), rows: vec![ - submitted_logical_locator_values("logical1", None, 10), - submitted_logical_locator_values("logical2", None, 10), + logical_locator_values("logical1", None, 10), + logical_locator_values("logical2", None, 10), ], }, EventCase { @@ -305,13 +292,12 @@ fn event_cases() -> Vec { TableDdlLocator::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "alter") .with_table_id(11), Some("drop_columns"), - EventContext::default(), ), payload: json!({ "version": TABLE_DDL_PAYLOAD_VERSION, "kind": "drop_columns", }), - rows: vec![submitted_table_locator_values(Some("alter"), Some(11))], + rows: vec![table_locator_values(Some("alter"), Some(11))], }, EventCase { event_type: TableDdlEventType::AlterLogicalTables, @@ -324,7 +310,6 @@ fn event_cases() -> Vec { ], 2, ["rename_table", "add_columns", "add_columns"], - EventContext::default(), ), payload: json!({ "version": TABLE_DDL_PAYLOAD_VERSION, @@ -332,8 +317,8 @@ fn event_cases() -> Vec { "kinds": ["add_columns", "rename_table"], }), rows: vec![ - submitted_logical_locator_values("logical1", None, 10), - submitted_logical_locator_values("logical2", None, 10), + logical_locator_values("logical1", None, 10), + logical_locator_values("logical2", None, 10), ], }, EventCase { @@ -342,33 +327,26 @@ fn event_cases() -> Vec { TableDdlLocator::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "drop") .with_table_id(12), true, - EventContext::default(), ), payload: json!({ "version": TABLE_DDL_PAYLOAD_VERSION, "drop_if_exists": true, }), - rows: vec![submitted_table_locator_values(Some("drop"), Some(12))], + rows: vec![table_locator_values(Some("drop"), Some(12))], }, #[cfg(feature = "enterprise")] EventCase { event_type: TableDdlEventType::UndropTable, - event: TableDdlEvent::undrop_table_submitted( - TableDdlLocator::from_table_id(13), - EventContext::default(), - ), + event: TableDdlEvent::undrop_table_submitted(TableDdlLocator::from_table_id(13)), payload: json!({"version": TABLE_DDL_PAYLOAD_VERSION}), - rows: vec![submitted_table_locator_values(None, Some(13))], + rows: vec![table_locator_values(None, Some(13))], }, #[cfg(feature = "enterprise")] EventCase { event_type: TableDdlEventType::PurgeDroppedTable, - event: TableDdlEvent::purge_dropped_table_submitted( - TableDdlLocator::from_table_id(14), - EventContext::default(), - ), + event: TableDdlEvent::purge_dropped_table_submitted(TableDdlLocator::from_table_id(14)), payload: json!({"version": TABLE_DDL_PAYLOAD_VERSION}), - rows: vec![submitted_table_locator_values(None, Some(14))], + rows: vec![table_locator_values(None, Some(14))], }, EventCase { event_type: TableDdlEventType::TruncateTable, @@ -376,41 +354,29 @@ fn event_cases() -> Vec { TableDdlLocator::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "truncate") .with_table_id(15), 4, - EventContext::default(), ), payload: json!({ "version": TABLE_DDL_PAYLOAD_VERSION, "time_range_count": 4, }), - rows: vec![submitted_table_locator_values(Some("truncate"), Some(15))], + rows: vec![table_locator_values(Some("truncate"), Some(15))], }, ] } fn procedure_cases() -> Vec { - let create_table = CreateTableProcedure::new( - test_create_table_task("create"), - QueryContext::default(), - EventContext::default(), - test_context(), - ) - .unwrap(); + let create_table = + CreateTableProcedure::new(test_create_table_task("create"), test_context()).unwrap(); let create_logical_tables = CreateLogicalTablesProcedure::new( vec![ test_create_logical_table_task("logical1"), test_create_logical_table_task("logical2"), ], 41, - EventContext::default(), test_context(), ); - let alter_table = AlterTableProcedure::new( - 42, - test_alter_table_task("alter"), - EventContext::default(), - test_context(), - ) - .unwrap(); + let alter_table = + AlterTableProcedure::new(42, test_alter_table_task("alter"), test_context()).unwrap(); let alter_logical_tables = AlterLogicalTablesProcedure::new( vec![ make_alter_logical_table_add_column_task( @@ -425,7 +391,6 @@ fn procedure_cases() -> Vec { ), ], 43, - EventContext::default(), test_context(), ); let drop_table = DropTableProcedure::new( @@ -437,20 +402,12 @@ fn procedure_cases() -> Vec { drop_if_exists: true, }, test_context(), - EventContext::default(), ); #[cfg(feature = "enterprise")] - let undrop_table = UndropTableProcedure::new( - UndropTableTask { table_id: 45 }, - test_context(), - EventContext::default(), - ); + let undrop_table = UndropTableProcedure::new(UndropTableTask { table_id: 45 }, test_context()); #[cfg(feature = "enterprise")] - let purge_dropped_table = PurgeDroppedTableProcedure::new_if_expired( - PurgeDroppedTableTask { table_id: 46 }, - test_context(), - EventContext::new(TriggerReason::ScheduledGc), - ); + let purge_dropped_table = + PurgeDroppedTableProcedure::new(PurgeDroppedTableTask { table_id: 46 }, test_context()); let truncate_table = truncate_procedure(TruncateTableTask { catalog: DEFAULT_CATALOG_NAME.to_string(), schema: DEFAULT_SCHEMA_NAME.to_string(), @@ -471,7 +428,7 @@ fn procedure_cases() -> Vec { "create_if_not_exists": false, "engine": "mito2", }), - rows: vec![submitted_table_locator_values(Some("create"), None)], + rows: vec![table_locator_values(Some("create"), None)], }, ProcedureCase { procedure: Box::new(create_logical_tables), @@ -481,8 +438,8 @@ fn procedure_cases() -> Vec { "table_count": 2, }), rows: vec![ - submitted_logical_locator_values("logical1", None, 41), - submitted_logical_locator_values("logical2", None, 41), + logical_locator_values("logical1", None, 41), + logical_locator_values("logical2", None, 41), ], }, ProcedureCase { @@ -492,7 +449,7 @@ fn procedure_cases() -> Vec { "version": TABLE_DDL_PAYLOAD_VERSION, "kind": "drop_columns", }), - rows: vec![submitted_table_locator_values(Some("alter"), Some(42))], + rows: vec![table_locator_values(Some("alter"), Some(42))], }, ProcedureCase { procedure: Box::new(alter_logical_tables), @@ -503,8 +460,8 @@ fn procedure_cases() -> Vec { "kinds": ["add_columns"], }), rows: vec![ - submitted_logical_locator_values("logical1", None, 43), - submitted_logical_locator_values("logical2", None, 43), + logical_locator_values("logical1", None, 43), + logical_locator_values("logical2", None, 43), ], }, ProcedureCase { @@ -514,25 +471,21 @@ fn procedure_cases() -> Vec { "version": TABLE_DDL_PAYLOAD_VERSION, "drop_if_exists": true, }), - rows: vec![submitted_table_locator_values(Some("drop"), Some(44))], + rows: vec![table_locator_values(Some("drop"), Some(44))], }, #[cfg(feature = "enterprise")] ProcedureCase { procedure: Box::new(undrop_table), event_type: "undrop_table", payload: json!({"version": TABLE_DDL_PAYLOAD_VERSION}), - rows: vec![submitted_table_locator_values(None, Some(45))], + rows: vec![table_locator_values(None, Some(45))], }, #[cfg(feature = "enterprise")] ProcedureCase { procedure: Box::new(purge_dropped_table), event_type: "purge_dropped_table", payload: json!({"version": TABLE_DDL_PAYLOAD_VERSION}), - rows: vec![submitted_table_locator_values_with_event_context( - None, - Some(46), - EventContext::new(TriggerReason::ScheduledGc), - )], + rows: vec![table_locator_values(None, Some(46))], }, ProcedureCase { procedure: Box::new(truncate_table), @@ -541,7 +494,7 @@ fn procedure_cases() -> Vec { "version": TABLE_DDL_PAYLOAD_VERSION, "time_range_count": 1, }), - rows: vec![submitted_table_locator_values(Some("truncate"), Some(47))], + rows: vec![table_locator_values(Some("truncate"), Some(47))], }, ] } @@ -562,7 +515,6 @@ fn expected_schema(event_type: TableDdlEventType) -> Vec<(&'static str, i32)> { ColumnDataType::Uint32 as i32, )); } - schema.push((EVENT_CONTEXT_COLUMN.name(), ColumnDataType::Binary as i32)); schema } @@ -575,14 +527,12 @@ fn table_locator_values(table_name: Option<&str>, table_id: Option) -> Vec< } else { (Value::default(), Value::default()) }; - let mut values = vec![ + vec![ catalog_name, schema_name, table_name.map(string_value).unwrap_or_default(), table_id.map(table_id_value).unwrap_or_default(), - ]; - values.push(Value::default()); - values + ] } fn logical_locator_values( @@ -591,44 +541,10 @@ fn logical_locator_values( physical_table_id: u32, ) -> Vec { let mut values = table_locator_values(Some(table_name), table_id); - let event_context = values.pop().unwrap(); values.push(table_id_value(physical_table_id)); - values.push(event_context); values } -fn submitted_table_locator_values(table_name: Option<&str>, table_id: Option) -> Vec { - submitted_table_locator_values_with_event_context(table_name, table_id, EventContext::default()) -} - -fn submitted_table_locator_values_with_event_context( - table_name: Option<&str>, - table_id: Option, - event_context: EventContext, -) -> Vec { - let mut values = table_locator_values(table_name, table_id); - *values.last_mut().unwrap() = event_context_value_for(event_context); - values -} - -fn submitted_logical_locator_values( - table_name: &str, - table_id: Option, - physical_table_id: u32, -) -> Vec { - let mut values = logical_locator_values(table_name, table_id, physical_table_id); - *values.last_mut().unwrap() = event_context_value(); - values -} - -fn event_context_value() -> Value { - event_context_value_for(EventContext::default()) -} - -fn event_context_value_for(event_context: EventContext) -> Value { - jsonb_value(&serde_json::to_value(event_context).unwrap()) -} - fn string_value(value: &str) -> Value { ValueData::StringValue(value.to_string()).into() } @@ -643,7 +559,6 @@ fn truncate_procedure(task: TruncateTableTask) -> TruncateTableProcedure { task, DeserializedValueWithBytes::from_inner(TableInfoValue::new(table_info)), test_context(), - EventContext::default(), ) } @@ -661,11 +576,12 @@ fn event_for_state( lifecycle_state: &ProcedureState, ) -> Box { procedure - .event(&EventRuntimeContext { + .event(&EventContext { procedure_id: ProcedureId::random(), lifecycle_state, trigger, event_type_filter: Arc::new(EventTypeFilter::All), + event_context: None, }) .unwrap() } diff --git a/src/common/meta/src/ddl/tests/event/test_util.rs b/src/common/meta/src/ddl/tests/event/test_util.rs index 281b5deb488..e20d294180e 100644 --- a/src/common/meta/src/ddl/tests/event/test_util.rs +++ b/src/common/meta/src/ddl/tests/event/test_util.rs @@ -15,32 +15,21 @@ use std::collections::HashSet; use std::sync::Arc; -use api::v1::Value; use common_event_recorder::EventTypeFilter; -use common_event_recorder::event_table::jsonb_value; -use common_procedure::{EventRuntimeContext, EventTrigger, Procedure, ProcedureId, ProcedureState}; - -use crate::rpc::ddl::EventContext; - -pub(crate) fn default_event_context_value() -> Value { - jsonb_value(&serde_json::to_value(EventContext::default()).unwrap()) -} - -pub(crate) fn procedure_trigger_value(trigger: &str) -> Value { - jsonb_value(&serde_json::json!({"type": trigger})) -} +use common_procedure::{EventContext, EventTrigger, Procedure, ProcedureId, ProcedureState}; pub(crate) fn assert_event_filter(procedure: &dyn Procedure, event_type: &str) { let state = ProcedureState::Running; - let runtime_context = |event_type_filter| EventRuntimeContext { + let event_context = |event_type_filter| EventContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Submitted, event_type_filter: Arc::new(event_type_filter), + event_context: None, }; let allowed = procedure - .event(&runtime_context(EventTypeFilter::Only(HashSet::from([ + .event(&event_context(EventTypeFilter::Only(HashSet::from([ event_type.to_string(), ])))) .unwrap(); @@ -49,7 +38,7 @@ pub(crate) fn assert_event_filter(procedure: &dyn Procedure, event_type: &str) { for denied in [HashSet::from(["other_event".to_string()]), HashSet::new()] { assert!( procedure - .event(&runtime_context(EventTypeFilter::Only(denied))) + .event(&event_context(EventTypeFilter::Only(denied))) .is_none() ); } diff --git a/src/common/meta/src/ddl/tests/event/view.rs b/src/common/meta/src/ddl/tests/event/view.rs index cabe7418ffb..5cb5a6bb7d9 100644 --- a/src/common/meta/src/ddl/tests/event/view.rs +++ b/src/common/meta/src/ddl/tests/event/view.rs @@ -19,16 +19,16 @@ use api::v1::{ColumnSchema, Row, Value}; use common_event_recorder::event_table::{ CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, - VIEW_NAME_COLUMN, + VIEW_NAME_COLUMN, jsonb_value, }; use common_event_recorder::testing::assert_event_contract; use common_event_recorder::{Event, EventTypeFilter}; use common_procedure::{ - ChildSubmissionOutcome, EventRuntimeContext, EventTrigger, Output, Procedure, ProcedureEvent, + ChildSubmissionOutcome, EventContext, EventTrigger, Output, Procedure, ProcedureEvent, ProcedureId, ProcedureState, RetryPhase, }; -use super::test_util::{assert_event_filter, default_event_context_value, procedure_trigger_value}; +use super::test_util::assert_event_filter; use crate::ddl::create_view::CreateViewProcedure; use crate::ddl::drop_view::DropViewProcedure; use crate::ddl::event::view::{ @@ -36,7 +36,6 @@ use crate::ddl::event::view::{ }; use crate::ddl::tests::create_view::test_create_view_task; use crate::ddl::tests::drop_view::new_drop_view_task; -use crate::rpc::ddl::EventContext; use crate::test_util::{MockDatanodeManager, new_ddl_context}; #[test] @@ -44,7 +43,7 @@ fn test_view_submitted_event_contracts() { let mut task = test_create_view_task("v_metrics"); task.create_view.or_replace = true; task.create_view.create_if_not_exists = true; - let create = CreateViewProcedure::new(task, EventContext::default(), test_context()); + let create = CreateViewProcedure::new(task, test_context()); let event = event_for(&create, EventTrigger::Submitted); assert_view_event_contract( @@ -72,11 +71,7 @@ fn test_view_submitted_event_contracts() { assert!(!payload.contains(omitted)); } - let drop = DropViewProcedure::new( - new_drop_view_task("view_name", 42, true), - EventContext::default(), - test_context(), - ); + let drop = DropViewProcedure::new(new_drop_view_task("view_name", 42, true), test_context()); let event = event_for(&drop, EventTrigger::Submitted); assert_view_event_contract( @@ -140,16 +135,8 @@ fn test_view_lifecycle_event_contracts() { #[test] fn test_view_procedures_preserve_lifecycle_locators() { - let create = CreateViewProcedure::new( - test_create_view_task("view_name"), - EventContext::default(), - test_context(), - ); - let drop = DropViewProcedure::new( - new_drop_view_task("view_name", 42, false), - EventContext::default(), - test_context(), - ); + let create = CreateViewProcedure::new(test_create_view_task("view_name"), test_context()); + let drop = DropViewProcedure::new(new_drop_view_task("view_name", 42, false), test_context()); let triggers = [ EventTrigger::Recovered, EventTrigger::ChildSubmitted { @@ -199,11 +186,7 @@ fn test_view_procedures_preserve_lifecycle_locators() { #[test] fn test_create_view_succeeded_output_mapping() { - let procedure = CreateViewProcedure::new( - test_create_view_task("view_name"), - EventContext::default(), - test_context(), - ); + let procedure = CreateViewProcedure::new(test_create_view_task("view_name"), test_context()); let state = ProcedureState::Done { output: Some(Arc::new(84_u32)), }; @@ -240,21 +223,14 @@ fn test_create_view_succeeded_output_mapping() { #[test] fn test_create_view_event_filter() { - let procedure = CreateViewProcedure::new( - test_create_view_task("view_name"), - EventContext::default(), - test_context(), - ); + let procedure = CreateViewProcedure::new(test_create_view_task("view_name"), test_context()); assert_event_filter(&procedure, CREATE_VIEW_EVENT_TYPE); } #[test] fn test_drop_view_event_filter() { - let procedure = DropViewProcedure::new( - new_drop_view_task("view_name", 42, false), - EventContext::default(), - test_context(), - ); + let procedure = + DropViewProcedure::new(new_drop_view_task("view_name", 42, false), test_context()); assert_event_filter(&procedure, DROP_VIEW_EVENT_TYPE); } @@ -273,7 +249,6 @@ fn test_view_event_procedure_envelope_contract() { referenced_table_count: 1, column_count: 1, }, - EventContext::default(), )), ProcedureState::Running, EventTrigger::Submitted, @@ -344,17 +319,11 @@ fn view_schema() -> Vec { SCHEMA_NAME_COLUMN.column_schema(), VIEW_NAME_COLUMN.column_schema(), VIEW_ID_COLUMN.column_schema(), - EVENT_CONTEXT_COLUMN.column_schema(), ] } fn assert_view_event_contract(event: &dyn Event, event_type: &str, locator: ViewEventLocator<'_>) { - let mut values = locator.values(); - values.push(if event.json_payload().unwrap().is_null() { - Value { value_data: None } - } else { - default_event_context_value() - }); + let values = locator.values(); assert_event_contract(event, event_type, &view_schema(), &[Row { values }]); } @@ -372,19 +341,16 @@ fn assert_procedure_event_contract( PROCEDURE_TRIGGER_COLUMN.column_schema(), ]; schema.extend(view_schema()); + schema.push(EVENT_CONTEXT_COLUMN.column_schema()); let mut values = vec![ ValueData::StringValue(event.procedure_id.to_string()).into(), ValueData::StringValue(state.to_string()).into(), ValueData::StringValue(String::new()).into(), - procedure_trigger_value(trigger), + jsonb_value(&serde_json::json!({"type": trigger})), ]; values.extend(locator.values()); - values.push(if trigger == "Submitted" { - default_event_context_value() - } else { - Value { value_data: None } - }); + values.push(Value { value_data: None }); assert_event_contract(event, event_type, &schema, &[Row { values }]); } @@ -403,11 +369,12 @@ fn event_for_state( lifecycle_state: &ProcedureState, ) -> Box { procedure - .event(&EventRuntimeContext { + .event(&EventContext { procedure_id: ProcedureId::random(), lifecycle_state, trigger, event_type_filter: Arc::new(EventTypeFilter::All), + event_context: None, }) .unwrap() } diff --git a/src/common/meta/src/ddl/truncate_table.rs b/src/common/meta/src/ddl/truncate_table.rs index 6a561b62afb..67d74361983 100644 --- a/src/common/meta/src/ddl/truncate_table.rs +++ b/src/common/meta/src/ddl/truncate_table.rs @@ -20,7 +20,7 @@ use api::v1::region::{ use async_trait::async_trait; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, Result as ProcedureResult, Status, }; use common_telemetry::debug; @@ -43,7 +43,7 @@ use crate::key::table_info::TableInfoValue; use crate::key::table_name::TableNameKey; use crate::lock_key::{CatalogLock, SchemaLock, TableLock}; use crate::metrics; -use crate::rpc::ddl::{EventContext, TruncateTableTask}; +use crate::rpc::ddl::TruncateTableTask; use crate::rpc::router::{find_leader_regions, find_leaders}; pub struct TruncateTableProcedure { @@ -89,10 +89,7 @@ impl Procedure for TruncateTableProcedure { LockKey::new(lock_key) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::TruncateTable.as_str()) @@ -103,11 +100,9 @@ impl Procedure for TruncateTableProcedure { let locator = TableDdlLocator::new(&task.catalog, &task.schema, &task.table) .with_table_id(task.table_id); let event = match &ctx.trigger { - EventTrigger::Submitted => TableDdlEvent::truncate_table_submitted( - locator, - task.time_ranges.len(), - self.data.event_context.clone(), - ), + EventTrigger::Submitted => { + TableDdlEvent::truncate_table_submitted(locator, task.time_ranges.len()) + } _ => TableDdlEvent::lifecycle(TableDdlEventType::TruncateTable, [locator]), }; @@ -122,11 +117,10 @@ impl TruncateTableProcedure { task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, context: DdlContext, - event_context: EventContext, ) -> Self { Self { context, - data: TruncateTableData::new(task, table_info_value, event_context), + data: TruncateTableData::new(task, table_info_value), } } @@ -233,21 +227,17 @@ pub struct TruncateTableData { state: TruncateTableState, task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, - #[serde(default)] - event_context: EventContext, } impl TruncateTableData { pub fn new( task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, - event_context: EventContext, ) -> Self { Self { state: TruncateTableState::Prepare, task, table_info_value, - event_context, } } diff --git a/src/common/meta/src/ddl/undrop_table.rs b/src/common/meta/src/ddl/undrop_table.rs index 3295b1163e2..73a361ffd0e 100644 --- a/src/common/meta/src/ddl/undrop_table.rs +++ b/src/common/meta/src/ddl/undrop_table.rs @@ -17,7 +17,7 @@ use api::v1::region::{ use async_trait::async_trait; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, + Context as ProcedureContext, EventContext, EventTrigger, LockKey, Procedure, Result as ProcedureResult, Status, }; use common_telemetry::tracing_context::TracingContext; @@ -43,7 +43,7 @@ use crate::instruction::CacheIdent; use crate::key::table_name::TableNameKey; use crate::key::table_route::TableRouteValue; use crate::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; -use crate::rpc::ddl::{EventContext, UndropTableTask}; +use crate::rpc::ddl::UndropTableTask; use crate::rpc::router::{ RegionRoute, find_follower_regions, find_followers, find_leader_regions, find_leaders, }; @@ -56,17 +56,16 @@ pub struct UndropTableProcedure { impl UndropTableProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::UndropTable"; - pub fn new(task: UndropTableTask, context: DdlContext, event_context: EventContext) -> Self { - Self::new_with_original_table_name(task, context, None, event_context) + pub fn new(task: UndropTableTask, context: DdlContext) -> Self { + Self::new_with_original_table_name(task, context, None) } pub(crate) fn new_with_original_table_name( task: UndropTableTask, context: DdlContext, table_name: Option, - event_context: EventContext, ) -> Self { - let mut data = UndropTableData::new(task, event_context); + let mut data = UndropTableData::new(task); data.table_name = table_name; Self { context, data } } @@ -444,10 +443,7 @@ impl Procedure for UndropTableProcedure { LockKey::new(lock_key) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx .event_type_filter .allows(TableDdlEventType::UndropTable.as_str()) @@ -456,9 +452,7 @@ impl Procedure for UndropTableProcedure { } let locator = self.event_locator(); let event = match &ctx.trigger { - EventTrigger::Submitted => { - TableDdlEvent::undrop_table_submitted(locator, self.data.event_context.clone()) - } + EventTrigger::Submitted => TableDdlEvent::undrop_table_submitted(locator), _ => TableDdlEvent::lifecycle(TableDdlEventType::UndropTable, [locator]), }; @@ -565,12 +559,10 @@ pub struct UndropTableData { drop_generation: Option, #[serde(default)] tombstone_identity_loaded: bool, - #[serde(default)] - event_context: EventContext, } impl UndropTableData { - fn new(task: UndropTableTask, event_context: EventContext) -> Self { + fn new(task: UndropTableTask) -> Self { Self { state: UndropTableState::Prepare, task, @@ -582,7 +574,6 @@ impl UndropTableData { retention_expires_at: None, drop_generation: None, tombstone_identity_loaded: false, - event_context, } } @@ -647,11 +638,7 @@ mod tests { #[test] fn test_map_restore_metadata_error_without_table_name() { let context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); - let procedure = UndropTableProcedure::new( - UndropTableTask { table_id: 42 }, - context, - EventContext::default(), - ); + let procedure = UndropTableProcedure::new(UndropTableTask { table_id: 42 }, context); let err = procedure.map_restore_metadata_error( error::TombstoneTargetAlreadyExistsSnafu { @@ -666,11 +653,8 @@ mod tests { #[test] fn test_recovered_legacy_snapshot_reloads_tombstone_identity() { let context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); - let mut procedure = UndropTableProcedure::new( - UndropTableTask { table_id: 42 }, - context.clone(), - EventContext::default(), - ); + let mut procedure = + UndropTableProcedure::new(UndropTableTask { table_id: 42 }, context.clone()); procedure.data.state = UndropTableState::OpenRegions; let mut data: serde_json::Value = serde_json::from_str(&procedure.dump().unwrap()).unwrap(); let data = data.as_object_mut().unwrap(); @@ -695,11 +679,7 @@ mod tests { let table_id = 1024; let region_id = RegionId::new(table_id, 1); - let mut procedure = UndropTableProcedure::new( - UndropTableTask { table_id }, - context, - EventContext::default(), - ); + let mut procedure = UndropTableProcedure::new(UndropTableTask { table_id }, context); procedure.data.table_route_value = Some(TableRouteValue::physical(vec![RegionRoute { region: Region::new_test(region_id), leader_peer: Some(Peer::empty(1)), @@ -738,11 +718,7 @@ mod tests { let table_id = 1024; let region_id = RegionId::new(table_id, 1); - let mut procedure = UndropTableProcedure::new( - UndropTableTask { table_id }, - context, - EventContext::default(), - ); + let mut procedure = UndropTableProcedure::new(UndropTableTask { table_id }, context); procedure.data.table_name = Some(TableName::new( DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index 137a16793d7..2eed4628ac3 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -78,8 +78,8 @@ use crate::rpc::ddl::trigger::DropTriggerTask; use crate::rpc::ddl::{ AlterDatabaseTask, AlterTableTask, CommentOnTask, CreateDatabaseTask, CreateFlowTask, CreateTableTask, CreateViewTask, DropDatabaseTask, DropFlowTask, DropTableTask, DropViewTask, - EventContext, PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest, SubmitDdlTaskResponse, - TruncateTableTask, UndropTableTask, + PersistentEventContext, PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest, + SubmitDdlTaskResponse, TruncateTableTask, UndropTableTask, }; const MAX_REGION_ROUTE_CHANGE_RETRIES: usize = 3; @@ -188,7 +188,6 @@ pub trait RepartitionProcedureFactory: Send + Sync { source: RepartitionSource, to_exprs: Vec, timeout: Option, - event_context: EventContext, ) -> std::result::Result; fn register_loaders( @@ -342,7 +341,7 @@ impl DdlManager { repartition: Repartition, wait: bool, timeout: Duration, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); @@ -378,14 +377,14 @@ impl DdlManager { source, into_partition_exprs, Some(timeout), - event_context, ) .context(CreateRepartitionProcedureSnafu)?; self.repartition_procedure_factory .ensure_gc_requirement() .await .context(PersistRepartitionGcRequirementSnafu)?; - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); if wait { self.execute_procedure_and_wait(procedure_with_id).await } else { @@ -401,7 +400,7 @@ impl DdlManager { &self, table_id: TableId, alter_table_task: AlterTableTask, - event_context: EventContext, + event_context: PersistentEventContext, ddl_options: DdlOptions, ) -> Result<(ProcedureId, Option)> { // make alter_table_task mutable so we can call .take() on its field @@ -456,12 +455,12 @@ impl DdlManager { let procedure = AlterTableProcedure::new_with_region_locks( table_id, alter_table_task.clone(), - event_context.clone(), region_ids_to_lock, context, )?; - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) + .with_event_context(event_context.clone()); let result = self.execute_procedure_and_wait(procedure_with_id).await?; if result .1 @@ -490,14 +489,18 @@ impl DdlManager { &self, create_table_task: CreateTableTask, query_context: QueryContext, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = - CreateTableProcedure::new(create_table_task, query_context, event_context, context)?; + let procedure = CreateTableProcedure::new_with_query_context( + create_table_task, + query_context, + context, + )?; - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -507,13 +510,14 @@ impl DdlManager { pub async fn submit_create_view_task( &self, create_view_task: CreateViewTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = CreateViewProcedure::new(create_view_task, event_context, context); + let procedure = CreateViewProcedure::new(create_view_task, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -524,18 +528,15 @@ impl DdlManager { &self, create_table_tasks: Vec, physical_table_id: TableId, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = CreateLogicalTablesProcedure::new( - create_table_tasks, - physical_table_id, - event_context, - context, - ); + let procedure = + CreateLogicalTablesProcedure::new(create_table_tasks, physical_table_id, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -546,18 +547,15 @@ impl DdlManager { &self, alter_table_tasks: Vec, physical_table_id: TableId, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = AlterLogicalTablesProcedure::new( - alter_table_tasks, - physical_table_id, - event_context, - context, - ); + let procedure = + AlterLogicalTablesProcedure::new(alter_table_tasks, physical_table_id, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -567,13 +565,14 @@ impl DdlManager { pub async fn submit_drop_table_task( &self, drop_table_task: DropTableTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = DropTableProcedure::new(drop_table_task, context, event_context); + let procedure = DropTableProcedure::new(drop_table_task, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -584,7 +583,7 @@ impl DdlManager { pub async fn submit_undrop_table_task( &self, undrop_table_task: UndropTableTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { #[cfg(not(feature = "enterprise"))] { @@ -610,9 +609,9 @@ impl DdlManager { undrop_table_task, context, Some(original_table_name), - event_context, ); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) + .with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -624,7 +623,7 @@ impl DdlManager { pub async fn submit_purge_dropped_table_task( &self, purge_dropped_table_task: PurgeDroppedTableTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { #[cfg(not(feature = "enterprise"))] { @@ -638,9 +637,9 @@ impl DdlManager { #[cfg(feature = "enterprise")] { let context = self.create_context(); - let procedure = - PurgeDroppedTableProcedure::new(purge_dropped_table_task, context, event_context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure = PurgeDroppedTableProcedure::new(purge_dropped_table_task, context); + let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) + .with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -666,12 +665,11 @@ impl DdlManager { #[cfg(feature = "enterprise")] { let context = self.create_context(); - let procedure = PurgeDroppedTableProcedure::new_if_expired( - purge_dropped_table_task, - context, - EventContext::new(TriggerReason::ScheduledGc), - ); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let event_context = PersistentEventContext::new(TriggerReason::ScheduledGc); + let procedure = + PurgeDroppedTableProcedure::new_if_expired(purge_dropped_table_task, context); + let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) + .with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -688,7 +686,7 @@ impl DdlManager { options, creator, }: CreateDatabaseTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = CreateDatabaseProcedure::new( @@ -697,10 +695,10 @@ impl DdlManager { create_if_not_exists, options, creator, - event_context, context, ); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -714,12 +712,12 @@ impl DdlManager { schema, drop_if_exists, }: DropDatabaseTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = - DropDatabaseProcedure::new(catalog, schema, drop_if_exists, event_context, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure = DropDatabaseProcedure::new(catalog, schema, drop_if_exists, context); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -727,11 +725,12 @@ impl DdlManager { pub async fn submit_alter_database( &self, alter_database_task: AlterDatabaseTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = AlterDatabaseProcedure::new(alter_database_task, event_context, context)?; - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure = AlterDatabaseProcedure::new(alter_database_task, context)?; + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -742,12 +741,12 @@ impl DdlManager { &self, create_flow: CreateFlowTask, query_context: QueryContext, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = - CreateFlowProcedure::new(create_flow, query_context, event_context, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure = CreateFlowProcedure::new(create_flow, query_context, context); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -757,11 +756,12 @@ impl DdlManager { pub async fn submit_drop_flow_task( &self, drop_flow: DropFlowTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = DropFlowProcedure::new(drop_flow, event_context, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure = DropFlowProcedure::new(drop_flow, context); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -771,11 +771,12 @@ impl DdlManager { pub async fn submit_drop_view_task( &self, drop_view: DropViewTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = DropViewProcedure::new(drop_view, event_context, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure = DropViewProcedure::new(drop_view, context); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -786,17 +787,13 @@ impl DdlManager { &self, truncate_table_task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = TruncateTableProcedure::new( - truncate_table_task, - table_info_value, - context, - event_context, - ); + let procedure = TruncateTableProcedure::new(truncate_table_task, table_info_value, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -806,6 +803,7 @@ impl DdlManager { pub async fn submit_comment_on_task( &self, mut comment_on_task: CommentOnTask, + event_context: PersistentEventContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); comment_on_task @@ -815,7 +813,8 @@ impl DdlManager { ) .await?; let procedure = CommentOnProcedure::new(comment_on_task, context); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -869,7 +868,10 @@ impl DdlManager { timeout, task, } = request; - let event_context = EventContext::from_query_context(&query_context); + let event_context = ctx + .event_context + .clone() + .unwrap_or_else(|| crate::rpc::ddl::event_context_from_query_context(&query_context)); let ddl_options = DdlOptions { wait, timeout }; async move { debug!("Submitting Ddl task: {:?}", task); @@ -924,7 +926,9 @@ impl DdlManager { DropView(drop_view_task) => { handle_drop_view_task(self, drop_view_task, event_context).await } - CommentOn(comment_on_task) => handle_comment_on_task(self, comment_on_task).await, + CommentOn(comment_on_task) => { + handle_comment_on_task(self, comment_on_task, event_context).await + } #[cfg(feature = "enterprise")] CreateTrigger(create_trigger_task) => { handle_create_trigger_task(self, create_trigger_task, query_context).await @@ -943,7 +947,7 @@ impl DdlManager { async fn handle_truncate_table_task( ddl_manager: &DdlManager, truncate_table_task: TruncateTableTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let table_id = truncate_table_task.table_id; let table_metadata_manager = &ddl_manager.table_metadata_manager(); @@ -983,7 +987,7 @@ async fn handle_alter_table_task( ddl_manager: &DdlManager, alter_table_task: AlterTableTask, ddl_options: DdlOptions, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let table_ref = alter_table_task.table_ref(); @@ -1030,7 +1034,7 @@ async fn handle_alter_table_task( async fn handle_drop_table_task( ddl_manager: &DdlManager, drop_table_task: DropTableTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let table_id = drop_table_task.table_id; let (id, _) = ddl_manager @@ -1048,7 +1052,7 @@ async fn handle_drop_table_task( async fn handle_undrop_table_task( ddl_manager: &DdlManager, undrop_table_task: UndropTableTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let table_id = undrop_table_task.table_id; let (id, _) = ddl_manager @@ -1066,7 +1070,7 @@ async fn handle_undrop_table_task( async fn handle_purge_dropped_table_task( ddl_manager: &DdlManager, purge_dropped_table_task: PurgeDroppedTableTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let (id, _) = ddl_manager .submit_purge_dropped_table_task(purge_dropped_table_task, event_context) @@ -1084,7 +1088,7 @@ async fn handle_create_table_task( ddl_manager: &DdlManager, create_table_task: CreateTableTask, query_context: QueryContext, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let (id, output) = ddl_manager .submit_create_table_task(create_table_task, query_context, event_context) @@ -1110,7 +1114,7 @@ async fn handle_create_table_task( async fn handle_create_logical_table_tasks( ddl_manager: &DdlManager, create_table_tasks: Vec, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { ensure!( !create_table_tasks.is_empty(), @@ -1155,7 +1159,7 @@ async fn handle_create_logical_table_tasks( async fn handle_create_database_task( ddl_manager: &DdlManager, create_database_task: CreateDatabaseTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let catalog = create_database_task.catalog.clone(); let schema = create_database_task.schema.clone(); @@ -1178,7 +1182,7 @@ async fn handle_create_database_task( async fn handle_drop_database_task( ddl_manager: &DdlManager, drop_database_task: DropDatabaseTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let (id, _) = ddl_manager .submit_drop_database(drop_database_task.clone(), event_context) @@ -1199,7 +1203,7 @@ async fn handle_drop_database_task( async fn handle_alter_database_task( ddl_manager: &DdlManager, alter_database_task: AlterDatabaseTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let (id, _) = ddl_manager .submit_alter_database(alter_database_task.clone(), event_context) @@ -1221,7 +1225,7 @@ async fn handle_alter_database_task( async fn handle_drop_flow_task( ddl_manager: &DdlManager, drop_flow_task: DropFlowTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let (id, _) = ddl_manager .submit_drop_flow_task(drop_flow_task.clone(), event_context) @@ -1266,7 +1270,7 @@ async fn handle_drop_trigger_task( async fn handle_drop_view_task( ddl_manager: &DdlManager, drop_view_task: DropViewTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let (id, _) = ddl_manager .submit_drop_view_task(drop_view_task.clone(), event_context) @@ -1289,7 +1293,7 @@ async fn handle_create_flow_task( ddl_manager: &DdlManager, create_flow_task: CreateFlowTask, query_context: QueryContext, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let (id, output) = ddl_manager .submit_create_flow_task(create_flow_task.clone(), query_context, event_context) @@ -1349,7 +1353,7 @@ async fn handle_create_trigger_task( async fn handle_alter_logical_table_tasks( ddl_manager: &DdlManager, alter_table_tasks: Vec, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { ensure!( !alter_table_tasks.is_empty(), @@ -1388,7 +1392,7 @@ async fn handle_alter_logical_table_tasks( async fn handle_create_view_task( ddl_manager: &DdlManager, create_view_task: CreateViewTask, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { let (id, output) = ddl_manager .submit_create_view_task(create_view_task, event_context) @@ -1414,9 +1418,10 @@ async fn handle_create_view_task( async fn handle_comment_on_task( ddl_manager: &DdlManager, comment_on_task: CommentOnTask, + event_context: PersistentEventContext, ) -> Result { let (id, _) = ddl_manager - .submit_comment_on_task(comment_on_task.clone()) + .submit_comment_on_task(comment_on_task.clone(), event_context) .await?; let procedure_id = id.to_string(); @@ -1469,7 +1474,7 @@ mod tests { use crate::procedure_executor::ExecutorContext; use crate::region_keeper::MemoryRegionKeeper; use crate::region_registry::LeaderRegionRegistry; - use crate::rpc::ddl::{CreatorGrantIntent, EventContext, UndropTableTask}; + use crate::rpc::ddl::{CreatorGrantIntent, PersistentEventContext, UndropTableTask}; #[cfg(not(feature = "enterprise"))] use crate::rpc::ddl::{DdlTask, PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest}; use crate::sequence::SequenceBuilder; @@ -1506,7 +1511,6 @@ mod tests { _source: RepartitionSource, _to_exprs: Vec, _timeout: Option, - _event_context: EventContext, ) -> std::result::Result { unimplemented!() } @@ -1680,7 +1684,10 @@ mod tests { let ddl_manager = build_soft_drop_test_ddl_manager(); let err = ddl_manager - .submit_undrop_table_task(UndropTableTask { table_id: 1024 }, EventContext::default()) + .submit_undrop_table_task( + UndropTableTask { table_id: 1024 }, + PersistentEventContext::default(), + ) .await .unwrap_err(); @@ -1694,7 +1701,10 @@ mod tests { let ddl_manager = build_soft_drop_test_ddl_manager(); let err = ddl_manager - .submit_undrop_table_task(UndropTableTask { table_id: 1024 }, EventContext::default()) + .submit_undrop_table_task( + UndropTableTask { table_id: 1024 }, + PersistentEventContext::default(), + ) .await .unwrap_err(); assert!(matches!(err, crate::error::Error::Unsupported { .. })); @@ -1702,7 +1712,7 @@ mod tests { let err = ddl_manager .submit_purge_dropped_table_task( PurgeDroppedTableTask { table_id: 1024 }, - EventContext::default(), + PersistentEventContext::default(), ) .await .unwrap_err(); diff --git a/src/common/meta/src/procedure_executor.rs b/src/common/meta/src/procedure_executor.rs index c0165dd80f4..b812aeb6aac 100644 --- a/src/common/meta/src/procedure_executor.rs +++ b/src/common/meta/src/procedure_executor.rs @@ -15,6 +15,7 @@ use std::sync::Arc; use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse}; +use common_event_recorder::PersistentEventContext; use common_procedure::{ProcedureId, ProcedureManagerRef}; use common_telemetry::tracing_context::W3cTrace; use snafu::{OptionExt, ResultExt}; @@ -30,9 +31,11 @@ use crate::rpc::procedure::{ }; /// The context of procedure executor. -#[derive(Debug, Default)] +#[derive(Debug, Clone, Default)] pub struct ExecutorContext { pub tracing_context: Option, + pub actor: Option, + pub event_context: Option, } /// The procedure executor that accepts ddl, region migration task etc. diff --git a/src/common/meta/src/rpc/ddl.rs b/src/common/meta/src/rpc/ddl.rs index 8aef3325e64..d7966804738 100644 --- a/src/common/meta/src/rpc/ddl.rs +++ b/src/common/meta/src/rpc/ddl.rs @@ -45,12 +45,12 @@ use base64::engine::general_purpose; use common_base::protocol::Channel; use common_catalog::{format_full_flow_name, format_full_table_name}; use common_error::ext::BoxedError; +pub use common_event_recorder::{PersistentEventContext, TriggerReason}; use common_time::{DatabaseTimeToLive, Timestamp}; use prost::Message; use serde::{Deserialize, Serialize}; use serde_with::{DefaultOnNull, serde_as}; use snafu::{OptionExt, ResultExt}; -use strum::{AsRefStr, EnumString}; use table::metadata::{TableId, TableInfo}; use table::requests::validate_database_option; use table::table_name::TableName; @@ -1668,77 +1668,19 @@ pub struct QueryContext { pub sst_min_sequences: HashMap, } -/// The stable context recorded for a procedure event. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct EventContext { - pub reason: TriggerReason, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub protocol: Option, - #[serde(default, skip_serializing_if = "serde_json::Map::is_empty")] - pub extensions: serde_json::Map, -} - -impl EventContext { - /// Creates an event context with no additional extensions. - pub fn new(reason: TriggerReason) -> Self { - Self { - reason, - protocol: None, - extensions: Default::default(), - } - } - - /// Adds the protocol that originated the operation. - pub fn with_protocol(mut self, protocol: impl Into) -> Self { - self.protocol = Some(protocol.into()); - self - } - - /// Builds an event context from frontend query context metadata. - pub fn from_query_context(query_context: &QueryContext) -> Self { - let reason = query_context - .extensions - .get(TRIGGER_REASON_EXTENSION_KEY) - .map(|reason| TriggerReason::from_extension(reason)) - .unwrap_or_default(); - let channel = Channel::from(u32::from(query_context.channel)); - let context = Self::new(reason); - if channel == Channel::Unknown { - context - } else { - context.with_protocol(channel.as_ref()) - } - } -} - -impl Default for EventContext { - fn default() -> Self { - Self::new(TriggerReason::default()) - } -} - -/// The stable classification of a procedure trigger. -#[derive( - Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, AsRefStr, EnumString, -)] -#[serde(rename_all = "snake_case")] -#[strum(serialize_all = "snake_case")] -pub enum TriggerReason { - Manual, - AutoCreate, - AutoAlter, - AutoRepartition, - AutoRebalance, - RegionFailover, - ScheduledGc, - #[default] - #[serde(other)] - Unknown, -} - -impl TriggerReason { - pub fn from_extension(value: &str) -> Self { - value.parse().unwrap_or_default() +/// Builds event context from frontend query context metadata. +pub fn event_context_from_query_context(query_context: &QueryContext) -> PersistentEventContext { + let reason = query_context + .extensions + .get(TRIGGER_REASON_EXTENSION_KEY) + .map(|reason| TriggerReason::from_extension(reason)) + .unwrap_or_default(); + let channel = Channel::from(u32::from(query_context.channel)); + let context = PersistentEventContext::new(reason); + if channel == Channel::Unknown { + context + } else { + context.with_protocol(channel.as_ref()) } } @@ -2323,7 +2265,7 @@ mod tests { #[test] fn test_event_context_serialization() { - let context = EventContext::new(TriggerReason::Manual).with_protocol("mysql"); + let context = PersistentEventContext::new(TriggerReason::Manual).with_protocol("mysql"); assert_eq!( serde_json::json!({ @@ -2335,7 +2277,7 @@ mod tests { assert_eq!( serde_json::json!({ "reason": "manual" }), - serde_json::to_value(EventContext::new(TriggerReason::Manual)).unwrap() + serde_json::to_value(PersistentEventContext::new(TriggerReason::Manual)).unwrap() ); } @@ -2348,7 +2290,7 @@ mod tests { ); query_context.channel = 4; - let event_context = EventContext::from_query_context(&query_context); + let event_context = event_context_from_query_context(&query_context); assert_eq!(event_context.reason, TriggerReason::AutoCreate); assert_eq!(event_context.protocol.as_deref(), Some("prometheus")); diff --git a/src/common/procedure-test/src/lib.rs b/src/common/procedure-test/src/lib.rs index f6d816381d1..2d342ed91ab 100644 --- a/src/common/procedure-test/src/lib.rs +++ b/src/common/procedure-test/src/lib.rs @@ -84,6 +84,7 @@ pub async fn execute_procedure_until_done(procedure: &mut dyn Procedure) -> Opti let ctx = Context { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, }; loop { @@ -110,6 +111,7 @@ pub async fn execute_procedure_once( let ctx = Context { procedure_id, provider: Arc::new(provider), + event_context: None, }; match procedure.execute(&ctx).await.unwrap() { @@ -137,6 +139,7 @@ pub async fn execute_until_suspended_or_done( let ctx = Context { procedure_id, provider: Arc::new(provider), + event_context: None, }; loop { @@ -155,6 +158,7 @@ pub fn new_test_procedure_context() -> Context { Context { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, } } diff --git a/src/common/procedure/src/event.rs b/src/common/procedure/src/event.rs index c3100b34cdd..a4a6c290bf3 100644 --- a/src/common/procedure/src/event.rs +++ b/src/common/procedure/src/event.rs @@ -19,13 +19,13 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, - jsonb_value, procedure_event_column_schemas, + EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, + PROCEDURE_TRIGGER_COLUMN, jsonb_value, nullable_json, procedure_event_column_schemas, }; use common_time::timestamp::{TimeUnit, Timestamp}; use snafu::ResultExt; -use crate::{EventTrigger, ProcedureId, ProcedureState}; +use crate::{EventTrigger, ProcedureContext, ProcedureId, ProcedureState}; pub const EVENTS_TABLE_PROCEDURE_ID_COLUMN_NAME: &str = PROCEDURE_ID_COLUMN.name(); pub const EVENTS_TABLE_PROCEDURE_STATE_COLUMN_NAME: &str = PROCEDURE_STATE_COLUMN.name(); @@ -43,6 +43,8 @@ pub struct ProcedureEvent { pub state: ProcedureState, /// The lifecycle trigger that caused the event to be emitted. pub trigger: EventTrigger, + /// Context associated with the root submission. + pub context: ProcedureContext, /// The event emitted by the procedure. It's generated by [Procedure::event]. pub internal_event: Box, } @@ -53,6 +55,22 @@ impl ProcedureEvent { internal_event: Box, state: ProcedureState, trigger: EventTrigger, + ) -> Self { + Self::new_with_context( + procedure_id, + internal_event, + state, + trigger, + ProcedureContext::default(), + ) + } + + pub fn new_with_context( + procedure_id: ProcedureId, + internal_event: Box, + state: ProcedureState, + trigger: EventTrigger, + context: ProcedureContext, ) -> Self { Self { procedure_id, @@ -60,6 +78,7 @@ impl ProcedureEvent { timestamp: Timestamp::current_time(TimeUnit::Nanosecond), state, trigger, + context, } } } @@ -79,7 +98,9 @@ impl Event for ProcedureEvent { fn extra_schema(&self) -> Vec { let mut schema = procedure_event_column_schemas(); - schema.append(&mut self.internal_event.extra_schema()); + let mut internal_schema = self.internal_event.extra_schema(); + schema.append(&mut internal_schema); + schema.push(EVENT_CONTEXT_COLUMN.column_schema()); schema } @@ -97,9 +118,15 @@ impl Event for ProcedureEvent { _ => String::new(), }; let trigger = serde_json::to_value(&self.trigger).context(SerializeEventSnafu)?; + let event_context = matches!(self.trigger, EventTrigger::Submitted) + .then_some(self.context.event_context.as_ref()) + .flatten() + .map(serde_json::to_value) + .transpose() + .context(SerializeEventSnafu)?; for internal_event_extra_row in internal_event_extra_rows.iter_mut() { - let mut values = Vec::with_capacity(4 + internal_event_extra_row.values.len()); + let mut values = Vec::with_capacity(5 + internal_event_extra_row.values.len()); values.extend([ ValueData::StringValue(procedure_id.clone()).into(), ValueData::StringValue(state.clone()).into(), @@ -107,6 +134,7 @@ impl Event for ProcedureEvent { jsonb_value(&trigger), ]); values.append(&mut internal_event_extra_row.values); + values.push(nullable_json(event_context.as_ref())); rows.push(Row { values }); } @@ -126,13 +154,15 @@ mod tests { use api::v1::{ColumnDataType, ColumnSchema, Row, SemanticType, Value}; use common_error::mock::MockError; use common_error::status_code::StatusCode; - use common_event_recorder::Event; - use common_event_recorder::event_table::{PROCEDURE_TRIGGER_COLUMN, jsonb_value}; + use common_event_recorder::event_table::{ + EVENT_CONTEXT_COLUMN, PROCEDURE_TRIGGER_COLUMN, jsonb_value, + }; + use common_event_recorder::{Event, PersistentEventContext, TriggerReason}; use serde_json::json; use crate::{ - ChildSubmissionOutcome, Error, EventTrigger, ProcedureEvent, ProcedureId, ProcedureState, - RetryPhase, + ChildSubmissionOutcome, Error, EventTrigger, ProcedureContext, ProcedureEvent, ProcedureId, + ProcedureState, RetryPhase, }; #[derive(Debug)] @@ -192,6 +222,7 @@ mod tests { ValueData::StringValue(String::new()).into(), jsonb_value(&json!({"type": "Submitted"})), ValueData::StringValue("test_event1".to_string()).into(), + Value { value_data: None }, ], }, Row { @@ -201,6 +232,7 @@ mod tests { ValueData::StringValue(String::new()).into(), jsonb_value(&json!({"type": "Submitted"})), ValueData::StringValue("test_event2".to_string()).into(), + Value { value_data: None }, ], }, Row { @@ -210,6 +242,7 @@ mod tests { ValueData::StringValue(String::new()).into(), jsonb_value(&json!({"type": "Submitted"})), Value { value_data: None }, + Value { value_data: None }, ], }, ] @@ -280,10 +313,41 @@ mod tests { semantic_type: SemanticType::Field.into(), ..Default::default() }, + EVENT_CONTEXT_COLUMN.column_schema(), ] ); } + #[test] + fn procedure_event_only_records_context_on_submission() { + let context = ProcedureContext::from_event_context(PersistentEventContext::new( + TriggerReason::AutoRepartition, + )); + let submitted = ProcedureEvent::new_with_context( + ProcedureId::random(), + Box::new(TestEvent {}), + ProcedureState::Running, + EventTrigger::Submitted, + context.clone(), + ); + let succeeded = ProcedureEvent::new_with_context( + ProcedureId::random(), + Box::new(TestEvent {}), + ProcedureState::Done { output: None }, + EventTrigger::Succeeded, + context, + ); + + assert_eq!( + submitted.extra_rows().unwrap()[0].values[5], + jsonb_value(&json!({"reason": "auto_repartition"})) + ); + assert_eq!( + succeeded.extra_rows().unwrap()[0].values[5], + Value { value_data: None } + ); + } + #[test] fn test_event_trigger_serialization() { let procedure_id = ProcedureId::parse_str("00000000-0000-0000-0000-000000000001").unwrap(); diff --git a/src/common/procedure/src/lib.rs b/src/common/procedure/src/lib.rs index 1d7a200ca58..a21e5d1c47a 100644 --- a/src/common/procedure/src/lib.rs +++ b/src/common/procedure/src/lib.rs @@ -30,8 +30,8 @@ pub use crate::error::{Error, Result}; pub use crate::event::ProcedureEvent; pub use crate::procedure::{ BoxedProcedure, BoxedProcedureLoader, ChildSubmissionOutcome, Context, ContextProvider, - ContextProviderRef, EventRuntimeContext, EventTrigger, LockKey, Output, ParseIdError, - PoisonKey, PoisonKeys, Procedure, ProcedureId, ProcedureInfo, ProcedureManager, + ContextProviderRef, EventContext, EventTrigger, LockKey, Output, ParseIdError, PoisonKey, + PoisonKeys, Procedure, ProcedureContext, ProcedureId, ProcedureInfo, ProcedureManager, ProcedureManagerRef, ProcedureState, ProcedureWithId, RetryPhase, Status, StringKey, }; pub use crate::watcher::Watcher; diff --git a/src/common/procedure/src/local.rs b/src/common/procedure/src/local.rs index 59eaee2d789..8fdd3167583 100644 --- a/src/common/procedure/src/local.rs +++ b/src/common/procedure/src/local.rs @@ -43,8 +43,8 @@ use crate::rwlock::{KeyRwLock, OwnedKeyRwLockGuard}; use crate::store::poison_store::PoisonStoreRef; use crate::store::{ProcedureMessage, ProcedureMessages, ProcedureStore, StateStoreRef}; use crate::{ - BoxedProcedure, ContextProvider, EventTrigger, LockKey, PoisonKey, ProcedureId, - ProcedureManager, ProcedureState, ProcedureWithId, StringKey, Watcher, + BoxedProcedure, ContextProvider, EventTrigger, LockKey, PoisonKey, ProcedureContext, + ProcedureId, ProcedureManager, ProcedureState, ProcedureWithId, StringKey, Watcher, }; /// The expired time of a procedure's metadata. @@ -81,6 +81,8 @@ pub(crate) struct ProcedureMeta { type_name: String, /// Parent procedure id. parent_id: Option, + /// Context associated with the root submission. + context: ProcedureContext, /// Notify to wait for subprocedures. child_notify: Notify, /// Lock required by this procedure. @@ -105,6 +107,7 @@ impl ProcedureMeta { id: ProcedureId, procedure_state: ProcedureState, parent_id: Option, + context: ProcedureContext, lock_key: LockKey, poison_keys: PoisonKeys, type_name: &str, @@ -113,6 +116,7 @@ impl ProcedureMeta { ProcedureMeta { id, parent_id, + context, child_notify: Notify::new(), lock_key, poison_keys, @@ -716,6 +720,7 @@ impl LocalManager { procedure_state: ProcedureState, step: u32, procedure: BoxedProcedure, + context: ProcedureContext, origin: RootSubmissionOrigin, ) -> Result { ensure!(self.manager_ctx.running(), ManagerNotStartSnafu); @@ -724,6 +729,7 @@ impl LocalManager { procedure_id, procedure_state, None, + context, procedure.lock_key(), procedure.poison_keys(), procedure.type_name(), @@ -837,6 +843,7 @@ impl LocalManager { procedure_state, loaded_procedure.step, loaded_procedure.procedure, + message.context.clone(), RootSubmissionOrigin::Recovery, ) { error!(e; "Failed to recover procedure {}", procedure_id); @@ -964,6 +971,7 @@ impl ProcedureManager for LocalManager { ProcedureState::Running, 0, procedure.procedure, + procedure.context, RootSubmissionOrigin::Fresh, ) } @@ -1021,6 +1029,7 @@ pub(crate) mod test_util { ProcedureId::random(), ProcedureState::Running, None, + ProcedureContext::default(), LockKey::default(), PoisonKeys::default(), "ProcedureAdapter", @@ -1043,7 +1052,10 @@ mod tests { use common_error::mock::MockError; use common_error::status_code::StatusCode; - use common_event_recorder::{Event, EventRecorder, EventTypeFilter, EventTypeFilterRef}; + use common_event_recorder::{ + Event, EventRecorder, EventTypeFilter, EventTypeFilterRef, PersistentEventContext, + TriggerReason, + }; use common_test_util::temp_dir::create_temp_dir; use tokio::sync::oneshot; use tokio::time::{sleep, timeout}; @@ -1053,8 +1065,8 @@ mod tests { use crate::store::state_store::ObjectStateStore; use crate::test_util::InMemoryPoisonStore; use crate::{ - ChildSubmissionOutcome, Context, EventRuntimeContext, EventTrigger, Procedure, - ProcedureEvent, Status, + ChildSubmissionOutcome, Context, EventContext, EventTrigger, Procedure, ProcedureEvent, + Status, }; fn new_test_manager_context() -> ManagerContext { @@ -1103,6 +1115,13 @@ mod tests { }) .collect() } + + fn procedure_context(&self, procedure_id: ProcedureId) -> Option { + self.events.lock().unwrap().iter().find_map(|event| { + let event = event.as_any().downcast_ref::().unwrap(); + (event.procedure_id == procedure_id).then(|| event.context.clone()) + }) + } } impl EventRecorder for CapturingEventRecorder { @@ -1296,7 +1315,7 @@ mod tests { self.poison_keys.clone() } - fn event(&self, _ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, _ctx: &EventContext<'_>) -> Option> { Some(Box::new(TestProcedureEvent)) } } @@ -1341,6 +1360,7 @@ mod tests { vec![ProcedureWithId { id: self.child_id, procedure: Box::new(ProcedureToLoad::new("child")), + context: ProcedureContext::default(), }], false, )) @@ -1354,7 +1374,7 @@ mod tests { LockKey::default() } - fn event(&self, _ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, _ctx: &EventContext<'_>) -> Option> { Some(Box::new(TestProcedureEvent)) } } @@ -1384,6 +1404,7 @@ mod tests { vec![ProcedureWithId { id: self.child_id, procedure: Box::new(ProcedureToLoad::new("child")), + context: ProcedureContext::default(), }], false, )) @@ -1397,7 +1418,7 @@ mod tests { LockKey::default() } - fn event(&self, _ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, _ctx: &EventContext<'_>) -> Option> { Some(Box::new(TestProcedureEvent)) } } @@ -1424,7 +1445,7 @@ mod tests { LockKey::default() } - fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { *self.captured_filter.lock().unwrap() = Some(ctx.event_type_filter.clone()); Some(Box::new(TestProcedureEvent)) } @@ -1458,6 +1479,7 @@ mod tests { procedure: Box::new(FilterCapturingProcedure { captured_filter: captured_filter.clone(), }), + context: ProcedureContext::default(), }) .await .unwrap(); @@ -1499,6 +1521,7 @@ mod tests { .submit(ProcedureWithId { id: ProcedureId::random(), procedure: Box::new(ProcedureToLoad::new("fresh submission")), + context: ProcedureContext::default(), }) .await .unwrap(); @@ -1533,6 +1556,7 @@ mod tests { ProcedureState::Running, 0, Box::new(ProcedureToLoad::new("submit root without runner")), + ProcedureContext::default(), RootSubmissionOrigin::Fresh, ) }); @@ -1584,6 +1608,9 @@ mod tests { child_id, child_submitted: false, }), + context: ProcedureContext::from_event_context(PersistentEventContext::new( + TriggerReason::AutoRepartition, + )), }) .await .unwrap(); @@ -1604,6 +1631,15 @@ mod tests { }, ))); assert!(procedure_events.contains(&(child_id, EventTrigger::Submitted))); + assert_eq!( + event_recorder + .procedure_context(child_id) + .unwrap() + .event_context + .unwrap() + .reason, + TriggerReason::AutoRepartition + ); } #[tokio::test] @@ -1637,6 +1673,7 @@ mod tests { proceed: Some(proceed_rx), child_submitted: false, }), + context: ProcedureContext::default(), }) .await .unwrap(); @@ -1711,12 +1748,15 @@ mod tests { let procedure = ProcedureToLoad::new("recovered submission"); let procedure_id = ProcedureId::random(); ProcedureStore::from_object_store(object_store) - .store_procedure( + .store_procedure_with_context( procedure_id, 0, procedure.type_name().to_string(), procedure.dump().unwrap(), None, + ProcedureContext::from_event_context(PersistentEventContext::new( + TriggerReason::AutoRebalance, + )), ) .await .unwrap(); @@ -1732,6 +1772,15 @@ mod tests { ); wait_for_trigger(&event_recorder, EventTrigger::Recovered).await; assert!(!event_recorder.triggers().contains(&EventTrigger::Submitted)); + assert_eq!( + event_recorder + .procedure_context(procedure_id) + .unwrap() + .event_context + .unwrap() + .reason, + TriggerReason::AutoRebalance + ); } #[tokio::test] @@ -1813,6 +1862,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .unwrap(); @@ -1965,6 +2015,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .is_ok() @@ -1986,6 +2037,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(ProcedureToLoad::new("submit")), + context: ProcedureContext::default(), }) .await .unwrap_err(); @@ -2054,6 +2106,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .unwrap() @@ -2094,6 +2147,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .unwrap_err(), @@ -2126,6 +2180,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .is_ok() @@ -2164,6 +2219,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .is_ok() @@ -2193,6 +2249,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .is_ok() @@ -2217,6 +2274,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .is_ok() @@ -2265,6 +2323,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .unwrap_err(); @@ -2285,6 +2344,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: ProcedureContext::default(), }) .await .is_ok() diff --git a/src/common/procedure/src/local/runner.rs b/src/common/procedure/src/local/runner.rs index a460b7376bf..30b775dbd7d 100644 --- a/src/common/procedure/src/local/runner.rs +++ b/src/common/procedure/src/local/runner.rs @@ -34,8 +34,8 @@ use crate::procedure::{Output, StringKey}; use crate::rwlock::OwnedKeyRwLockGuard; use crate::store::{ProcedureMessage, ProcedureStore}; use crate::{ - BoxedProcedure, ChildSubmissionOutcome, Context, Error, EventRuntimeContext, EventTrigger, - Procedure, ProcedureId, ProcedureState, ProcedureWithId, RetryPhase, Status, + BoxedProcedure, ChildSubmissionOutcome, Context, Error, EventContext, EventTrigger, Procedure, + ProcedureId, ProcedureState, ProcedureWithId, RetryPhase, Status, }; /// A guard to cleanup procedure state. @@ -227,6 +227,7 @@ impl Runner { let ctx = Context { procedure_id: self.meta.id, provider: self.manager_ctx.clone(), + event_context: self.meta.context.event_context.clone(), }; self.rolling_back = false; @@ -508,6 +509,7 @@ impl Runner { procedure_id, procedure_state, Some(self.meta.id), + self.meta.context.clone(), procedure.lock_key(), procedure.poison_keys(), procedure.type_name(), @@ -662,12 +664,13 @@ impl Runner { let data = self.procedure.dump()?; self.store - .store_procedure( + .store_procedure_with_context( self.meta.id, self.step, type_name, data, self.meta.parent_id, + self.meta.context.clone(), ) .await .map_err(|e| { @@ -708,6 +711,7 @@ impl Runner { parent_id: self.meta.parent_id, step: self.step, error: Some(error), + context: self.meta.context.clone(), }; self.store .rollback_procedure(self.meta.id, message) @@ -776,15 +780,22 @@ impl Runner { pub(crate) fn build_event(&self, trigger: EventTrigger) -> Option { let recorder = self.event_recorder.as_ref()?; let state = self.meta.state(); - let context = EventRuntimeContext { + let context = EventContext { procedure_id: self.meta.id, lifecycle_state: &state, trigger: trigger.clone(), event_type_filter: recorder.event_type_filter(), + event_context: self.meta.context.event_context.as_ref(), }; - self.procedure - .event(&context) - .map(|event| ProcedureEvent::new(self.meta.id, event, state, trigger)) + self.procedure.event(&context).map(|event| { + ProcedureEvent::new_with_context( + self.meta.id, + event, + state, + trigger, + self.meta.context.clone(), + ) + }) } /// Builds and dispatches an event from the live procedure. Delivery is best effort and is @@ -876,6 +887,7 @@ mod tests { Context { procedure_id, provider, + event_context: None, } } @@ -914,6 +926,7 @@ mod tests { Context { procedure_id, provider: Arc::new(MockProvider), + event_context: None, } } @@ -1100,6 +1113,7 @@ mod tests { ProcedureWithId { id: procedure_id, procedure: Box::new(child), + context: Default::default(), } } @@ -1480,6 +1494,7 @@ mod tests { subprocedures: vec![ProcedureWithId { id: child_id, procedure: Box::new(fail), + context: Default::default(), }], persist: true, }) @@ -1672,6 +1687,7 @@ mod tests { subprocedures: vec![ProcedureWithId { id: child_id, procedure: Box::new(fail), + context: Default::default(), }], persist: true, }) @@ -2192,6 +2208,7 @@ mod tests { subprocedures: vec![ProcedureWithId { id: child_id, procedure: Box::new(child), + context: Default::default(), }], persist: false, }) @@ -2248,6 +2265,7 @@ mod tests { subprocedures: vec![ProcedureWithId { id: child_id, procedure: Box::new(child), + context: Default::default(), }], persist: false, }) diff --git a/src/common/procedure/src/procedure.rs b/src/common/procedure/src/procedure.rs index a8af03e0da0..cd11aea9215 100644 --- a/src/common/procedure/src/procedure.rs +++ b/src/common/procedure/src/procedure.rs @@ -19,7 +19,7 @@ use std::str::FromStr; use std::sync::Arc; use async_trait::async_trait; -use common_event_recorder::{Event, EventTypeFilterRef}; +use common_event_recorder::{Event, EventTypeFilterRef, PersistentEventContext}; use serde::{Deserialize, Serialize}; use smallvec::{SmallVec, smallvec}; use snafu::{ResultExt, Snafu}; @@ -32,6 +32,26 @@ use crate::watcher::Watcher; pub type Output = Arc; +/// Context attached to a procedure submission and inherited by its children. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct ProcedureContext { + /// Context describing why and how the procedure was submitted. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub event_context: Option, +} + +impl ProcedureContext { + pub fn from_event_context(event_context: PersistentEventContext) -> Self { + Self { + event_context: Some(event_context), + } + } + + pub fn is_empty(&self) -> bool { + self.event_context.is_none() + } +} + /// Procedure execution status. #[derive(Debug)] pub enum Status { @@ -175,6 +195,8 @@ pub struct Context { pub procedure_id: ProcedureId, /// [ProcedureManager] context provider. pub provider: ContextProviderRef, + /// Event context inherited from the root submission. + pub event_context: Option, } impl Context { @@ -237,13 +259,13 @@ pub trait Procedure: Send { /// [`Event::extra_schema`] values. The event recorder batches events by type /// and rejects incompatible schemas; use a distinct event type for a /// different schema. - fn event(&self, _ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, _ctx: &EventContext<'_>) -> Option> { None } } /// Framework-owned context supplied when a procedure builds a lifecycle event. -pub struct EventRuntimeContext<'a> { +pub struct EventContext<'a> { /// Id of the procedure associated with the event. pub procedure_id: ProcedureId, /// Current framework state of the procedure. @@ -252,6 +274,8 @@ pub struct EventRuntimeContext<'a> { pub trigger: EventTrigger, /// Event types retained by the configured recorder. pub event_type_filter: EventTypeFilterRef, + /// Event context inherited from the root submission. + pub event_context: Option<&'a PersistentEventContext>, } /// Lifecycle action that causes the framework to invoke [`Procedure::event`]. @@ -336,7 +360,7 @@ impl Procedure for Box { (**self).poison_keys() } - fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { (**self).event(ctx) } } @@ -459,6 +483,8 @@ pub struct ProcedureWithId { /// Id of the procedure. pub id: ProcedureId, pub procedure: BoxedProcedure, + /// Context associated with this procedure submission. + pub context: ProcedureContext, } impl ProcedureWithId { @@ -468,8 +494,15 @@ impl ProcedureWithId { ProcedureWithId { id: ProcedureId::random(), procedure, + context: ProcedureContext::default(), } } + + /// Attaches event context to this procedure submission. + pub fn with_event_context(mut self, event_context: PersistentEventContext) -> Self { + self.context.event_context = Some(event_context); + self + } } impl fmt::Debug for ProcedureWithId { @@ -721,11 +754,12 @@ mod tests { #[test] fn test_default_procedure_event_hook() { let state = ProcedureState::Running; - let context = EventRuntimeContext { + let context = EventContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Succeeded, event_type_filter: Arc::new(common_event_recorder::EventTypeFilter::All), + event_context: None, }; assert!(DefaultEventProcedure.event(&context).is_none()); diff --git a/src/common/procedure/src/store.rs b/src/common/procedure/src/store.rs index 906682ed6e6..4867046f8de 100644 --- a/src/common/procedure/src/store.rs +++ b/src/common/procedure/src/store.rs @@ -20,9 +20,9 @@ use futures::TryStreamExt; use serde::{Deserialize, Serialize}; use snafu::ResultExt; -use crate::ProcedureId; use crate::error::{Result, ToJsonSnafu}; pub(crate) use crate::store::state_store::StateStoreRef; +use crate::{ProcedureContext, ProcedureId}; pub mod poison_store; pub mod state_store; @@ -55,6 +55,9 @@ pub struct ProcedureMessage { /// Errors raised during the procedure. #[serde(default, skip_serializing_if = "Option::is_none")] pub error: Option, + /// Context associated with the root submission. + #[serde(default, skip_serializing_if = "ProcedureContext::is_empty")] + pub context: ProcedureContext, } /// A collection of all procedures' messages. @@ -88,6 +91,7 @@ impl ProcedureStore { } /// Dump the `procedure` to the storage. + #[cfg(test)] pub(crate) async fn store_procedure( &self, procedure_id: ProcedureId, @@ -95,6 +99,26 @@ impl ProcedureStore { type_name: String, data: String, parent_id: Option, + ) -> Result<()> { + self.store_procedure_with_context( + procedure_id, + step, + type_name, + data, + parent_id, + ProcedureContext::default(), + ) + .await + } + + pub(crate) async fn store_procedure_with_context( + &self, + procedure_id: ProcedureId, + step: u32, + type_name: String, + data: String, + parent_id: Option, + context: ProcedureContext, ) -> Result<()> { let message = ProcedureMessage { type_name, @@ -102,6 +126,7 @@ impl ProcedureStore { parent_id, step, error: None, + context, }; let key = ParsedKey { prefix: &self.proc_path, @@ -356,6 +381,7 @@ mod tests { } use async_trait::async_trait; + use common_event_recorder::{PersistentEventContext, TriggerReason}; use common_test_util::temp_dir::{TempDir, create_temp_dir}; use object_store::services::Fs as Builder; @@ -462,6 +488,7 @@ mod tests { parent_id: None, step: 4, error: None, + context: ProcedureContext::default(), }; let json = serde_json::to_string(&message).unwrap(); @@ -477,6 +504,21 @@ mod tests { json, r#"{"type_name":"TestMessage","data":"no parent id","parent_id":"9f805a1f-05f7-490c-9f91-bd56e3cc54c1","step":4}"# ); + + message.context = ProcedureContext::from_event_context(PersistentEventContext::new( + TriggerReason::AutoRebalance, + )); + let json = serde_json::to_string(&message).unwrap(); + assert_eq!( + json, + r#"{"type_name":"TestMessage","data":"no parent id","parent_id":"9f805a1f-05f7-490c-9f91-bd56e3cc54c1","step":4,"context":{"event_context":{"reason":"auto_rebalance"}}}"# + ); + + let legacy: ProcedureMessage = serde_json::from_str( + r#"{"type_name":"TestMessage","data":"legacy","parent_id":null,"step":1}"#, + ) + .unwrap(); + assert_eq!(legacy.context, ProcedureContext::default()); } struct MockProcedure { @@ -541,6 +583,7 @@ mod tests { parent_id: None, step: 0, error: None, + context: ProcedureContext::default(), }; assert_eq!(expect, *msg); } @@ -589,6 +632,7 @@ mod tests { parent_id: None, step: 1, error: None, + context: ProcedureContext::default(), }; store .rollback_procedure(procedure_id, message) diff --git a/src/common/procedure/src/watcher.rs b/src/common/procedure/src/watcher.rs index 4ecd75ded0f..4717fcd0010 100644 --- a/src/common/procedure/src/watcher.rs +++ b/src/common/procedure/src/watcher.rs @@ -124,6 +124,7 @@ mod tests { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(MockProcedure { error: true }), + context: Default::default(), }) .await .unwrap(); diff --git a/src/meta-srv/src/event/gc.rs b/src/meta-srv/src/event/gc.rs index c467bc665c6..3575ac20c2f 100644 --- a/src/meta-srv/src/event/gc.rs +++ b/src/meta-srv/src/event/gc.rs @@ -21,10 +21,9 @@ use api::v1::{ColumnSchema, Row, Value}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - EVENT_CONTEXT_COLUMN, GC_REPORT_COLUMN, REGION_ID_COLUMN, REGION_NUMBER_COLUMN, - TABLE_ID_COLUMN, column_schemas, nullable_json, + GC_REPORT_COLUMN, REGION_ID_COLUMN, REGION_NUMBER_COLUMN, TABLE_ID_COLUMN, column_schemas, + nullable_json, }; -use common_meta::rpc::ddl::EventContext; use serde::Serialize; use snafu::ResultExt; use store_api::storage::{GcReport, IndexVersion, RegionId}; @@ -67,7 +66,6 @@ struct BatchGcRegionRow { #[derive(Debug)] pub(crate) struct BatchGcEvent { payload: Option, - event_context: Option, // None emits a procedure-level lifecycle row with null Region dimensions. regions: Option>, } @@ -77,7 +75,6 @@ impl BatchGcEvent { regions: &[RegionId], full_file_listing: bool, timeout: Duration, - event_context: Option, ) -> Self { Self { payload: Some(BatchGcPayload { @@ -86,7 +83,6 @@ impl BatchGcEvent { full_file_listing, timeout, }), - event_context, regions: None, } } @@ -111,7 +107,6 @@ impl BatchGcEvent { (!regions.is_empty()).then_some(Self { payload: None, - event_context: None, regions: Some(regions), }) } @@ -136,22 +131,8 @@ impl Event for BatchGcEvent { } fn extra_rows(&self) -> Result> { - let event_context = self - .event_context - .as_ref() - .map(serde_json::to_value) - .transpose() - .context(SerializeEventSnafu)?; let Some(regions) = &self.regions else { - return Ok(vec![Row { - values: vec![ - Value { value_data: None }, - Value { value_data: None }, - Value { value_data: None }, - Value { value_data: None }, - nullable_json(event_context.as_ref()), - ], - }]); + return Ok(vec![null_row()]); }; regions @@ -164,7 +145,6 @@ impl Event for BatchGcEvent { ValueData::U32Value(region.region_id.table_id()).into(), ValueData::U32Value(region.region_id.region_number()).into(), nullable_json(Some(&report)), - nullable_json(None), ], }) }) @@ -182,11 +162,9 @@ fn schema() -> Vec { &TABLE_ID_COLUMN, ®ION_NUMBER_COLUMN, &GC_REPORT_COLUMN, - &EVENT_CONTEXT_COLUMN, ]) } -#[cfg(test)] fn null_row() -> Row { Row { values: (0..schema().len()) @@ -233,17 +211,17 @@ mod tests { use std::sync::Arc; use api::v1::ColumnSchema; - use common_event_recorder::EventTypeFilter; use common_event_recorder::event_table::{ - PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, + EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, jsonb_value, }; use common_event_recorder::testing::assert_event_contract; + use common_event_recorder::{EventTypeFilter, PersistentEventContext, TriggerReason}; use common_meta::key::TableMetadataManager; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::sequence::SequenceBuilder; use common_procedure::{ - EventRuntimeContext, EventTrigger, Procedure, ProcedureEvent, ProcedureId, ProcedureState, + EventContext, EventTrigger, Procedure, ProcedureEvent, ProcedureId, ProcedureState, RetryPhase, }; use store_api::storage::FileId; @@ -256,12 +234,8 @@ mod tests { fn test_batch_gc_lifecycle_event_contract() { let first = RegionId::new(1024, 1); let second = RegionId::new(1024, 2); - let event = BatchGcEvent::with_config( - &[second, first, second], - true, - Duration::from_secs(10), - None, - ); + let event = + BatchGcEvent::with_config(&[second, first, second], true, Duration::from_secs(10)); assert_event_contract(&event, BATCH_GC_EVENT_TYPE, &schema(), &[null_row()]); assert_eq!( @@ -356,22 +330,35 @@ mod tests { fn test_batch_gc_event_filter() { let procedure = batch_gc_procedure(); let running = ProcedureState::Running; - let runtime_context = |trigger, lifecycle_state, event_type_filter| EventRuntimeContext { + let manual_context = PersistentEventContext::new(TriggerReason::Manual); + let event_context = |trigger, lifecycle_state, event_type_filter| EventContext { procedure_id: ProcedureId::random(), lifecycle_state, trigger, event_type_filter: Arc::new(event_type_filter), + event_context: None, }; assert!( procedure - .event(&runtime_context( + .event(&event_context( EventTrigger::Submitted, &running, EventTypeFilter::All, )) .is_none() ); + assert!( + procedure + .event(&EventContext { + procedure_id: ProcedureId::random(), + lifecycle_state: &running, + trigger: EventTrigger::Submitted, + event_type_filter: Arc::new(EventTypeFilter::All), + event_context: Some(&manual_context), + }) + .is_some() + ); let report = GcReport { processed_regions: HashSet::from([RegionId::new(1024, 1)]), @@ -382,7 +369,7 @@ mod tests { }; assert!( procedure - .event(&runtime_context( + .event(&event_context( EventTrigger::Succeeded, &done, EventTypeFilter::All, @@ -392,7 +379,7 @@ mod tests { assert!( procedure - .event(&runtime_context( + .event(&event_context( EventTrigger::Recovered, &running, EventTypeFilter::All, @@ -401,7 +388,7 @@ mod tests { ); let retrying = procedure - .event(&runtime_context( + .event(&event_context( EventTrigger::Retrying { phase: RetryPhase::Execute, attempt: 1, @@ -423,7 +410,7 @@ mod tests { assert!( procedure - .event(&runtime_context( + .event(&event_context( EventTrigger::Recovered, &running, EventTypeFilter::Only(HashSet::from(["another_event".to_string()])), @@ -432,7 +419,7 @@ mod tests { ); assert!( procedure - .event(&runtime_context( + .event(&event_context( EventTrigger::Recovered, &running, EventTypeFilter::Only(HashSet::new()), @@ -443,7 +430,7 @@ mod tests { let missing = ProcedureState::Done { output: None }; assert!( procedure - .event(&runtime_context( + .event(&event_context( EventTrigger::Succeeded, &missing, EventTypeFilter::All, @@ -455,7 +442,7 @@ mod tests { }; assert!( procedure - .event(&runtime_context( + .event(&event_context( EventTrigger::Succeeded, &wrong, EventTypeFilter::All, @@ -468,11 +455,12 @@ mod tests { fn test_batch_gc_terminal_event_with_report() { let mut procedure = batch_gc_procedure(); let running = ProcedureState::Running; - let runtime_context = |trigger| EventRuntimeContext { + let event_context = |trigger| EventContext { procedure_id: ProcedureId::random(), lifecycle_state: &running, trigger, event_type_filter: Arc::new(EventTypeFilter::All), + event_context: None, }; let report = GcReport { need_retry_regions: HashSet::from([RegionId::new(1024, 2)]), @@ -481,7 +469,7 @@ mod tests { procedure.set_gc_report_for_test(report); let retrying = procedure - .event(&runtime_context(EventTrigger::Retrying { + .event(&event_context(EventTrigger::Retrying { phase: RetryPhase::Execute, attempt: 2, })) @@ -498,7 +486,7 @@ mod tests { assert_eq!(retrying.extra_rows().unwrap(), vec![null_row()]); for trigger in [EventTrigger::Failed, EventTrigger::Poisoned] { - let event = procedure.event(&runtime_context(trigger)).unwrap(); + let event = procedure.event(&event_context(trigger)).unwrap(); assert_eq!(event.json_payload().unwrap(), serde_json::Value::Null); assert_eq!( event.extra_rows().unwrap(), @@ -513,7 +501,7 @@ mod tests { let procedure = batch_gc_procedure(); let failed_without_report = procedure - .event(&runtime_context(EventTrigger::Failed)) + .event(&event_context(EventTrigger::Failed)) .unwrap(); assert_eq!( failed_without_report.json_payload().unwrap(), @@ -548,6 +536,7 @@ mod tests { ); let mut event_schema = procedure_schema(); event_schema.extend(schema()); + event_schema.push(EVENT_CONTEXT_COLUMN.column_schema()); let mut values = vec![ ValueData::StringValue(procedure_id.to_string()).into(), ValueData::StringValue("Done".to_string()).into(), @@ -564,6 +553,7 @@ mod tests { ) .values, ); + values.push(Value { value_data: None }); assert_event_contract( &event, @@ -589,7 +579,6 @@ mod tests { ValueData::U32Value(region_id.table_id()).into(), ValueData::U32Value(region_id.region_number()).into(), nullable_json(report.as_ref()), - nullable_json(None), ], } } @@ -607,7 +596,6 @@ mod tests { true, Duration::from_secs(10), HashMap::new(), - EventContext::default(), ) } } diff --git a/src/meta-srv/src/event/region_migration.rs b/src/meta-srv/src/event/region_migration.rs index 65b89df2515..cfbe114077e 100644 --- a/src/meta-srv/src/event/region_migration.rs +++ b/src/meta-srv/src/event/region_migration.rs @@ -17,14 +17,13 @@ use std::time::Duration; use api::v1::value::ValueData; use api::v1::{ColumnSchema, Row}; -use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - EVENT_CONTEXT_COLUMN, REGION_ID_COLUMN, REGION_MIGRATION_DST_NODE_ID_COLUMN, - REGION_MIGRATION_DST_PEER_ADDR_COLUMN, REGION_MIGRATION_SRC_NODE_ID_COLUMN, - REGION_MIGRATION_SRC_PEER_ADDR_COLUMN, REGION_MIGRATION_TRIGGER_REASON_COLUMN, - REGION_NUMBER_COLUMN, TABLE_ID_COLUMN, column_schemas, nullable_json, + REGION_ID_COLUMN, REGION_MIGRATION_DST_NODE_ID_COLUMN, REGION_MIGRATION_DST_PEER_ADDR_COLUMN, + REGION_MIGRATION_SRC_NODE_ID_COLUMN, REGION_MIGRATION_SRC_PEER_ADDR_COLUMN, + REGION_MIGRATION_TRIGGER_REASON_COLUMN, REGION_NUMBER_COLUMN, TABLE_ID_COLUMN, column_schemas, }; +use common_event_recorder::{Event, TriggerReason}; use serde::Serialize; use snafu::ResultExt; use store_api::storage::RegionId; @@ -50,7 +49,6 @@ pub(crate) struct RegionMigrationEvent { dst_peer_addr: String, // The timeout of the region migration. timeout: Duration, - event_context: Option, } #[derive(Debug, Serialize)] @@ -60,16 +58,15 @@ struct Payload { } impl RegionMigrationEvent { - pub fn from_persistent_ctx(ctx: &PersistentContext, include_event_context: bool) -> Self { + pub fn from_persistent_ctx(ctx: &PersistentContext, trigger_reason: TriggerReason) -> Self { Self { region_ids: ctx.region_ids.clone(), - trigger_reason: ctx.trigger_reason, + trigger_reason: RegionMigrationTriggerReason::from_trigger_reason(trigger_reason), src_node_id: ctx.from_peer.id, src_peer_addr: ctx.from_peer.addr.clone(), dst_node_id: ctx.to_peer.id, dst_peer_addr: ctx.to_peer.addr.clone(), timeout: ctx.timeout, - event_context: include_event_context.then(|| ctx.event_context.clone()), } } } @@ -89,7 +86,6 @@ impl Event for RegionMigrationEvent { ®ION_MIGRATION_SRC_PEER_ADDR_COLUMN, ®ION_MIGRATION_DST_NODE_ID_COLUMN, ®ION_MIGRATION_DST_PEER_ADDR_COLUMN, - &EVENT_CONTEXT_COLUMN, ]) } @@ -106,12 +102,6 @@ impl Event for RegionMigrationEvent { ValueData::StringValue(self.src_peer_addr.clone()).into(), ValueData::U64Value(self.dst_node_id).into(), ValueData::StringValue(self.dst_peer_addr.clone()).into(), - self.event_context - .as_ref() - .map(serde_json::to_value) - .transpose() - .context(SerializeEventSnafu) - .map(|value| nullable_json(value.as_ref()))?, ], }); } diff --git a/src/meta-srv/src/event/repartition.rs b/src/meta-srv/src/event/repartition.rs index 958bc4e11d8..1faf3f18212 100644 --- a/src/meta-srv/src/event/repartition.rs +++ b/src/meta-srv/src/event/repartition.rs @@ -21,11 +21,11 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, PARENT_PROCEDURE_ID_COLUMN, - REPARTITION_GROUP_ID_COLUMN, SCHEMA_NAME_COLUMN, SOURCE_PARTITION_EXPR_COLUMN, - SOURCE_REGION_ID_COLUMN, SOURCE_REGION_NUMBER_COLUMN, TABLE_ID_COLUMN, TABLE_NAME_COLUMN, - TARGET_PARTITION_EXPR_COLUMN, TARGET_REGION_ID_COLUMN, TARGET_REGION_NUMBER_COLUMN, - column_schemas, nullable_json, nullable_string, nullable_value, + CATALOG_NAME_COLUMN, PARENT_PROCEDURE_ID_COLUMN, REPARTITION_GROUP_ID_COLUMN, + SCHEMA_NAME_COLUMN, SOURCE_PARTITION_EXPR_COLUMN, SOURCE_REGION_ID_COLUMN, + SOURCE_REGION_NUMBER_COLUMN, TABLE_ID_COLUMN, TABLE_NAME_COLUMN, TARGET_PARTITION_EXPR_COLUMN, + TARGET_REGION_ID_COLUMN, TARGET_REGION_NUMBER_COLUMN, column_schemas, nullable_string, + nullable_value, }; use serde::Serialize; use snafu::ResultExt; @@ -61,7 +61,6 @@ pub(crate) struct RepartitionEvent { table_name: Option, table_id: Option, payload: Option, - event_context: Option, } impl RepartitionEvent { @@ -91,7 +90,6 @@ impl RepartitionEvent { target_partition_columns: intent.target_partition_columns().map(ToOwned::to_owned), timeout: persistent_ctx.timeout, }), - event_context: Some(persistent_ctx.event_context.clone()), } } @@ -102,17 +100,10 @@ impl RepartitionEvent { table_name: Some(persistent_ctx.table_name.clone()), table_id: Some(persistent_ctx.table_id), payload: None, - event_context: None, } } fn schema() -> Vec { - let mut schema = Self::base_schema(); - schema.push(EVENT_CONTEXT_COLUMN.column_schema()); - schema - } - - fn base_schema() -> Vec { column_schemas([ &CATALOG_NAME_COLUMN, &SCHEMA_NAME_COLUMN, @@ -141,19 +132,12 @@ impl Event for RepartitionEvent { } fn extra_rows(&self) -> Result> { - let event_context = self - .event_context - .as_ref() - .map(serde_json::to_value) - .transpose() - .context(SerializeEventSnafu)?; Ok(vec![Row { values: vec![ nullable_string(self.catalog_name.as_deref()), nullable_string(self.schema_name.as_deref()), nullable_string(self.table_name.as_deref()), nullable_value(self.table_id.map(ValueData::U32Value)), - nullable_json(event_context.as_ref()), ], }]) } @@ -323,7 +307,7 @@ impl RepartitionGroupEvent { } fn schema() -> Vec { - let mut schema = RepartitionEvent::base_schema(); + let mut schema = RepartitionEvent::schema(); schema.extend(column_schemas([ &PARENT_PROCEDURE_ID_COLUMN, &REPARTITION_GROUP_ID_COLUMN, @@ -377,14 +361,13 @@ mod tests { use std::time::Duration; use api::v1::value::ValueData; - use api::v1::{ColumnSchema, Row}; + use api::v1::{ColumnSchema, Row, Value}; use common_event_recorder::Event; use common_event_recorder::event_table::{ - PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, + EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, jsonb_value, }; use common_event_recorder::testing::assert_event_contract; - use common_meta::rpc::ddl::EventContext; use common_procedure::{EventTrigger, ProcedureEvent, ProcedureId, ProcedureState}; use table::table_name::TableName; use uuid::Uuid; @@ -402,7 +385,6 @@ mod tests { TableName::new("greptime", "public", "repartition_events"), 1024, Some(Duration::from_secs(30)), - EventContext::default(), ) } @@ -438,9 +420,6 @@ mod tests { ValueData::StringValue("public".to_string()).into(), ValueData::StringValue("repartition_events".to_string()).into(), ValueData::U32Value(1024).into(), - jsonb_value(&serde_json::json!({ - "reason": "unknown", - })), ], }], ); @@ -654,7 +633,6 @@ mod tests { ValueData::StringValue("public".to_string()).into(), ValueData::StringValue("repartition_events".to_string()).into(), ValueData::U32Value(1024).into(), - Default::default(), ], }], ); @@ -690,6 +668,7 @@ mod tests { ); let mut schema = procedure_schema(); schema.extend(parent_schema()); + schema.push(EVENT_CONTEXT_COLUMN.column_schema()); assert_event_contract( &event, @@ -705,9 +684,7 @@ mod tests { ValueData::StringValue("public".to_string()).into(), ValueData::StringValue("repartition_events".to_string()).into(), ValueData::U32Value(1024).into(), - jsonb_value(&serde_json::json!({ - "reason": "unknown", - })), + Value { value_data: None }, ], }], ); @@ -719,12 +696,11 @@ mod tests { &SCHEMA_NAME_COLUMN, &TABLE_NAME_COLUMN, &TABLE_ID_COLUMN, - &EVENT_CONTEXT_COLUMN, ]) } fn group_schema() -> Vec { - let mut schema = RepartitionEvent::base_schema(); + let mut schema = parent_schema(); schema.extend(column_schemas([ &PARENT_PROCEDURE_ID_COLUMN, &REPARTITION_GROUP_ID_COLUMN, diff --git a/src/meta-srv/src/gc/ctx.rs b/src/meta-srv/src/gc/ctx.rs index 41027a65390..c9cc8049867 100644 --- a/src/meta-srv/src/gc/ctx.rs +++ b/src/meta-srv/src/gc/ctx.rs @@ -29,7 +29,7 @@ use common_meta::key::DroppedTableName; use common_meta::key::TableMetadataManagerRef; use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; -use common_meta::rpc::ddl::EventContext; +use common_meta::rpc::ddl::PersistentEventContext; #[cfg(feature = "enterprise")] use common_meta::rpc::ddl::PurgeDroppedTableTask; use common_procedure::{ProcedureManagerRef, ProcedureWithId, watcher}; @@ -68,7 +68,7 @@ pub(crate) trait SchedulerCtx: Send + Sync { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result; #[cfg(feature = "enterprise")] @@ -256,7 +256,7 @@ impl SchedulerCtx for DefaultGcSchedulerCtx { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { self.gc_regions_inner( region_ids, @@ -310,7 +310,7 @@ impl DefaultGcSchedulerCtx { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { debug!( "Sending GC instruction for {} regions (full_file_listing: {})", @@ -326,9 +326,9 @@ impl DefaultGcSchedulerCtx { full_file_listing, timeout, region_routes_override, - event_context, ); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); let id = procedure_with_id.id; diff --git a/src/meta-srv/src/gc/handler.rs b/src/meta-srv/src/gc/handler.rs index 4532e9a8ca3..16b29a39c7e 100644 --- a/src/meta-srv/src/gc/handler.rs +++ b/src/meta-srv/src/gc/handler.rs @@ -18,7 +18,7 @@ use std::time::Instant; use common_catalog::consts::MITO_ENGINE; use common_meta::datanode::{RegionManifestInfo, RegionStat}; use common_meta::peer::Peer; -use common_meta::rpc::ddl::EventContext; +use common_meta::rpc::ddl::PersistentEventContext; use common_telemetry::tracing::Instrument as _; use common_telemetry::{debug, error, info, warn}; use futures::StreamExt; @@ -37,7 +37,10 @@ use crate::gc::tracker::RegionGcInfo; use crate::metrics::METRIC_META_GC_CANDIDATE_REGIONS; impl GcScheduler { - pub(crate) async fn trigger_gc(&self, event_context: EventContext) -> Result { + pub(crate) async fn trigger_gc( + &self, + event_context: PersistentEventContext, + ) -> Result { let start_time = Instant::now(); info!("Starting GC cycle"); @@ -212,7 +215,7 @@ impl GcScheduler { datanode_to_candidates: HashMap>, force_full_listing_by_peer: HashMap>, region_routes_override_by_peer: HashMap, - event_context: EventContext, + event_context: PersistentEventContext, ) -> GcJobReport { let mut per_datanode_reports = HashMap::new(); let mut failed_datanodes: HashMap<_, Vec<_>> = HashMap::new(); @@ -283,7 +286,7 @@ impl GcScheduler { candidates: Vec<(TableId, GcCandidate)>, force_full_listing: HashSet, region_routes_override: Region2Peers, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { info!( "Starting GC for datanode {} with {} candidate regions", diff --git a/src/meta-srv/src/gc/mock.rs b/src/meta-srv/src/gc/mock.rs index a941476ab44..94749dfb469 100644 --- a/src/meta-srv/src/gc/mock.rs +++ b/src/meta-srv/src/gc/mock.rs @@ -30,7 +30,7 @@ use common_meta::datanode::{RegionManifestInfo, RegionStat}; use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; use common_meta::peer::Peer; -use common_meta::rpc::ddl::EventContext; +use common_meta::rpc::ddl::PersistentEventContext; use common_meta::rpc::router::{Region, RegionRoute}; use common_telemetry::debug; use ordered_float::OrderedFloat; @@ -75,7 +75,7 @@ pub struct MockSchedulerCtx { pub candidates: Arc>>>>, pub get_table_to_region_stats_calls: Arc>, pub gc_regions_calls: Arc>, - pub gc_event_contexts: Arc>>, + pub gc_event_contexts: Arc>>, // Error injection fields for testing pub get_table_to_region_stats_error: Arc>>, pub get_table_route_error: Arc>>, @@ -214,7 +214,7 @@ impl SchedulerCtx for MockSchedulerCtx { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - event_context: EventContext, + event_context: PersistentEventContext, ) -> Result { *self.gc_regions_calls.lock().unwrap() += 1; self.gc_event_contexts.lock().unwrap().push(event_context); diff --git a/src/meta-srv/src/gc/mock/basic.rs b/src/meta-srv/src/gc/mock/basic.rs index 45d64910561..ed051cc839e 100644 --- a/src/meta-srv/src/gc/mock/basic.rs +++ b/src/meta-srv/src/gc/mock/basic.rs @@ -17,7 +17,7 @@ use std::sync::{Arc, Mutex}; use std::time::Instant; use common_meta::peer::Peer; -use common_meta::rpc::ddl::{EventContext, TriggerReason}; +use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; use common_telemetry::init_default_ut_logging; use store_api::region_engine::RegionRole; use store_api::storage::{FileId, FileRefsManifest, GcReport, RegionId}; @@ -36,7 +36,7 @@ async fn test_parallel_process_datanodes_empty() { HashMap::new(), HashMap::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await; @@ -105,7 +105,7 @@ async fn test_parallel_process_datanodes_with_candidates() { datanode_to_candidates, HashMap::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await; @@ -189,7 +189,7 @@ async fn test_handle_tick() { assert_eq!(*ctx.gc_regions_calls.lock().unwrap(), 1); assert_eq!( ctx.gc_event_contexts.lock().unwrap().as_slice(), - &[EventContext::new(TriggerReason::ScheduledGc)] + &[PersistentEventContext::new(TriggerReason::ScheduledGc)] ); let tracker = scheduler.region_gc_tracker.lock().await; @@ -249,6 +249,6 @@ async fn test_handle_manual_gc_without_regions_records_manual_event_context() { assert_eq!( ctx.gc_event_contexts.lock().unwrap().as_slice(), - &[EventContext::new(TriggerReason::Manual)] + &[PersistentEventContext::new(TriggerReason::Manual)] ); } diff --git a/src/meta-srv/src/gc/mock/concurrent.rs b/src/meta-srv/src/gc/mock/concurrent.rs index 7dca378b71d..5bce9a1be8b 100644 --- a/src/meta-srv/src/gc/mock/concurrent.rs +++ b/src/meta-srv/src/gc/mock/concurrent.rs @@ -18,7 +18,7 @@ use std::time::{Duration, Instant}; use common_meta::key::table_route::PhysicalTableRouteValue; use common_meta::peer::Peer; -use common_meta::rpc::ddl::EventContext; +use common_meta::rpc::ddl::PersistentEventContext; use common_meta::rpc::router::{Region, RegionRoute}; use common_telemetry::{info, init_default_ut_logging}; use store_api::region_engine::RegionRole; @@ -102,7 +102,7 @@ async fn test_concurrent_table_processing_limits() { datanode_to_candidates, HashMap::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await; @@ -190,7 +190,7 @@ async fn test_datanode_processes_tables_with_partial_gc_failures() { datanode_to_candidates, HashMap::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await; @@ -296,7 +296,7 @@ async fn test_region_gc_concurrency_limit() { candidates.into_iter().map(|c| (table_id, c)).collect(), HashSet::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await .unwrap(); @@ -413,7 +413,7 @@ async fn test_region_gc_concurrency_with_partial_failures() { datanode_to_candidates, HashMap::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await; @@ -557,7 +557,7 @@ async fn test_region_gc_concurrency_with_retryable_errors() { datanode_to_candidates, HashMap::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await; diff --git a/src/meta-srv/src/gc/mock/full_list.rs b/src/meta-srv/src/gc/mock/full_list.rs index a5d5012c6d3..8193a9094e1 100644 --- a/src/meta-srv/src/gc/mock/full_list.rs +++ b/src/meta-srv/src/gc/mock/full_list.rs @@ -17,7 +17,7 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use common_meta::peer::Peer; -use common_meta::rpc::ddl::EventContext; +use common_meta::rpc::ddl::PersistentEventContext; use common_telemetry::init_default_ut_logging; use store_api::region_engine::RegionRole; use store_api::storage::{FileId, FileRefsManifest, GcReport, RegionId}; @@ -83,7 +83,7 @@ async fn test_full_file_listing_first_time_gc() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await .unwrap(); @@ -157,7 +157,7 @@ async fn test_full_file_listing_interval_enforcement() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await .unwrap(); @@ -184,7 +184,7 @@ async fn test_full_file_listing_interval_enforcement() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await .unwrap(); @@ -260,7 +260,7 @@ async fn test_full_file_listing_no_interval_passed() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await .unwrap(); @@ -284,7 +284,7 @@ async fn test_full_file_listing_no_interval_passed() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await .unwrap(); diff --git a/src/meta-srv/src/gc/mock/misc.rs b/src/meta-srv/src/gc/mock/misc.rs index 9b3c35cbe9a..575a5c5fc08 100644 --- a/src/meta-srv/src/gc/mock/misc.rs +++ b/src/meta-srv/src/gc/mock/misc.rs @@ -17,7 +17,7 @@ use std::sync::{Arc, Mutex}; use std::time::Instant; use common_meta::peer::Peer; -use common_meta::rpc::ddl::EventContext; +use common_meta::rpc::ddl::PersistentEventContext; use common_telemetry::init_default_ut_logging; use store_api::storage::{FileRefsManifest, GcReport, RegionId}; @@ -76,7 +76,7 @@ async fn test_empty_file_refs_manifest() { datanode_to_candidates, HashMap::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await; @@ -169,7 +169,7 @@ async fn test_multiple_regions_per_table() { datanode_to_candidates, HashMap::new(), HashMap::new(), - EventContext::default(), + PersistentEventContext::default(), ) .await; diff --git a/src/meta-srv/src/gc/procedure.rs b/src/meta-srv/src/gc/procedure.rs index d745e88e2c5..f9f597577ae 100644 --- a/src/meta-srv/src/gc/procedure.rs +++ b/src/meta-srv/src/gc/procedure.rs @@ -24,11 +24,11 @@ use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; use common_meta::lock_key::{RegionLock, TableLock}; use common_meta::peer::Peer; -use common_meta::rpc::ddl::{EventContext, TriggerReason}; +use common_meta::rpc::ddl::TriggerReason; use common_procedure::error::ToJsonSnafu; use common_procedure::{ - Context as ProcedureContext, Error as ProcedureError, EventRuntimeContext, EventTrigger, - LockKey, Procedure, ProcedureState, Result as ProcedureResult, Status, + Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, + Procedure, ProcedureState, Result as ProcedureResult, Status, }; use common_telemetry::tracing::Instrument as _; use common_telemetry::tracing_context::TracingContext; @@ -200,8 +200,6 @@ pub struct BatchGcData { /// mailbox timeout duration timeout: Duration, gc_report: Option, - #[serde(default)] - event_context: EventContext, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] @@ -219,7 +217,6 @@ pub enum State { impl BatchGcProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::BatchGcProcedure"; - #[allow(clippy::too_many_arguments)] pub fn new( mailbox: MailboxRef, table_metadata_manager: TableMetadataManagerRef, @@ -228,7 +225,6 @@ impl BatchGcProcedure { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, - event_context: EventContext, ) -> Self { Self { mailbox, @@ -244,7 +240,6 @@ impl BatchGcProcedure { related_regions: HashMap::new(), file_refs: FileRefsManifest::default(), gc_report: None, - event_context, }, } } @@ -275,7 +270,6 @@ impl BatchGcProcedure { related_regions: HashMap::new(), file_refs, gc_report: Some(GcReport::default()), - event_context: EventContext::default(), }, } } @@ -1072,23 +1066,21 @@ impl Procedure for BatchGcProcedure { LockKey::new(lock_key) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(BATCH_GC_EVENT_TYPE) { return None; } let event = match &ctx.trigger { // Keep scheduled GC low-noise; record submitted manual requests for auditability. - EventTrigger::Submitted => (self.data.event_context.reason == TriggerReason::Manual) + EventTrigger::Submitted => ctx + .event_context + .is_some_and(|context| context.reason == TriggerReason::Manual) .then(|| { BatchGcEvent::with_config( &self.data.regions, self.data.full_file_listing, self.data.timeout, - Some(self.data.event_context.clone()), ) })?, EventTrigger::Recovered | EventTrigger::ChildSubmitted { .. } => return None, @@ -1106,7 +1098,6 @@ impl Procedure for BatchGcProcedure { &self.data.regions, self.data.full_file_listing, self.data.timeout, - None, ), EventTrigger::Failed | EventTrigger::Poisoned => self .data @@ -1118,7 +1109,6 @@ impl Procedure for BatchGcProcedure { &self.data.regions, self.data.full_file_listing, self.data.timeout, - None, ) }), }; @@ -1211,16 +1201,6 @@ mod tests { assert!(report.need_retry_regions.contains(®ion_id)); } - #[test] - fn test_batch_gc_data_defaults_missing_event_context() { - let mut data = serde_json::to_value(&batch_gc_procedure().data).unwrap(); - data.as_object_mut().unwrap().remove("event_context"); - - let recovered: BatchGcData = serde_json::from_value(data).unwrap(); - - assert_eq!(recovered.event_context, EventContext::default()); - } - #[tokio::test] async fn test_send_gc_instructions_preserves_partial_report() { let first_region = RegionId::new(1024, 1); @@ -1255,7 +1235,6 @@ mod tests { true, Duration::from_secs(10), HashMap::new(), - EventContext::default(), ); procedure.data.region_routes = HashMap::from([ (first_region, (first_peer, vec![])), @@ -1296,7 +1275,6 @@ mod tests { true, Duration::from_secs(10), HashMap::new(), - EventContext::default(), ) } } diff --git a/src/meta-srv/src/gc/scheduler.rs b/src/meta-srv/src/gc/scheduler.rs index 92b362ad947..b89c0d06be9 100644 --- a/src/meta-srv/src/gc/scheduler.rs +++ b/src/meta-srv/src/gc/scheduler.rs @@ -18,7 +18,7 @@ use std::time::{Duration, Instant}; use common_meta::DatanodeId; use common_meta::key::runtime_switch::RuntimeSwitchManagerRef; -use common_meta::rpc::ddl::{EventContext, TriggerReason}; +use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; use common_telemetry::tracing::Instrument as _; use common_telemetry::{error, info}; use snafu::ResultExt; @@ -218,7 +218,7 @@ impl GcScheduler { } let span = common_telemetry::tracing::info_span!("meta_gc_handle_tick"); let report = self - .trigger_gc(EventContext::new(TriggerReason::ScheduledGc)) + .trigger_gc(PersistentEventContext::new(TriggerReason::ScheduledGc)) .instrument(span) .await?; @@ -305,7 +305,7 @@ impl GcScheduler { // No specific regions, use default tick behavior let Some(regions) = region_ids else { let report = self - .trigger_gc(EventContext::new(TriggerReason::Manual)) + .trigger_gc(PersistentEventContext::new(TriggerReason::Manual)) .await?; info!("Finished manual gc request"); return Ok(report); @@ -353,7 +353,7 @@ impl GcScheduler { full_listing, gc_timeout, Region2Peers::new(), - EventContext::new(TriggerReason::Manual), + PersistentEventContext::new(TriggerReason::Manual), ) .await?; combined_report.merge(report); @@ -367,7 +367,7 @@ impl GcScheduler { true, gc_timeout, dropped_routes_override, - EventContext::new(TriggerReason::Manual), + PersistentEventContext::new(TriggerReason::Manual), ) .await?; combined_report.merge(report); @@ -490,7 +490,7 @@ mod tests { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - _event_context: EventContext, + _event_context: PersistentEventContext, ) -> Result { self.gc_regions_calls.fetch_add(1, Ordering::Relaxed); panic!("gc_regions should not be called in maintenance mode") @@ -564,7 +564,7 @@ mod tests { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - _event_context: EventContext, + _event_context: PersistentEventContext, ) -> Result { Ok(GcReport::default()) } @@ -898,7 +898,7 @@ mod tests { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - _event_context: EventContext, + _event_context: PersistentEventContext, ) -> Result { crate::error::UnexpectedSnafu { violated: "mock gc failure".to_string(), diff --git a/src/meta-srv/src/procedure/region_migration.rs b/src/meta-srv/src/procedure/region_migration.rs index 25fc7d6a309..ad1553c92d7 100644 --- a/src/meta-srv/src/procedure/region_migration.rs +++ b/src/meta-srv/src/procedure/region_migration.rs @@ -45,13 +45,12 @@ use common_meta::kv_backend::{KvBackendRef, ResettableKvBackendRef}; use common_meta::lock_key::{CatalogLock, RegionLock, SchemaLock, TableLock}; use common_meta::peer::Peer; use common_meta::region_keeper::{MemoryRegionKeeperRef, OperatingRegionGuard}; -use common_meta::rpc::ddl::{EventContext, TriggerReason}; +use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; use common_procedure::error::{ Error as ProcedureError, FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu, }; use common_procedure::{ - Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, Status, - StringKey, + Context as ProcedureContext, EventContext, LockKey, Procedure, Status, StringKey, }; use common_telemetry::{debug, error, info}; use manager::RegionMigrationProcedureGuard; @@ -123,11 +122,9 @@ pub struct PersistentContext { /// The timeout for downgrading leader region and upgrading candidate region operations. #[serde(with = "humantime_serde", default = "default_timeout")] pub(crate) timeout: Duration, - /// The trigger reason of region migration. + /// The trigger reason persisted for compatibility with versions without procedure context. #[serde(default)] - pub(crate) trigger_reason: RegionMigrationTriggerReason, - #[serde(default)] - pub(crate) event_context: EventContext, + trigger_reason: RegionMigrationTriggerReason, } impl PersistentContext { @@ -137,7 +134,7 @@ impl PersistentContext { to_peer: Peer, region_ids: Vec, timeout: Duration, - event_context: EventContext, + trigger_reason: RegionMigrationTriggerReason, ) -> Self { #[allow(deprecated)] Self { @@ -148,17 +145,7 @@ impl PersistentContext { to_peer, region_ids, timeout, - trigger_reason: RegionMigrationTriggerReason::from_trigger_reason(event_context.reason), - event_context, - } - } - - fn normalize_trigger_reasons(&mut self) { - if self.trigger_reason != RegionMigrationTriggerReason::Unknown { - self.event_context.reason = self.trigger_reason.to_trigger_reason(); - } else if self.event_context.reason != TriggerReason::Unknown { - self.trigger_reason = - RegionMigrationTriggerReason::from_trigger_reason(self.event_context.reason); + trigger_reason, } } } @@ -173,7 +160,7 @@ impl RegionMigrationTriggerReason { } } - fn from_trigger_reason(reason: TriggerReason) -> Self { + pub(crate) fn from_trigger_reason(reason: TriggerReason) -> Self { match reason { TriggerReason::Manual => Self::Manual, TriggerReason::AutoRebalance => Self::AutoRebalance, @@ -457,6 +444,14 @@ pub struct Context { } impl Context { + pub(crate) fn trigger_reason( + &self, + event_context: Option<&PersistentEventContext>, + ) -> RegionMigrationTriggerReason { + event_context + .map(|ctx| RegionMigrationTriggerReason::from_trigger_reason(ctx.reason)) + .unwrap_or(self.persistent_ctx.trigger_reason) + } /// Returns the next operation's timeout. pub fn next_operation_timeout(&self) -> Option { self.persistent_ctx @@ -852,10 +847,9 @@ impl RegionMigrationProcedure { tracker: RegionMigrationProcedureTracker, ) -> ProcedureResult { let RegionMigrationDataOwned { - mut persistent_ctx, + persistent_ctx, state, } = serde_json::from_str(json).context(FromJsonSnafu)?; - persistent_ctx.normalize_trigger_reasons(); let guards = persistent_ctx .region_ids .iter() @@ -999,14 +993,16 @@ impl Procedure for RegionMigrationProcedure { LockKey::new(self.context.persistent_ctx.lock_key()) } - fn event(&self, ctx: &EventRuntimeContext<'_>) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(REGION_MIGRATION_EVENT_TYPE) { return None; } Some(Box::new(RegionMigrationEvent::from_persistent_ctx( &self.context.persistent_ctx, - matches!(ctx.trigger, EventTrigger::Submitted), + self.context + .trigger_reason(ctx.event_context) + .to_trigger_reason(), ))) } } @@ -1063,7 +1059,7 @@ mod tests { let procedure = RegionMigrationProcedure::new(persistent_context, context, vec![]); let serialized = procedure.dump().unwrap(); - let expected = r#"{"persistent_ctx":{"catalog_and_schema":[["greptime","public"]],"from_peer":{"id":1,"addr":""},"to_peer":{"id":2,"addr":""},"region_ids":[4398046511105],"timeout":"10s","trigger_reason":"Unknown","event_context":{"reason":"unknown"}},"state":{"region_migration_state":"RegionMigrationStart"}}"#; + let expected = r#"{"persistent_ctx":{"catalog_and_schema":[["greptime","public"]],"from_peer":{"id":1,"addr":""},"to_peer":{"id":2,"addr":""},"region_ids":[4398046511105],"timeout":"10s","trigger_reason":"Unknown"},"state":{"region_migration_state":"RegionMigrationStart"}}"#; assert_eq!(expected, serialized); } @@ -1087,11 +1083,12 @@ mod tests { for trigger in triggers { let event = procedure - .event(&EventRuntimeContext { + .event(&EventContext { procedure_id: common_procedure::ProcedureId::random(), lifecycle_state: &state, trigger, event_type_filter: Arc::new(common_event_recorder::EventTypeFilter::All), + event_context: None, }) .unwrap(); assert_eq!(event.event_type(), "region_migration"); @@ -1111,61 +1108,98 @@ mod tests { to_peer: Peer::empty(2), region_ids: vec![RegionId::new(1024, 1)], timeout: Duration::from_secs(10), - trigger_reason: RegionMigrationTriggerReason::default(), - event_context: EventContext::default(), + trigger_reason: RegionMigrationTriggerReason::Unknown, }; // NOTES: Changes it will break backward compatibility. let serialized = r#"{"catalog":"greptime","schema":"public","from_peer":{"id":1,"addr":""},"to_peer":{"id":2,"addr":""},"region_id":4398046511105}"#; let deserialized: PersistentContext = serde_json::from_str(serialized).unwrap(); assert_eq!(persistent_ctx, deserialized); - assert_eq!(deserialized.event_context, EventContext::default()); } #[test] - fn test_new_persistent_context_keeps_event_context() { - let mut event_context = - EventContext::new(TriggerReason::AutoRebalance).with_protocol("internal"); - event_context - .extensions - .insert("scheduler".into(), serde_json::json!("load_based")); - - let context = PersistentContext::new( - vec![("greptime".into(), "public".into())], - Peer::empty(1), - Peer::empty(2), - vec![RegionId::new(1024, 1)], - Duration::from_secs(10), - event_context.clone(), + fn test_legacy_trigger_reason_survives_recovery_and_repersistence() { + let serialized = r#"{"persistent_ctx":{"catalog":"greptime","schema":"public","from_peer":{"id":1,"addr":""},"to_peer":{"id":2,"addr":""},"region_id":4398046511105,"trigger_reason":"Failover"},"state":{"region_migration_state":"RegionMigrationStart"}}"#; + let env = TestingEnv::new(); + let procedure = RegionMigrationProcedure::from_json( + serialized, + env.context_factory(), + RegionMigrationProcedureTracker::default(), + ) + .unwrap(); + assert_eq!( + procedure.context.persistent_ctx.trigger_reason, + RegionMigrationTriggerReason::Failover + ); + let repersisted = procedure.dump().unwrap(); + assert_eq!( + serde_json::from_str::(&repersisted).unwrap()["persistent_ctx"]["trigger_reason"], + "Failover" ); + let recovered = RegionMigrationProcedure::from_json( + &repersisted, + env.context_factory(), + RegionMigrationProcedureTracker::default(), + ) + .unwrap(); assert_eq!( - context.trigger_reason, + recovered.context.trigger_reason(None), + RegionMigrationTriggerReason::Failover + ); + let event_context = PersistentEventContext::new(TriggerReason::AutoRebalance); + assert_eq!( + recovered.context.trigger_reason(Some(&event_context)), RegionMigrationTriggerReason::AutoRebalance ); - assert_eq!(context.event_context, event_context); + + let state = common_procedure::ProcedureState::Running; + let event = recovered + .event(&EventContext { + procedure_id: common_procedure::ProcedureId::random(), + lifecycle_state: &state, + trigger: common_procedure::EventTrigger::Recovered, + event_type_filter: Arc::new(common_event_recorder::EventTypeFilter::All), + event_context: None, + }) + .unwrap(); + assert_eq!( + event.extra_rows().unwrap()[0].values[3].value_data, + Some(api::v1::value::ValueData::StringValue( + "Failover".to_string() + )) + ); } #[test] - fn test_normalize_trigger_reasons() { - let mut context = new_persistent_context(); - context.trigger_reason = RegionMigrationTriggerReason::Manual; - context.event_context = EventContext::default(); - context.normalize_trigger_reasons(); - assert_eq!(context.event_context.reason, TriggerReason::Manual); + fn test_migration_reason_uses_event_context() { + let env = TestingEnv::new(); + let context = env.context_factory().new_context(new_persistent_context()); + let event_context = PersistentEventContext::new(TriggerReason::RegionFailover); - context.trigger_reason = RegionMigrationTriggerReason::Unknown; - context.event_context = EventContext::new(TriggerReason::RegionFailover); - context.normalize_trigger_reasons(); assert_eq!( - context.trigger_reason, + context.trigger_reason(Some(&event_context)), RegionMigrationTriggerReason::Failover ); - context.trigger_reason = RegionMigrationTriggerReason::Manual; - context.event_context = EventContext::new(TriggerReason::RegionFailover); - context.normalize_trigger_reasons(); - assert_eq!(context.event_context.reason, TriggerReason::Manual); + let procedure = + RegionMigrationProcedure::new(new_persistent_context(), env.context_factory(), vec![]); + let state = common_procedure::ProcedureState::Running; + let event = procedure + .event(&EventContext { + procedure_id: common_procedure::ProcedureId::random(), + lifecycle_state: &state, + trigger: common_procedure::EventTrigger::Submitted, + event_type_filter: Arc::new(common_event_recorder::EventTypeFilter::All), + event_context: Some(&event_context), + }) + .unwrap(); + assert_eq!( + event.extra_rows().unwrap()[0].values[3].value_data, + Some(api::v1::value::ValueData::StringValue( + "Failover".to_string() + )) + ); } #[derive(Debug, Serialize, Deserialize, Default)] diff --git a/src/meta-srv/src/procedure/region_migration/close_downgraded_region.rs b/src/meta-srv/src/procedure/region_migration/close_downgraded_region.rs index 5cdcad4503b..29cd1804d03 100644 --- a/src/meta-srv/src/procedure/region_migration/close_downgraded_region.rs +++ b/src/meta-srv/src/procedure/region_migration/close_downgraded_region.rs @@ -40,7 +40,7 @@ impl State for CloseDowngradedRegion { async fn next( &mut self, ctx: &mut Context, - _procedure_ctx: &ProcedureContext, + procedure_ctx: &ProcedureContext, ) -> Result<(Box, Status)> { if let Err(err) = self.close_downgraded_leader_region(ctx).await { let downgrade_leader_datanode = &ctx.persistent_ctx.from_peer; @@ -52,7 +52,7 @@ impl State for CloseDowngradedRegion { ctx.persistent_ctx.region_ids, ctx.persistent_ctx.from_peer, ctx.persistent_ctx.to_peer, - ctx.persistent_ctx.trigger_reason, + ctx.trigger_reason(procedure_ctx.event_context.as_ref()), ctx.volatile_ctx.metrics, ); Ok((Box::new(RegionMigrationEnd), Status::done())) diff --git a/src/meta-srv/src/procedure/region_migration/downgrade_leader_region.rs b/src/meta-srv/src/procedure/region_migration/downgrade_leader_region.rs index c58c1f8cd1f..97c3a8e2f31 100644 --- a/src/meta-srv/src/procedure/region_migration/downgrade_leader_region.rs +++ b/src/meta-srv/src/procedure/region_migration/downgrade_leader_region.rs @@ -382,7 +382,6 @@ mod tests { use common_meta::key::table_route::TableRouteValue; use common_meta::key::test_utils::new_test_table_info; use common_meta::peer::Peer; - use common_meta::rpc::ddl::{EventContext, TriggerReason}; use common_meta::rpc::router::{Region, RegionRoute}; use common_meta::wal_provider::RegionWalOptions; use store_api::storage::RegionId; @@ -391,7 +390,9 @@ mod tests { use super::*; use crate::error::Error; use crate::procedure::region_migration::test_util::{TestingEnv, new_procedure_context}; - use crate::procedure::region_migration::{ContextFactory, PersistentContext}; + use crate::procedure::region_migration::{ + ContextFactory, PersistentContext, RegionMigrationTriggerReason, + }; use crate::procedure::test_util::{ new_close_region_reply, new_downgrade_region_reply, send_mock_reply, }; @@ -403,7 +404,7 @@ mod tests { Peer::empty(2), vec![RegionId::new(1024, 1)], Duration::from_millis(1000), - EventContext::new(TriggerReason::Manual), + RegionMigrationTriggerReason::Unknown, ) } diff --git a/src/meta-srv/src/procedure/region_migration/manager.rs b/src/meta-srv/src/procedure/region_migration/manager.rs index 3de79735e3b..e1d719d9747 100644 --- a/src/meta-srv/src/procedure/region_migration/manager.rs +++ b/src/meta-srv/src/procedure/region_migration/manager.rs @@ -21,7 +21,7 @@ use std::time::Duration; use common_meta::key::table_info::TableInfoValue; use common_meta::key::table_route::TableRouteValue; use common_meta::peer::Peer; -use common_meta::rpc::ddl::EventContext; +use common_meta::rpc::ddl::PersistentEventContext; use common_meta::rpc::router::RegionRoute; use common_procedure::{ProcedureId, ProcedureManagerRef, ProcedureWithId, watcher}; use common_telemetry::{error, info, warn}; @@ -479,6 +479,7 @@ impl RegionMigrationManager { procedure_guards: Vec, catalog_and_schema: Vec<(String, String)>, ) -> Result { + let event_context = PersistentEventContext::new(task.trigger_reason.to_trigger_reason()); let procedure = RegionMigrationProcedure::new( PersistentContext::new( catalog_and_schema, @@ -486,12 +487,13 @@ impl RegionMigrationManager { task.to_peer.clone(), task.region_ids.clone(), task.timeout, - EventContext::new(task.trigger_reason.to_trigger_reason()), + task.trigger_reason, ), self.context_factory.clone(), procedure_guards, ); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); let procedure_id = procedure_with_id.id; info!("Starting region migration procedure {procedure_id} for {task}"); let procedure_manager = self.procedure_manager.clone(); @@ -575,6 +577,7 @@ impl RegionMigrationManager { timeout, trigger_reason, } = task.clone(); + let event_context = PersistentEventContext::new(trigger_reason.to_trigger_reason()); let procedure = RegionMigrationProcedure::new( PersistentContext::new( vec![(catalog_name, schema_name)], @@ -582,12 +585,13 @@ impl RegionMigrationManager { to_peer, vec![region_id], timeout, - EventContext::new(trigger_reason.to_trigger_reason()), + trigger_reason, ), self.context_factory.clone(), vec![guard], ); - let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)); + let procedure_with_id = + ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); let procedure_id = procedure_with_id.id; info!("Starting region migration procedure {procedure_id} for {task}"); let procedure_manager = self.procedure_manager.clone(); diff --git a/src/meta-srv/src/procedure/region_migration/migration_abort.rs b/src/meta-srv/src/procedure/region_migration/migration_abort.rs index f3ad8052de5..981c2a20252 100644 --- a/src/meta-srv/src/procedure/region_migration/migration_abort.rs +++ b/src/meta-srv/src/procedure/region_migration/migration_abort.rs @@ -41,7 +41,7 @@ impl State for RegionMigrationAbort { async fn next( &mut self, ctx: &mut Context, - _procedure_ctx: &ProcedureContext, + procedure_ctx: &ProcedureContext, ) -> Result<(Box, Status)> { warn!( "Region migration is aborted: {}, regions: {:?}, from_peer: {}, to_peer: {}, trigger_reason: {}, {}", @@ -49,7 +49,7 @@ impl State for RegionMigrationAbort { ctx.persistent_ctx.region_ids, ctx.persistent_ctx.from_peer, ctx.persistent_ctx.to_peer, - ctx.persistent_ctx.trigger_reason, + ctx.trigger_reason(procedure_ctx.event_context.as_ref()), ctx.volatile_ctx.metrics, ); error::MigrationAbortSnafu { diff --git a/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs b/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs index 51d7054c7c5..0cb1131e545 100644 --- a/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs +++ b/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs @@ -47,9 +47,12 @@ impl State for OpenCandidateRegion { async fn next( &mut self, ctx: &mut Context, - _procedure_ctx: &ProcedureContext, + procedure_ctx: &ProcedureContext, ) -> Result<(Box, Status)> { - let instruction = self.build_open_region_instruction(ctx).await?; + let trigger_reason = ctx.trigger_reason(procedure_ctx.event_context.as_ref()); + let instruction = self + .build_open_region_instruction(ctx, trigger_reason) + .await?; let now = Instant::now(); self.open_candidate_region(ctx, instruction).await?; ctx.update_open_candidate_region_elapsed(now); @@ -67,11 +70,15 @@ impl OpenCandidateRegion { /// /// Abort(non-retry): /// - Datanode Table is not found. - async fn build_open_region_instruction(&self, ctx: &mut Context) -> Result { + async fn build_open_region_instruction( + &self, + ctx: &mut Context, + trigger_reason: RegionMigrationTriggerReason, + ) -> Result { let region_ids = ctx.persistent_ctx.region_ids.clone(); let from_peer_id = ctx.persistent_ctx.from_peer.id; let to_peer_id = ctx.persistent_ctx.to_peer.id; - let reason = match ctx.persistent_ctx.trigger_reason { + let reason = match trigger_reason { RegionMigrationTriggerReason::Failover => OpenRegionReason::RegionFailover, _ => OpenRegionReason::RegionMigration, }; @@ -280,7 +287,7 @@ mod tests { let mut ctx = env.context_factory().new_context(persistent_context); let err = state - .build_open_region_instruction(&mut ctx) + .build_open_region_instruction(&mut ctx, RegionMigrationTriggerReason::Unknown) .await .unwrap_err(); @@ -291,7 +298,7 @@ mod tests { #[tokio::test] async fn test_build_open_region_instruction_reason() { let state = OpenCandidateRegion; - let mut persistent_context = new_persistent_context(); + let persistent_context = new_persistent_context(); let from_peer_id = persistent_context.from_peer.id; let region_id = persistent_context.region_ids[0]; let env = TestingEnv::new(); @@ -314,7 +321,10 @@ mod tests { let mut ctx = env .context_factory() .new_context(persistent_context.clone()); - let instruction = state.build_open_region_instruction(&mut ctx).await.unwrap(); + let instruction = state + .build_open_region_instruction(&mut ctx, RegionMigrationTriggerReason::Unknown) + .await + .unwrap(); let open_regions = instruction.into_open_regions().unwrap(); assert_eq!( Some(OpenRegionReason::RegionMigration), @@ -325,9 +335,11 @@ mod tests { open_regions[0].requirements ); - persistent_context.trigger_reason = RegionMigrationTriggerReason::Failover; let mut ctx = env.context_factory().new_context(persistent_context); - let instruction = state.build_open_region_instruction(&mut ctx).await.unwrap(); + let instruction = state + .build_open_region_instruction(&mut ctx, RegionMigrationTriggerReason::Failover) + .await + .unwrap(); let open_regions = instruction.into_open_regions().unwrap(); assert_eq!( Some(OpenRegionReason::RegionFailover), diff --git a/src/meta-srv/src/procedure/region_migration/test_util.rs b/src/meta-srv/src/procedure/region_migration/test_util.rs index 70342608646..44180dc6f8c 100644 --- a/src/meta-srv/src/procedure/region_migration/test_util.rs +++ b/src/meta-srv/src/procedure/region_migration/test_util.rs @@ -27,7 +27,6 @@ use common_meta::kv_backend::KvBackendRef; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::peer::Peer; use common_meta::region_keeper::{MemoryRegionKeeper, MemoryRegionKeeperRef}; -use common_meta::rpc::ddl::EventContext; use common_meta::rpc::router::RegionRoute; use common_meta::sequence::SequenceBuilder; use common_meta::state_store::KvStateStore; @@ -53,7 +52,8 @@ use crate::procedure::region_migration::open_candidate_region::OpenCandidateRegi use crate::procedure::region_migration::update_metadata::UpdateMetadata; use crate::procedure::region_migration::upgrade_candidate_region::UpgradeCandidateRegion; use crate::procedure::region_migration::{ - Context, ContextFactory, DefaultContextFactory, PersistentContext, State, VolatileContext, + Context, ContextFactory, DefaultContextFactory, PersistentContext, + RegionMigrationTriggerReason, State, VolatileContext, }; use crate::procedure::test_util::{MailboxContext, send_mock_reply}; use crate::service::mailbox::Channel; @@ -157,6 +157,7 @@ impl TestingEnv { ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, } } @@ -190,7 +191,7 @@ pub fn new_persistent_context(from: u64, to: u64, region_id: RegionId) -> Persis Peer::empty(to), vec![region_id], Duration::from_secs(10), - EventContext::default(), + RegionMigrationTriggerReason::Unknown, ) } @@ -577,5 +578,6 @@ pub fn new_procedure_context() -> ProcedureContext { ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, } } diff --git a/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs b/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs index 16d0270354b..86e1c35c10c 100644 --- a/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs +++ b/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs @@ -356,7 +356,6 @@ mod tests { use common_meta::key::topic_name::TopicNameKey; use common_meta::key::topic_region::{ReplayCheckpoint, TopicRegionKey, TopicRegionValue}; use common_meta::peer::Peer; - use common_meta::rpc::ddl::{EventContext, TriggerReason}; use common_meta::rpc::router::{Region, RegionRoute}; use common_meta::wal_provider::RegionWalOptions; use common_wal::options::KafkaWalOptions; @@ -365,7 +364,9 @@ mod tests { use super::*; use crate::error::Error; use crate::procedure::region_migration::test_util::{TestingEnv, new_procedure_context}; - use crate::procedure::region_migration::{ContextFactory, PersistentContext}; + use crate::procedure::region_migration::{ + ContextFactory, PersistentContext, RegionMigrationTriggerReason, + }; use crate::procedure::test_util::{ new_close_region_reply, new_upgrade_region_reply, send_mock_reply, }; @@ -377,7 +378,7 @@ mod tests { Peer::empty(2), vec![RegionId::new(1024, 1)], Duration::from_millis(1000), - EventContext::new(TriggerReason::Manual), + RegionMigrationTriggerReason::Unknown, ) } diff --git a/src/meta-srv/src/procedure/repartition.rs b/src/meta-srv/src/procedure/repartition.rs index d8b585b28ee..ee8b3ee749b 100644 --- a/src/meta-srv/src/procedure/repartition.rs +++ b/src/meta-srv/src/procedure/repartition.rs @@ -44,12 +44,11 @@ use common_meta::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock}; use common_meta::node_manager::NodeManagerRef; use common_meta::region_keeper::{MemoryRegionKeeperRef, OperatingRegionGuard}; use common_meta::region_registry::LeaderRegionRegistryRef; -use common_meta::rpc::ddl::EventContext; use common_meta::rpc::router::{RegionRoute, operating_leader_region_roles}; use common_meta::wal_provider::RegionWalOptions; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - BoxedProcedure, Context as ProcedureContext, Error as ProcedureError, EventRuntimeContext, + BoxedProcedure, Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, Procedure, ProcedureManagerRef, Result as ProcedureResult, Status, StringKey, }; @@ -104,8 +103,6 @@ pub struct PersistentContext { #[serde(default)] /// Records table-level partition metadata updated by this repartition. pub partition_metadata_update: Option, - #[serde(default)] - pub event_context: EventContext, } fn default_timeout() -> Duration { @@ -113,7 +110,7 @@ fn default_timeout() -> Duration { } impl PersistentContext { - /// Creates a new [PersistentContext] with the table, timeout and event context. + /// Creates a new [PersistentContext] with the given table name, table id and timeout. /// /// If the timeout is not provided, the default timeout will be used. pub fn new( @@ -124,7 +121,6 @@ impl PersistentContext { }: TableName, table_id: TableId, timeout: Option, - event_context: EventContext, ) -> Self { Self { catalog_name, @@ -136,7 +132,6 @@ impl PersistentContext { unknown_procedures: vec![], timeout: timeout.unwrap_or_else(default_timeout), partition_metadata_update: None, - event_context, } } @@ -797,10 +792,7 @@ impl Procedure for RepartitionProcedure { LockKey::new(self.context.persistent_ctx.lock_key()) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(REPARTITION_EVENT_TYPE) { return None; } @@ -869,7 +861,6 @@ impl RepartitionProcedureFactory for GcDisabledRepartitionProcedureFactory { _source: RepartitionSource, _to_exprs: Vec, _timeout: Option, - _event_context: EventContext, ) -> std::result::Result { Err(BoxedError::new( error::InvalidArgumentsSnafu { @@ -908,9 +899,8 @@ impl RepartitionProcedureFactory for DefaultRepartitionProcedureFactory { source: RepartitionSource, to_exprs: Vec, timeout: Option, - event_context: EventContext, ) -> std::result::Result { - let persistent_ctx = PersistentContext::new(table_name, table_id, timeout, event_context); + let persistent_ctx = PersistentContext::new(table_name, table_id, timeout); let from = match source { RepartitionSource::Partitioned { exprs, @@ -1143,7 +1133,6 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, - EventContext::default(), ); Context::new( @@ -1175,7 +1164,6 @@ mod tests { }, vec![], None, - EventContext::default(), ) .err() .expect("GC-disabled factory must reject repartition"); @@ -1277,53 +1265,58 @@ mod tests { let all = Arc::new(EventTypeFilter::All); let submitted = procedure - .event(&EventRuntimeContext { + .event(&EventContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Submitted, event_type_filter: all.clone(), + event_context: None, }) .unwrap(); assert_eq!(submitted.event_type(), REPARTITION_EVENT_TYPE); assert_ne!(submitted.json_payload().unwrap(), serde_json::Value::Null); let allowed = procedure - .event(&EventRuntimeContext { + .event(&EventContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Submitted, event_type_filter: Arc::new(EventTypeFilter::Only(HashSet::from([ REPARTITION_EVENT_TYPE.to_string(), ]))), + event_context: None, }) .unwrap(); assert_eq!(allowed.event_type(), REPARTITION_EVENT_TYPE); let succeeded = procedure - .event(&EventRuntimeContext { + .event(&EventContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Succeeded, event_type_filter: all, + event_context: None, }) .unwrap(); assert_eq!(succeeded.json_payload().unwrap(), serde_json::Value::Null); - let filtered = procedure.event(&EventRuntimeContext { + let filtered = procedure.event(&EventContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Submitted, event_type_filter: Arc::new(EventTypeFilter::Only(HashSet::from([ "another_event".to_string() ]))), + event_context: None, }); assert!(filtered.is_none()); - let empty = procedure.event(&EventRuntimeContext { + let empty = procedure.event(&EventContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Submitted, event_type_filter: Arc::new(EventTypeFilter::Only(HashSet::new())), + event_context: None, }); assert!(empty.is_none()); } @@ -1382,7 +1375,7 @@ mod tests { } #[test] - fn test_persistent_context_serde_defaults_for_new_fields() { + fn test_persistent_context_partition_metadata_update_serde_default() { let json = r#"{ "catalog_name":"test_catalog", "schema_name":"test_schema", @@ -1394,10 +1387,7 @@ mod tests { let persistent_ctx: PersistentContext = serde_json::from_str(json).unwrap(); - assert!(persistent_ctx.failed_procedures.is_empty()); - assert!(persistent_ctx.unknown_procedures.is_empty()); assert!(persistent_ctx.partition_metadata_update.is_none()); - assert_eq!(persistent_ctx.event_context, EventContext::default()); } #[tokio::test] @@ -1474,7 +1464,6 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, - EventContext::default(), ); persistent_ctx.plans = vec![with_rollback_metadata( test_plan(table_id), @@ -1541,7 +1530,6 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, - EventContext::default(), ); persistent_ctx.plans = vec![test_plan(table_id)]; let context = Context::new( @@ -1595,7 +1583,6 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, - EventContext::default(), ); let failed_plan = test_plan(table_id); let failed_plan = with_rollback_metadata( @@ -1729,7 +1716,6 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, - EventContext::default(), ); persistent_ctx.plans = vec![failed_plan, succeeded_plan.clone()]; persistent_ctx.failed_procedures = vec![ProcedureMeta { @@ -1839,7 +1825,6 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, - EventContext::default(), ); persistent_ctx.plans = vec![plan.clone()]; persistent_ctx.unknown_procedures = vec![ProcedureMeta { @@ -1901,7 +1886,6 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, - EventContext::default(), ); persistent_ctx.plans = vec![with_rollback_metadata( test_plan(table_id), @@ -2030,7 +2014,6 @@ mod tests { TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, - EventContext::default(), ); persistent_ctx.plans = vec![failed_merge_plan, succeeded_split_plan.clone()]; persistent_ctx.failed_procedures = vec![ProcedureMeta { diff --git a/src/meta-srv/src/procedure/repartition/allocate_region.rs b/src/meta-srv/src/procedure/repartition/allocate_region.rs index aa36e95be23..857ac946708 100644 --- a/src/meta-srv/src/procedure/repartition/allocate_region.rs +++ b/src/meta-srv/src/procedure/repartition/allocate_region.rs @@ -777,6 +777,7 @@ mod tests { concurrent_region_route, region_wal_options: test_region_wal_options(&[1, 2]), }), + event_context: None, }; let mut state = ExecutePlan; diff --git a/src/meta-srv/src/procedure/repartition/collect.rs b/src/meta-srv/src/procedure/repartition/collect.rs index 56d9a007a16..b71c3d6650f 100644 --- a/src/meta-srv/src/procedure/repartition/collect.rs +++ b/src/meta-srv/src/procedure/repartition/collect.rs @@ -136,7 +136,6 @@ mod tests { use common_error::mock::MockError; use common_error::status_code::StatusCode; - use common_meta::rpc::ddl::EventContext; use common_meta::test_util::MockDatanodeManager; use common_procedure::{ Context as ProcedureContext, ContextProvider, Error as ProcedureError, ProcedureId, @@ -194,7 +193,6 @@ mod tests { table::table_name::TableName::new("test_catalog", "test_schema", "test_table"), 1024, None, - EventContext::default(), ); let mut ctx = crate::procedure::repartition::Context::new( &ddl_ctx, @@ -229,7 +227,6 @@ mod tests { table::table_name::TableName::new("test_catalog", "test_schema", "test_table"), 1024, None, - EventContext::default(), ); let mut ctx = crate::procedure::repartition::Context::new( &ddl_ctx, @@ -249,6 +246,7 @@ mod tests { receiver: rx, inner: MockContextProvider::default(), }), + event_context: None, }; let mut state = Collect { inflight_procedures: vec![ProcedureMeta { diff --git a/src/meta-srv/src/procedure/repartition/gc_requirement.rs b/src/meta-srv/src/procedure/repartition/gc_requirement.rs index 11e82d294ac..1ae24e40316 100644 --- a/src/meta-srv/src/procedure/repartition/gc_requirement.rs +++ b/src/meta-srv/src/procedure/repartition/gc_requirement.rs @@ -244,6 +244,7 @@ mod tests { block_after_persist: true, persisted_tx: Some(persisted_tx), }), + context: Default::default(), }) .await .unwrap(); diff --git a/src/meta-srv/src/procedure/repartition/group.rs b/src/meta-srv/src/procedure/repartition/group.rs index a3c38b59557..63e05bcb150 100644 --- a/src/meta-srv/src/procedure/repartition/group.rs +++ b/src/meta-srv/src/procedure/repartition/group.rs @@ -38,8 +38,8 @@ use common_meta::peer::Peer; use common_meta::rpc::router::RegionRoute; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, Error as ProcedureError, EventRuntimeContext, EventTrigger, - LockKey, Procedure, ProcedureId, Result as ProcedureResult, Status, StringKey, + Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, + Procedure, ProcedureId, Result as ProcedureResult, Status, StringKey, }; use common_telemetry::{error, info}; use serde::{Deserialize, Serialize}; @@ -262,10 +262,7 @@ impl Procedure for RepartitionGroupProcedure { LockKey::new(self.context.persistent_ctx.lock_key()) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(REPARTITION_GROUP_EVENT_TYPE) { return None; } @@ -634,9 +631,7 @@ mod tests { use common_event_recorder::EventTypeFilter; use common_meta::key::TableMetadataManager; use common_meta::kv_backend::test_util::MockKvBackendBuilder; - use common_procedure::{ - EventRuntimeContext, EventTrigger, Procedure, ProcedureId, ProcedureState, - }; + use common_procedure::{EventContext, EventTrigger, Procedure, ProcedureId, ProcedureState}; use crate::error::Error; use crate::event::repartition::REPARTITION_GROUP_EVENT_TYPE; @@ -713,15 +708,16 @@ mod tests { context: env.create_context(new_persistent_context(1024, vec![], vec![])), }; let state = ProcedureState::Running; - let runtime_context = |event_type_filter| EventRuntimeContext { + let event_context = |event_type_filter| EventContext { procedure_id: ProcedureId::random(), lifecycle_state: &state, trigger: EventTrigger::Submitted, event_type_filter: Arc::new(event_type_filter), + event_context: None, }; let allowed = procedure - .event(&runtime_context(EventTypeFilter::Only(HashSet::from([ + .event(&event_context(EventTypeFilter::Only(HashSet::from([ REPARTITION_GROUP_EVENT_TYPE.to_string(), ])))) .unwrap(); @@ -729,14 +725,14 @@ mod tests { assert!( procedure - .event(&runtime_context(EventTypeFilter::Only(HashSet::from([ + .event(&event_context(EventTypeFilter::Only(HashSet::from([ "another_event".to_string(), ])))) .is_none() ); assert!( procedure - .event(&runtime_context(EventTypeFilter::Only(HashSet::new()))) + .event(&event_context(EventTypeFilter::Only(HashSet::new()))) .is_none() ); } diff --git a/src/meta-srv/src/procedure/repartition/test_util.rs b/src/meta-srv/src/procedure/repartition/test_util.rs index a6f2330b83a..2ce494ee7c9 100644 --- a/src/meta-srv/src/procedure/repartition/test_util.rs +++ b/src/meta-srv/src/procedure/repartition/test_util.rs @@ -24,7 +24,6 @@ use common_meta::kv_backend::KvBackendRef; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::node_manager::NodeManagerRef; use common_meta::peer::Peer; -use common_meta::rpc::ddl::EventContext; use common_meta::rpc::router::{Region, RegionRoute}; use common_meta::sequence::SequenceBuilder; use common_meta::test_util::new_ddl_context_with_kv_backend; @@ -103,6 +102,7 @@ impl TestingEnv { ProcedureContext { procedure_id: ProcedureId::random(), provider: Arc::new(MockContextProvider::default()), + event_context: None, } } @@ -225,7 +225,6 @@ pub fn new_parent_context( TableName::new("test_catalog", "test_schema", "test_table"), table_id, None, - EventContext::default(), ); ParentContext::new( @@ -266,6 +265,7 @@ pub fn procedure_context_with_receivers( receivers, inner: MockContextProvider::default(), }), + event_context: None, } } diff --git a/src/meta-srv/src/procedure/tests.rs b/src/meta-srv/src/procedure/tests.rs index dd855c66aa0..4f70dcca95c 100644 --- a/src/meta-srv/src/procedure/tests.rs +++ b/src/meta-srv/src/procedure/tests.rs @@ -29,7 +29,7 @@ use common_meta::ddl::test_util::create_table::{ }; use common_meta::key::table_route::{PhysicalTableRouteValue, TableRouteValue}; use common_meta::node_manager::NodeManagerRef; -use common_meta::rpc::ddl::{CreateTableTask, EventContext, QueryContext}; +use common_meta::rpc::ddl::CreateTableTask; use common_meta::rpc::router::{RegionRoute, find_leaders}; use common_procedure::Status; use store_api::storage::RegionId; @@ -91,8 +91,6 @@ fn create_table_task(table_name: Option<&str>) -> CreateTableTask { fn test_region_request_builder() { let mut procedure = CreateTableProcedure::new( create_table_task(None), - QueryContext::default(), - EventContext::default(), test_data::new_ddl_context(Arc::new(NodeClients::default())), ) .unwrap(); @@ -190,8 +188,6 @@ async fn test_on_datanode_create_regions() { let mut procedure = CreateTableProcedure::new( create_table_task(None), - QueryContext::default(), - EventContext::default(), test_data::new_ddl_context(node_manager), ) .unwrap(); @@ -264,12 +260,8 @@ async fn test_on_datanode_create_logical_regions() { .unwrap() .0; let _ = kv_backend.txn(physical_route_txn).await.unwrap(); - let mut procedure = CreateLogicalTablesProcedure::new( - vec![task1, task2, task3], - physical_table_id, - EventContext::default(), - ctx, - ); + let mut procedure = + CreateLogicalTablesProcedure::new(vec![task1, task2, task3], physical_table_id, ctx); let expected_created_regions = Arc::new(Mutex::new(HashMap::from([(1, 3), (2, 3), (3, 3)]))); diff --git a/src/meta-srv/src/procedure/wal_prune.rs b/src/meta-srv/src/procedure/wal_prune.rs index 128fc58c44d..5aca85a51a6 100644 --- a/src/meta-srv/src/procedure/wal_prune.rs +++ b/src/meta-srv/src/procedure/wal_prune.rs @@ -25,8 +25,8 @@ use common_meta::lock_key::RemoteWalLock; use common_meta::region_registry::LeaderRegionRegistryRef; use common_procedure::error::ToJsonSnafu; use common_procedure::{ - Context as ProcedureContext, Error as ProcedureError, EventRuntimeContext, EventTrigger, - LockKey, Procedure, ProcedureState, Result as ProcedureResult, Status, StringKey, + Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, + Procedure, ProcedureState, Result as ProcedureResult, Status, StringKey, }; use common_telemetry::{info, warn}; use manager::{WalPruneProcedureGuard, WalPruneProcedureTracker}; @@ -213,10 +213,7 @@ impl Procedure for WalPruneProcedure { LockKey::new(vec![lock_key]) } - fn event( - &self, - ctx: &EventRuntimeContext<'_>, - ) -> Option> { + fn event(&self, ctx: &EventContext<'_>) -> Option> { if !ctx.event_type_filter.allows(WAL_PRUNE_EVENT_TYPE) { return None; } @@ -506,11 +503,12 @@ mod tests { let mut procedure = WalPruneProcedure::new(context, None, "test_topic".to_string(), 42, false); let running = ProcedureState::Running; - let runtime_context = |trigger, lifecycle_state, event_type_filter| EventRuntimeContext { + let runtime_context = |trigger, lifecycle_state, event_type_filter| EventContext { procedure_id: ProcedureId::random(), lifecycle_state, trigger, event_type_filter: Arc::new(event_type_filter), + event_context: None, }; for trigger in [EventTrigger::Submitted, EventTrigger::Recovered] { diff --git a/src/meta-srv/src/service/procedure.rs b/src/meta-srv/src/service/procedure.rs index b20d7f64775..1b5c4c42dbc 100644 --- a/src/meta-srv/src/service/procedure.rs +++ b/src/meta-srv/src/service/procedure.rs @@ -107,6 +107,7 @@ impl procedure_service_server::ProcedureService for Metasrv { .submit_ddl_task( &ExecutorContext { tracing_context: Some(header.tracing_context), + ..Default::default() }, SubmitDdlTaskRequest { query_context, diff --git a/src/standalone/src/procedure.rs b/src/standalone/src/procedure.rs index 85973e9f643..ec63b11703b 100644 --- a/src/standalone/src/procedure.rs +++ b/src/standalone/src/procedure.rs @@ -20,7 +20,6 @@ use common_meta::ddl::DdlContext; use common_meta::ddl_manager::{RepartitionProcedureFactory, RepartitionSource}; use common_meta::key::runtime_switch::RuntimeSwitchManager; use common_meta::kv_backend::KvBackendRef; -use common_meta::rpc::ddl::EventContext; use common_meta::state_store::KvStateStore; use common_procedure::local::{EventRecorderHandle, LocalManager, ManagerConfig}; use common_procedure::options::ProcedureConfig; @@ -74,7 +73,6 @@ impl RepartitionProcedureFactory for StandaloneRepartitionProcedureFactory { _source: RepartitionSource, _to_exprs: Vec, _timeout: Option, - _event_context: EventContext, ) -> std::result::Result { Err(BoxedError::new(NoSupportRepartitionProcedureSnafu.build())) } diff --git a/tests-integration/src/tests/gc.rs b/tests-integration/src/tests/gc.rs index 5bc20e2e2cd..badaef5f7ac 100644 --- a/tests-integration/src/tests/gc.rs +++ b/tests-integration/src/tests/gc.rs @@ -17,7 +17,6 @@ use std::sync::Arc; use std::time::Duration; use common_meta::key::TableMetadataManagerRef; -use common_meta::rpc::ddl::EventContext; use common_procedure::ProcedureWithId; use common_telemetry::info; use common_test_util::recordbatch::check_output_stream; @@ -214,7 +213,6 @@ async fn test_gc_basic(store_type: &StorageType) { false, // full_file_listing Duration::from_secs(10), // timeout Default::default(), - EventContext::default(), ); // Submit the procedure to the procedure manager diff --git a/tests-integration/tests/gc_event.rs b/tests-integration/tests/gc_event.rs index b62ba1d3649..9ecedc77775 100644 --- a/tests-integration/tests/gc_event.rs +++ b/tests-integration/tests/gc_event.rs @@ -15,7 +15,7 @@ use std::sync::Arc; use std::time::Duration; -use common_meta::rpc::ddl::{EventContext, TriggerReason}; +use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; use common_procedure::{ProcedureId, ProcedureWithId, watcher}; use common_test_util::temp_dir::create_temp_dir; use meta_srv::gc::{BatchGcProcedure, GcSchedulerOptions}; @@ -129,7 +129,6 @@ async fn test_batch_gc_event() { false, Duration::from_secs(10), Default::default(), - EventContext::new(TriggerReason::Manual), ); let procedure_id = ProcedureId::parse_str("00000000-0000-0000-0000-00000000bac0").unwrap(); let mut watcher = cluster @@ -138,6 +137,9 @@ async fn test_batch_gc_event() { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), + context: common_procedure::ProcedureContext::from_event_context( + PersistentEventContext::new(TriggerReason::Manual), + ), }) .await .unwrap(); diff --git a/tests-integration/tests/repartition.rs b/tests-integration/tests/repartition.rs index 3dac860066b..1990cc2c43d 100644 --- a/tests-integration/tests/repartition.rs +++ b/tests-integration/tests/repartition.rs @@ -19,7 +19,6 @@ use std::time::Duration; use client::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_error::root_source; use common_meta::key::table_name::TableNameKey; -use common_meta::rpc::ddl::EventContext; use common_procedure::{ProcedureWithId, watcher}; use common_query::Output; use common_telemetry::info; @@ -817,7 +816,6 @@ async fn trigger_table_gc(metasrv: &Arc, table_name: &str) { false, // full_file_listing Duration::from_secs(10), // timeout Default::default(), - EventContext::default(), ); // Submit the procedure to the procedure manager diff --git a/tests-integration/tests/table_ddl_event.rs b/tests-integration/tests/table_ddl_event.rs index 5c23e101de2..5194640880c 100644 --- a/tests-integration/tests/table_ddl_event.rs +++ b/tests-integration/tests/table_ddl_event.rs @@ -377,7 +377,7 @@ async fn test_table_ddl_procedure_events() { .ddl_manager() .submit_undrop_table_task( common_meta::rpc::ddl::UndropTableTask { table_id }, - common_meta::rpc::ddl::EventContext::new( + common_meta::rpc::ddl::PersistentEventContext::new( common_meta::rpc::ddl::TriggerReason::Manual, ), ) @@ -419,7 +419,7 @@ async fn test_table_ddl_procedure_events() { .ddl_manager() .submit_purge_dropped_table_task( common_meta::rpc::ddl::PurgeDroppedTableTask { table_id }, - common_meta::rpc::ddl::EventContext::new( + common_meta::rpc::ddl::PersistentEventContext::new( common_meta::rpc::ddl::TriggerReason::Manual, ), ) diff --git a/tests/conf/frontend-test.toml.template b/tests/conf/frontend-test.toml.template index 25d44ff6e4f..0042716bd80 100644 --- a/tests/conf/frontend-test.toml.template +++ b/tests/conf/frontend-test.toml.template @@ -1,3 +1,6 @@ [grpc] bind_addr = "{addrs.grpc_addr}" server_addr = "{addrs.grpc_addr}" + +[event_recorder] +event_types = [] From 943eee852f292bfffd7240867dc27248216c353f Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 11 Aug 2026 20:50:50 +0800 Subject: [PATCH 28/92] feat(event): record admin function executions (#8835) * feat(event): record admin function executions Signed-off-by: WenyXu * fix(event): handle admin function recording edge cases Signed-off-by: WenyXu * feat(event): record actor for admin functions Signed-off-by: WenyXu * fix(event): preserve admin function event values Signed-off-by: WenyXu * fix(event): preserve non-finite admin results Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- Cargo.lock | 1 + config/config.md | 3 +- config/frontend.example.toml | 4 + config/standalone.example.toml | 2 +- src/cmd/tests/load_config_test.rs | 5 + src/common/event-recorder/src/event_table.rs | 63 ++ src/frontend/src/frontend.rs | 6 +- src/frontend/src/instance/builder.rs | 13 +- src/operator/Cargo.toml | 1 + src/operator/src/error.rs | 4 + src/operator/src/statement.rs | 16 + src/operator/src/statement/admin.rs | 116 +++- src/operator/src/statement/admin/event.rs | 569 ++++++++++++++++++ src/operator/src/statement/admin/layer.rs | 440 ++++++++++++++ .../tests/admin_function_event.rs | 193 ++++++ tests-integration/tests/main.rs | 1 + 16 files changed, 1421 insertions(+), 16 deletions(-) create mode 100644 src/operator/src/statement/admin/event.rs create mode 100644 src/operator/src/statement/admin/layer.rs create mode 100644 tests-integration/tests/admin_function_event.rs diff --git a/Cargo.lock b/Cargo.lock index 18248bd442d..4d0c271882b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9892,6 +9892,7 @@ dependencies = [ "common-catalog", "common-datasource", "common-error", + "common-event-recorder", "common-frontend", "common-function", "common-grpc", diff --git a/config/config.md b/config/config.md index 40272ddcfda..c7737a164b8 100644 --- a/config/config.md +++ b/config/config.md @@ -234,7 +234,7 @@ | `tracing.tokio_console_addr` | String | Unset | The tokio console address. | | `event_recorder` | -- | -- | Configuration options for the event recorder. | | `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. | -| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `create_database`,
`alter_database`, `drop_database`, `create_flow`, `drop_flow`,
`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,
`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,
`create_view`, `drop_view`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | +| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `create_database`,
`alter_database`, `drop_database`, `create_flow`, `drop_flow`,
`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,
`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,
`create_view`, `drop_view`, `admin_function`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | | `memory` | -- | -- | The memory options. | | `memory.enable_heap_profiling` | Bool | `true` | Whether to enable heap profiling activation during startup.
When enabled, heap profiling will be activated if the `MALLOC_CONF` environment variable
is set to "prof:true,prof_active:false". The official image adds this env variable.
Default is true. | @@ -372,6 +372,7 @@ | `memory.enable_heap_profiling` | Bool | `true` | Whether to enable heap profiling activation during startup.
When enabled, heap profiling will be activated if the `MALLOC_CONF` environment variable
is set to "prof:true,prof_active:false". The official image adds this env variable.
Default is true. | | `event_recorder` | -- | -- | Configuration options for the event recorder. | | `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. | +| `event_recorder.event_types` | Array | -- | Event types to record. Current available event type: `admin_function`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | ### Metasrv diff --git a/config/frontend.example.toml b/config/frontend.example.toml index a5907545eba..260984a4ab3 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -405,3 +405,7 @@ enable_heap_profiling = true [event_recorder] ## TTL for the events table that will be used to store the events. Default is `90d`. ttl = "90d" +## Event types to record. Current available event type: `admin_function`. +## When omitted, all current and future event types are recorded. +## Set to an empty array to disable event recording. +#+ event_types = ["admin_function"] diff --git a/config/standalone.example.toml b/config/standalone.example.toml index d256d681f9f..c8dff4944c1 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -922,7 +922,7 @@ ttl = "90d" ## `alter_database`, `drop_database`, `create_flow`, `drop_flow`, ## `create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`, ## `drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`, -## `create_view`, `drop_view`. +## `create_view`, `drop_view`, `admin_function`. ## When omitted, all current and future event types are recorded. ## Set to an empty array to disable event recording. #+ event_types = ["create_database"] diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index aa1d98cf4b9..fd32680f4c6 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -445,6 +445,11 @@ fn test_load_event_types_from_env() { .component .event_recorder .event_types, + GreptimeOptions::::load_layered_options(None, env_prefix) + .unwrap() + .component + .event_recorder + .event_types, ] { assert!(event_types.allows("region_migration")); assert!(!event_types.allows("other_event")); diff --git a/src/common/event-recorder/src/event_table.rs b/src/common/event-recorder/src/event_table.rs index 0ec38c8eabc..1ea4369a00a 100644 --- a/src/common/event-recorder/src/event_table.rs +++ b/src/common/event-recorder/src/event_table.rs @@ -88,6 +88,9 @@ pub const TIMESTAMP_COLUMN: EventTableColumn = EventTableColumn::new( ColumnDataType::TimestampNanosecond, SemanticType::Timestamp, ); +/// The effective database user that initiated an event. +pub const ACTOR_COLUMN: EventTableColumn = + EventTableColumn::new("actor", ColumnDataType::String, SemanticType::Field); /// The canonical procedure identifier envelope column. pub const PROCEDURE_ID_COLUMN: EventTableColumn = EventTableColumn::new("procedure_id", ColumnDataType::String, SemanticType::Field); @@ -144,6 +147,25 @@ pub const LATEST_OFFSET_COLUMN: EventTableColumn = pub const GC_REPORT_COLUMN: EventTableColumn = EventTableColumn::json_binary("gc_report", ColumnDataType::Binary, SemanticType::Field); +/// The canonical ADMIN function name field. +pub const ADMIN_FUNCTION_NAME_COLUMN: EventTableColumn = EventTableColumn::new( + "admin_function_name", + ColumnDataType::String, + SemanticType::Field, +); +/// The canonical ADMIN function execution status field. +pub const ADMIN_FUNCTION_STATUS_COLUMN: EventTableColumn = EventTableColumn::new( + "admin_function_status", + ColumnDataType::String, + SemanticType::Field, +); +/// The canonical ADMIN function output field. +pub const ADMIN_FUNCTION_OUTPUT_COLUMN: EventTableColumn = EventTableColumn::json_binary( + "admin_function_output", + ColumnDataType::Binary, + SemanticType::Field, +); + /// The canonical table name field for table DDL events. pub const TABLE_NAME_COLUMN: EventTableColumn = EventTableColumn::new("table_name", ColumnDataType::String, SemanticType::Field); @@ -389,6 +411,47 @@ mod tests { ); } + #[test] + fn admin_function_schema_preserves_names_types_semantics_extensions_and_order() { + assert_eq!( + column_schemas([ + &ACTOR_COLUMN, + &ADMIN_FUNCTION_NAME_COLUMN, + &ADMIN_FUNCTION_STATUS_COLUMN, + &ADMIN_FUNCTION_OUTPUT_COLUMN, + ]), + vec![ + ColumnSchema { + column_name: "actor".to_string(), + datatype: ColumnDataType::String.into(), + semantic_type: SemanticType::Field.into(), + ..Default::default() + }, + ColumnSchema { + column_name: "admin_function_name".to_string(), + datatype: ColumnDataType::String.into(), + semantic_type: SemanticType::Field.into(), + ..Default::default() + }, + ColumnSchema { + column_name: "admin_function_status".to_string(), + datatype: ColumnDataType::String.into(), + semantic_type: SemanticType::Field.into(), + ..Default::default() + }, + ColumnSchema { + column_name: "admin_function_output".to_string(), + datatype: ColumnDataType::Binary.into(), + semantic_type: SemanticType::Field.into(), + datatype_extension: Some(ColumnDataTypeExtension { + type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())), + }), + ..Default::default() + }, + ] + ); + } + #[test] fn table_dimension_schema_preserves_names_types_semantics_and_order() { assert_eq!( diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index 7a09b40e736..b8fea7b6043 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -116,7 +116,11 @@ impl Default for FrontendOptions { impl Configurable for FrontendOptions { fn env_list_keys() -> Option<&'static [&'static str]> { - Some(&["heartbeat_env_vars", "meta_client.metasrv_addrs"]) + Some(&[ + "heartbeat_env_vars", + "meta_client.metasrv_addrs", + "event_recorder.event_types", + ]) } } diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index 85a2b8ff183..c49a8394beb 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -23,7 +23,7 @@ use catalog::process_manager::ProcessManagerRef; use catalog::system_schema::semantic_graph::EntityGraphProviderRef; use common_base::Plugins; use common_datasource::object_store::LocalFileAccess; -use common_event_recorder::EventRecorderImpl; +use common_event_recorder::{EventRecorderImpl, EventRecorderRef}; use common_meta::cache::{LayeredCacheRegistryRef, TableRouteCacheRef}; use common_meta::cache_invalidator::{CacheInvalidatorRef, DummyCacheInvalidator}; use common_meta::key::TableMetadataManager; @@ -40,8 +40,8 @@ use operator::request::Requester; #[cfg(feature = "enterprise")] use operator::statement::CreateDatabaseHandlerRef; use operator::statement::{ - ExecutorConfigureContext, StatementExecutor, StatementExecutorConfiguratorRef, - StatementExecutorRef, + AdminEventRecorderHandle, AdminFunctionRecordingLayer, ExecutorConfigureContext, + StatementExecutor, StatementExecutorConfiguratorRef, StatementExecutorRef, }; use operator::table::TableMutationOperator; use partition::cache::PartitionInfoCacheRef; @@ -323,6 +323,10 @@ impl FrontendBuilder { statement_executor }; + let admin_event_recorder = AdminEventRecorderHandle::default(); + let statement_executor = statement_executor.with_admin_function_layer(Arc::new( + AdminFunctionRecordingLayer::new(admin_event_recorder.clone()), + )); let statement_executor = Arc::new(statement_executor); let pipeline_operator = Arc::new(PipelineOperator::new( @@ -337,13 +341,14 @@ impl FrontendBuilder { let slow_query_recorder = Arc::new(EventRecorderImpl::new(Box::new( EventHandlerImpl::new(statement_executor.clone(), self.options.slow_query.ttl), ))); - let event_recorder = Arc::new(EventRecorderImpl::with_event_type_filter( + let event_recorder: EventRecorderRef = Arc::new(EventRecorderImpl::with_event_type_filter( Box::new(EventHandlerImpl::new( statement_executor.clone(), self.options.event_recorder.ttl, )), self.options.event_recorder.event_types.clone(), )); + admin_event_recorder.install(&event_recorder); Ok(Instance { frontend_peer_addr, diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index bd34e21ed24..087ea461f10 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -33,6 +33,7 @@ common-base.workspace = true common-catalog.workspace = true common-datasource.workspace = true common-error.workspace = true +common-event-recorder.workspace = true common-frontend.workspace = true common-function.workspace = true common-grpc.workspace = true diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs index b5e76f0f2cc..f1e94b2c616 100644 --- a/src/operator/src/error.rs +++ b/src/operator/src/error.rs @@ -65,6 +65,9 @@ pub enum Error { location: Location, }, + #[snafu(display("Admin function execution was cancelled"))] + AdminFunctionCancelled, + #[snafu(display("Expected {expected} args, but actual {actual}"))] FunctionArityMismatch { expected: usize, actual: usize }, @@ -1023,6 +1026,7 @@ impl ErrorExt for Error { | Error::EncodeJson { .. } | Error::DeserializePartitionExpr { .. } | Error::SerializePartitionExpr { .. } => StatusCode::Unexpected, + Error::AdminFunctionCancelled => StatusCode::Cancelled, Error::ViewNotFound { .. } | Error::ViewInfoNotFound { .. } | Error::TableNotFound { .. } => StatusCode::TableNotFound, diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index fea54bb27c3..478f06e83f7 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -81,6 +81,11 @@ use table::requests::{CopyDatabaseRequest, CopyDirection, CopyQueryToRequest, Co use table::table_name::TableName; use table::table_reference::TableReference; +pub use self::admin::{ + AdminEventRecorderHandle, AdminFunctionLayer, AdminFunctionLayerRef, + AdminFunctionRecordingLayer, AdminFunctionRequest, AdminFunctionResponse, AdminFunctionService, + AdminFunctionServiceRef, +}; use self::set::{ set_bytea_output, set_datestyle, set_intervalstyle, set_timezone, validate_client_encoding, }; @@ -139,6 +144,7 @@ pub struct StatementExecutor { inserter: InserterRef, process_manager: Option, origin_frontend_addr: String, + admin_function_service: AdminFunctionServiceRef, pub(crate) local_file_access: LocalFileAccess, #[cfg(feature = "enterprise")] create_database_handler: Option, @@ -180,6 +186,7 @@ impl StatementExecutor { origin_frontend_addr: String, local_file_access: LocalFileAccess, ) -> Self { + let admin_function_service = admin::new_admin_function_service(query_engine.clone()); Self { catalog_manager, query_engine, @@ -192,6 +199,7 @@ impl StatementExecutor { inserter, process_manager, origin_frontend_addr, + admin_function_service, local_file_access, #[cfg(feature = "enterprise")] create_database_handler: None, @@ -200,6 +208,14 @@ impl StatementExecutor { } } + /// Adds a layer around the ADMIN function execution service. + /// + /// The last added layer is the outermost layer. + pub fn with_admin_function_layer(mut self, layer: AdminFunctionLayerRef) -> Self { + self.admin_function_service = layer.layer(self.admin_function_service); + self + } + #[cfg(feature = "enterprise")] pub fn with_trigger_querier(mut self, querier: TriggerQuerierRef) -> Self { self.trigger_querier = Some(querier); diff --git a/src/operator/src/statement/admin.rs b/src/operator/src/statement/admin.rs index 7494977b214..2b4a909e3cd 100644 --- a/src/operator/src/statement/admin.rs +++ b/src/operator/src/statement/admin.rs @@ -12,6 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod event; +mod layer; + use std::sync::Arc; use common_function::function::FunctionContext; @@ -28,6 +31,10 @@ use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::value::Value; use datatypes::vectors::VectorRef; +pub use layer::{ + AdminEventRecorderHandle, AdminFunctionLayer, AdminFunctionLayerRef, + AdminFunctionRecordingLayer, +}; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt, ensure}; use sql::ast::{Expr, FunctionArg, FunctionArgExpr, FunctionArguments, Value as SqlValue}; @@ -38,14 +45,49 @@ use crate::statement::StatementExecutor; const DUMMY_COLUMN: &str = ""; -impl StatementExecutor { - /// Execute the [`Admin`] statement and returns the output. - #[tracing::instrument(skip_all)] - pub(super) async fn execute_admin_command( - &self, - stmt: Admin, - query_ctx: QueryContextRef, - ) -> Result { +/// A request to execute one ADMIN function statement. +#[derive(Clone)] +pub struct AdminFunctionRequest { + /// The parsed ADMIN statement to execute. + pub statement: Admin, + /// The query context of the request. + pub query_ctx: QueryContextRef, +} + +/// The client output and immediate typed result of one ADMIN function execution. +pub struct AdminFunctionResponse { + /// The output returned to the client. + pub output: Output, + /// The typed immediate result exposed to outer layers. + pub immediate_result: Option, +} + +/// Executes an ADMIN function request. +#[async_trait::async_trait] +pub trait AdminFunctionService: Send + Sync { + /// Executes an ADMIN function request. + async fn call(&self, request: AdminFunctionRequest) -> Result; +} + +/// A shared ADMIN function service. +pub type AdminFunctionServiceRef = Arc; + +#[derive(Clone)] +struct CoreAdminFunctionService { + query_engine: query::QueryEngineRef, +} + +impl CoreAdminFunctionService { + fn new(query_engine: query::QueryEngineRef) -> Self { + Self { query_engine } + } + + async fn execute(&self, request: AdminFunctionRequest) -> Result { + let AdminFunctionRequest { + statement: stmt, + query_ctx, + } = request; + let Admin::Func(func) = &stmt; // the function name should be in lower case. let func_name = func.name.to_string().to_lowercase(); @@ -141,6 +183,7 @@ impl StatementExecutor { (&result_columnar).try_into().context(CastSnafu)?; let result_vector: VectorRef = result_columnar.try_into_vector(1).context(CastSnafu)?; + let immediate_result = immediate_result(&result_vector); let column_schemas = vec![ColumnSchema::new( // Use statement as the result column name @@ -154,7 +197,46 @@ impl StatementExecutor { let batches = RecordBatches::try_new(schema, vec![batch]).context(error::BuildRecordBatchSnafu)?; - Ok(Output::new_with_record_batches(batches)) + Ok(AdminFunctionResponse { + output: Output::new_with_record_batches(batches), + immediate_result, + }) + } +} + +fn immediate_result(result_vector: &VectorRef) -> Option { + (!result_vector.is_empty()).then(|| result_vector.get(0)) +} + +#[async_trait::async_trait] +impl AdminFunctionService for CoreAdminFunctionService { + async fn call(&self, request: AdminFunctionRequest) -> Result { + self.execute(request).await + } +} + +/// Creates the core ADMIN function service. +pub(crate) fn new_admin_function_service( + query_engine: query::QueryEngineRef, +) -> AdminFunctionServiceRef { + Arc::new(CoreAdminFunctionService::new(query_engine)) +} + +impl StatementExecutor { + /// Executes the [`Admin`] statement and returns the output. + #[tracing::instrument(skip_all)] + pub(crate) async fn execute_admin_command( + &self, + stmt: Admin, + query_ctx: QueryContextRef, + ) -> Result { + self.admin_function_service + .call(AdminFunctionRequest { + statement: stmt, + query_ctx, + }) + .await + .map(|response| response.output) } } @@ -306,3 +388,19 @@ fn try_get_data_type_for_sql_value(value: &SqlValue) -> Result { .fail(), } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datatypes::vectors::{Int32Vector, VectorRef}; + + use crate::statement::admin::immediate_result; + + #[test] + fn empty_admin_function_result_has_no_immediate_value() { + let result_vector: VectorRef = Arc::new(Int32Vector::from(vec![])); + + assert_eq!(immediate_result(&result_vector), None); + } +} diff --git a/src/operator/src/statement/admin/event.rs b/src/operator/src/statement/admin/event.rs new file mode 100644 index 00000000000..bafa2ad8824 --- /dev/null +++ b/src/operator/src/statement/admin/event.rs @@ -0,0 +1,569 @@ +// Copyright 2023 Greptime Team +// +// 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. + +use std::any::Any; + +use api::v1::value::ValueData; +use api::v1::{ColumnSchema, Row}; +use common_event_recorder::Event; +use common_event_recorder::error::Result as EventResult; +use common_event_recorder::event_table::{ + ACTOR_COLUMN, ADMIN_FUNCTION_NAME_COLUMN, ADMIN_FUNCTION_OUTPUT_COLUMN, + ADMIN_FUNCTION_STATUS_COLUMN, column_schemas, jsonb_value, +}; +use datatypes::value::Value; +use serde_json::{Value as JsonValue, json}; +use sql::ast::{Expr, FunctionArg, FunctionArgExpr, FunctionArguments, Value as SqlValue}; +use sql::statements::admin::Admin; + +use crate::error::Error; +use crate::statement::admin::AdminFunctionRequest; + +/// Event type emitted for an ADMIN function execution. +pub(crate) const ADMIN_FUNCTION_EVENT_TYPE: &str = "admin_function"; + +const PAYLOAD_VERSION: u8 = 1; +const VERSION_KEY: &str = "version"; +const ARGUMENTS_KEY: &str = "arguments"; +const ARGUMENT_TYPE_KEY: &str = "type"; +const ARGUMENT_VALUE_KEY: &str = "value"; +const LIST_ARGUMENT_TYPE: &str = "list"; +const SUBQUERY_ARGUMENT_TYPE: &str = "subquery"; +const RESULT_KEY: &str = "result"; +const ERROR_KEY: &str = "error"; +const SUCCEEDED_STATUS: &str = "Succeeded"; +const FAILED_STATUS: &str = "Failed"; +const UNSUPPORTED: &str = ""; + +/// The ADMIN function input captured before execution. +#[derive(Debug, Clone)] +pub(crate) struct AdminFunctionEventInput { + actor: String, + function: String, + arguments: Option, +} + +impl AdminFunctionEventInput { + /// Creates an event input from an ADMIN function request. + pub(crate) fn from_request(request: &AdminFunctionRequest) -> Self { + let Admin::Func(function) = &request.statement; + let function_name = function.name.to_string().to_lowercase(); + + let arguments = match &function.args { + FunctionArguments::List(arguments) => Some(json!({ + (ARGUMENT_TYPE_KEY): LIST_ARGUMENT_TYPE, + (ARGUMENT_VALUE_KEY): arguments.args.iter().map(argument_to_json).collect::>(), + })), + FunctionArguments::Subquery(query) => Some(json!({ + (ARGUMENT_TYPE_KEY): SUBQUERY_ARGUMENT_TYPE, + (ARGUMENT_VALUE_KEY): query.to_string(), + })), + FunctionArguments::None => None, + }; + + Self { + actor: request.query_ctx.current_user().username().to_string(), + function: function_name, + arguments, + } + } +} + +/// An event describing the outcome of an ADMIN function execution. +#[derive(Debug)] +pub(crate) struct AdminFunctionEvent { + actor: String, + function_name: String, + status: &'static str, + output: JsonValue, + payload: JsonValue, +} + +impl AdminFunctionEvent { + /// Creates a successful ADMIN function event. + pub(crate) fn success(input: AdminFunctionEventInput, result: Option<&Value>) -> Self { + let AdminFunctionEventInput { + actor, + function, + arguments, + } = input; + Self { + actor, + function_name: function, + status: SUCCEEDED_STATUS, + output: result.map_or_else( + || json!({}), + |result| { + json!({ + (RESULT_KEY): value_to_json(result), + }) + }, + ), + payload: input_payload(arguments), + } + } + + /// Creates a failed ADMIN function event with the debug representation of the error. + pub(crate) fn failure(input: AdminFunctionEventInput, error: &Error) -> Self { + let AdminFunctionEventInput { + actor, + function, + arguments, + } = input; + Self { + actor, + function_name: function, + status: FAILED_STATUS, + output: json!({ + (ERROR_KEY): format!("{error:?}"), + }), + payload: input_payload(arguments), + } + } + + /// Creates a failed ADMIN function event for a cancelled execution. + pub(crate) fn cancelled(input: AdminFunctionEventInput) -> Self { + Self::failure(input, &Error::AdminFunctionCancelled) + } +} + +impl Event for AdminFunctionEvent { + fn event_type(&self) -> &str { + ADMIN_FUNCTION_EVENT_TYPE + } + + fn json_payload(&self) -> EventResult { + Ok(self.payload.clone()) + } + + fn extra_schema(&self) -> Vec { + column_schemas([ + &ACTOR_COLUMN, + &ADMIN_FUNCTION_NAME_COLUMN, + &ADMIN_FUNCTION_STATUS_COLUMN, + &ADMIN_FUNCTION_OUTPUT_COLUMN, + ]) + } + + fn extra_rows(&self) -> EventResult> { + Ok(vec![Row { + values: vec![ + ValueData::StringValue(self.actor.clone()).into(), + ValueData::StringValue(self.function_name.clone()).into(), + ValueData::StringValue(self.status.to_string()).into(), + jsonb_value(&self.output), + ], + }]) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +fn input_payload(arguments: Option) -> JsonValue { + let mut payload = json!({ VERSION_KEY: PAYLOAD_VERSION }); + if let Some(arguments) = arguments { + payload[ARGUMENTS_KEY] = arguments; + } + payload +} + +fn argument_to_json(argument: &FunctionArg) -> JsonValue { + match argument { + FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(value))) => { + sql_value_to_json(&value.value) + } + _ => JsonValue::String(argument.to_string()), + } +} + +fn sql_value_to_json(value: &SqlValue) -> JsonValue { + match value { + SqlValue::Number(value, _) => { + serde_json::from_str(value).unwrap_or_else(|_| JsonValue::String(value.clone())) + } + SqlValue::Boolean(value) => JsonValue::Bool(*value), + SqlValue::Null => JsonValue::Null, + SqlValue::SingleQuotedString(value) | SqlValue::DoubleQuotedString(value) => { + JsonValue::String(value.clone()) + } + SqlValue::Placeholder(value) => JsonValue::String(value.clone()), + _ => value + .clone() + .into_string() + .map(JsonValue::String) + .unwrap_or_else(|| JsonValue::String(value.to_string())), + } +} + +fn value_to_json(value: &Value) -> JsonValue { + match value { + Value::Float32(value) if !value.0.is_finite() => JsonValue::String(value.to_string()), + Value::Float64(value) if !value.0.is_finite() => JsonValue::String(value.to_string()), + _ => JsonValue::try_from(value.clone()) + .unwrap_or_else(|_| JsonValue::String(UNSUPPORTED.to_string())), + } +} + +#[cfg(test)] +mod tests { + use api::v1::Row; + use api::v1::value::ValueData; + use common_event_recorder::Event; + use common_event_recorder::event_table::{ + ACTOR_COLUMN, ADMIN_FUNCTION_NAME_COLUMN, ADMIN_FUNCTION_OUTPUT_COLUMN, + ADMIN_FUNCTION_STATUS_COLUMN, column_schemas, jsonb_value, + }; + use datatypes::value::Value; + use serde_json::json; + use session::context::QueryContext; + use sql::ast::{ + Expr, FunctionArg, FunctionArgExpr, FunctionArguments, Ident, Value as SqlValue, + }; + use sql::dialect::GreptimeDbDialect; + use sql::parser::{ParseOptions, ParserContext}; + use sql::statements::admin::Admin; + use sql::statements::statement::Statement; + use sqlparser::ast::{DollarQuotedString, FunctionArgOperator, FunctionArgumentList}; + + use crate::error::Error; + use crate::statement::admin::AdminFunctionRequest; + use crate::statement::admin::event::{ + ADMIN_FUNCTION_EVENT_TYPE, AdminFunctionEvent, AdminFunctionEventInput, sql_value_to_json, + value_to_json, + }; + + fn request(sql: &str) -> AdminFunctionRequest { + let Statement::Admin(statement) = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap() + .remove(0) + else { + panic!("expected ADMIN statement") + }; + AdminFunctionRequest { + statement, + query_ctx: QueryContext::arc(), + } + } + + fn request_with_arguments(arguments: FunctionArguments) -> AdminFunctionRequest { + let mut request = request("ADMIN plugin_function()"); + let Admin::Func(function) = &mut request.statement; + function.args = arguments; + request + } + + fn assert_admin_columns( + event: &AdminFunctionEvent, + function_name: &str, + status: &str, + output: serde_json::Value, + ) { + assert_eq!(event.event_type(), ADMIN_FUNCTION_EVENT_TYPE); + assert_eq!( + event.extra_schema(), + column_schemas([ + &ACTOR_COLUMN, + &ADMIN_FUNCTION_NAME_COLUMN, + &ADMIN_FUNCTION_STATUS_COLUMN, + &ADMIN_FUNCTION_OUTPUT_COLUMN, + ]) + ); + assert_eq!( + event.extra_rows().unwrap(), + vec![Row { + values: vec![ + ValueData::StringValue("greptime".to_string()).into(), + ValueData::StringValue(function_name.to_string()).into(), + ValueData::StringValue(status.to_string()).into(), + jsonb_value(&output), + ], + }] + ); + } + + #[test] + fn success_uses_separate_output_columns() { + let input = AdminFunctionEventInput::from_request(&request( + "ADMIN flush_table('greptime.public.demo')", + )); + let event = AdminFunctionEvent::success(input, Some(&Value::UInt64(3))); + + assert_admin_columns(&event, "flush_table", "Succeeded", json!({"result": 3})); + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "list", "value": ["greptime.public.demo"]}, + }) + ); + } + + #[test] + fn successful_null_is_explicit() { + let input = AdminFunctionEventInput::from_request(&request( + "ADMIN migrate_region(NULL, NULL, NULL)", + )); + let event = AdminFunctionEvent::success(input, Some(&Value::Null)); + + assert_admin_columns( + &event, + "migrate_region", + "Succeeded", + json!({"result": null}), + ); + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "list", "value": [null, null, null]}, + }) + ); + } + + #[test] + fn unknown_function_values_are_recorded() { + let input = AdminFunctionEventInput::from_request(&request( + "ADMIN plugin_function('plugin-value', NULL)", + )); + let event = AdminFunctionEvent::success(input, Some(&Value::from("plugin-result"))); + + assert_admin_columns( + &event, + "plugin_function", + "Succeeded", + json!({"result": "plugin-result"}), + ); + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "list", "value": ["plugin-value", null]}, + }) + ); + } + + #[test] + fn failure_payload_has_error_only() { + let input = AdminFunctionEventInput::from_request(&request( + "ADMIN flush_table('greptime.public.demo')", + )); + let error = Error::BuildAdminFunctionArgs { + msg: "invalid input".to_string(), + }; + let expected_error = format!("{error:?}"); + let event = AdminFunctionEvent::failure(input, &error); + + assert_admin_columns( + &event, + "flush_table", + "Failed", + json!({"error": expected_error}), + ); + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "list", "value": ["greptime.public.demo"]}, + }) + ); + } + + #[test] + fn cancellation_uses_debug_error_format() { + let input = AdminFunctionEventInput::from_request(&request("ADMIN flush_table('demo')")); + let event = AdminFunctionEvent::cancelled(input); + + assert_admin_columns( + &event, + "flush_table", + "Failed", + json!({"error": format!("{:?}", Error::AdminFunctionCancelled)}), + ); + } + + #[test] + fn subquery_arguments_use_display() { + let Statement::Query(query) = ParserContext::create_with_dialect( + "SELECT 1", + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap() + .remove(0) else { + panic!("expected query statement") + }; + let input = AdminFunctionEventInput::from_request(&request_with_arguments( + FunctionArguments::Subquery(Box::new(query.inner)), + )); + let event = AdminFunctionEvent::failure( + input, + &Error::BuildAdminFunctionArgs { + msg: "subquery is not executable".to_string(), + }, + ); + + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "subquery", "value": "SELECT 1"}, + }) + ); + } + + #[test] + fn no_arguments_omits_arguments_field_and_empty_result() { + let input = + AdminFunctionEventInput::from_request(&request_with_arguments(FunctionArguments::None)); + let event = AdminFunctionEvent::success(input, None); + + assert_admin_columns(&event, "plugin_function", "Succeeded", json!({})); + assert_eq!(event.json_payload().unwrap(), json!({"version": 1})); + } + + #[test] + fn empty_argument_list_is_recorded() { + let input = AdminFunctionEventInput::from_request(&request("ADMIN plugin_function()")); + let event = AdminFunctionEvent::success(input, Some(&Value::UInt64(0))); + + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "list", "value": []}, + }) + ); + } + + #[test] + fn records_extended_sql_values_and_non_literal_arguments() { + let input = AdminFunctionEventInput::from_request(&request( + "ADMIN plugin_function(1, true, NULL, X'48656c6c6f', table_name)", + )); + let event = AdminFunctionEvent::success(input, Some(&Value::UInt64(0))); + + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": { + "type": "list", + "value": [1, true, null, "48656c6c6f", "table_name"], + }, + }) + ); + } + + #[test] + fn string_literals_preserve_logical_content() { + let values = [ + SqlValue::SingleQuotedString("single".to_string()), + SqlValue::DoubleQuotedString("double".to_string()), + SqlValue::DollarQuotedString(DollarQuotedString { + value: "dollar".to_string(), + tag: None, + }), + SqlValue::SingleQuotedRawStringLiteral("raw".to_string()), + SqlValue::SingleQuotedByteStringLiteral("bytes".to_string()), + SqlValue::NationalStringLiteral("national".to_string()), + SqlValue::UnicodeStringLiteral("unicode".to_string()), + SqlValue::HexStringLiteral("48656c6c6f".to_string()), + ]; + + let expected = [ + "single", + "double", + "dollar", + "raw", + "bytes", + "national", + "unicode", + "48656c6c6f", + ]; + for (value, expected) in values.iter().zip(expected) { + assert_eq!(sql_value_to_json(value), json!(expected)); + } + } + + #[test] + fn non_literal_arguments_use_sql_display() { + let arguments = FunctionArgumentList { + duplicate_treatment: None, + args: vec![ + FunctionArg::Named { + name: Ident::new("named"), + arg: FunctionArgExpr::Expr(Expr::Identifier(Ident::new("value"))), + operator: FunctionArgOperator::RightArrow, + }, + FunctionArg::Unnamed(FunctionArgExpr::Wildcard), + ], + clauses: vec![], + }; + let input = AdminFunctionEventInput::from_request(&request_with_arguments( + FunctionArguments::List(arguments), + )); + let event = AdminFunctionEvent::success(input, Some(&Value::UInt64(0))); + + assert_eq!( + event.json_payload().unwrap(), + json!({ + "version": 1, + "arguments": {"type": "list", "value": ["named => value", "*"]}, + }) + ); + assert_eq!( + sql_value_to_json(&SqlValue::Placeholder("$1".to_string())), + json!("$1") + ); + } + + #[test] + fn serializes_binary_result() { + assert_eq!( + value_to_json(&Value::from(vec![1_u8, 2, 3])), + json!([1, 2, 3]) + ); + } + + #[test] + fn serializes_non_finite_float_results_as_strings() { + assert_eq!( + value_to_json(&Value::Float32(f32::NAN.into())), + json!("NaN") + ); + assert_eq!( + value_to_json(&Value::Float32(f32::INFINITY.into())), + json!("inf") + ); + assert_eq!( + value_to_json(&Value::Float32(f32::NEG_INFINITY.into())), + json!("-inf") + ); + assert_eq!( + value_to_json(&Value::Float64(f64::NAN.into())), + json!("NaN") + ); + assert_eq!( + value_to_json(&Value::Float64(f64::INFINITY.into())), + json!("inf") + ); + assert_eq!( + value_to_json(&Value::Float64(f64::NEG_INFINITY.into())), + json!("-inf") + ); + } +} diff --git a/src/operator/src/statement/admin/layer.rs b/src/operator/src/statement/admin/layer.rs new file mode 100644 index 00000000000..195d47d48aa --- /dev/null +++ b/src/operator/src/statement/admin/layer.rs @@ -0,0 +1,440 @@ +// Copyright 2023 Greptime Team +// +// 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. + +use std::sync::{Arc, RwLock, Weak}; + +use common_event_recorder::{EventRecorder, EventRecorderRef}; + +use crate::error::Result; +use crate::statement::admin::event::{ + ADMIN_FUNCTION_EVENT_TYPE, AdminFunctionEvent, AdminFunctionEventInput, +}; +use crate::statement::admin::{ + AdminFunctionRequest, AdminFunctionResponse, AdminFunctionService, AdminFunctionServiceRef, +}; + +/// A composable layer over ADMIN function execution. +pub trait AdminFunctionLayer: Send + Sync { + /// Wraps `inner` and returns the resulting ADMIN function service. + fn layer(&self, inner: AdminFunctionServiceRef) -> AdminFunctionServiceRef; +} + +/// A shared composable ADMIN function layer. +pub type AdminFunctionLayerRef = Arc; + +/// A late-bound weak reference to the event recorder. +/// +/// The weak reference avoids retaining the frontend event handler, which owns a +/// statement executor itself. +#[derive(Clone, Default)] +pub struct AdminEventRecorderHandle { + recorder: Arc>>>, +} + +impl AdminEventRecorderHandle { + /// Installs the event recorder as a weak reference. + pub fn install(&self, recorder: &EventRecorderRef) { + *self.recorder.write().unwrap_or_else(|e| e.into_inner()) = Some(Arc::downgrade(recorder)); + } + + fn upgrade(&self) -> Option { + self.recorder + .read() + .unwrap_or_else(|e| e.into_inner()) + .as_ref() + .and_then(Weak::upgrade) + } +} + +/// A layer that records ADMIN function inputs and outcomes as events. +#[derive(Clone)] +pub struct AdminFunctionRecordingLayer { + recorder: AdminEventRecorderHandle, +} + +impl AdminFunctionRecordingLayer { + /// Creates a recording layer backed by the late-bound recorder handle. + pub fn new(recorder: AdminEventRecorderHandle) -> Self { + Self { recorder } + } +} + +impl AdminFunctionLayer for AdminFunctionRecordingLayer { + fn layer(&self, inner: AdminFunctionServiceRef) -> AdminFunctionServiceRef { + Arc::new(AdminFunctionRecordingService { + inner, + recorder: self.recorder.clone(), + }) + } +} + +#[derive(Clone)] +struct AdminFunctionRecordingService { + inner: AdminFunctionServiceRef, + recorder: AdminEventRecorderHandle, +} + +#[async_trait::async_trait] +impl AdminFunctionService for AdminFunctionRecordingService { + async fn call(&self, request: AdminFunctionRequest) -> Result { + // Avoid building the event input when this event type is disabled. Do + // not retain the recorder across ADMIN execution because its handler can + // own the statement executor that contains this layer. + let enabled = self.recorder.upgrade().is_some_and(|recorder| { + recorder + .event_type_filter() + .allows(ADMIN_FUNCTION_EVENT_TYPE) + }); + let mut recording = AdminFunctionRecordingGuard::new( + enabled.then(|| AdminFunctionEventInput::from_request(&request)), + self.recorder.clone(), + ); + let result = self.inner.call(request).await; + recording.record_result(&result); + result + } +} + +/// Records an ADMIN event after completion or when an in-flight call is dropped. +struct AdminFunctionRecordingGuard { + input: Option, + recorder: AdminEventRecorderHandle, +} + +impl AdminFunctionRecordingGuard { + fn new(input: Option, recorder: AdminEventRecorderHandle) -> Self { + Self { input, recorder } + } + + fn record_result(&mut self, result: &Result) { + let Some(input) = self.input.take() else { + return; + }; + let Some(recorder) = self.recorder.upgrade() else { + return; + }; + let event = match result { + Ok(response) => AdminFunctionEvent::success(input, response.immediate_result.as_ref()), + Err(error) => AdminFunctionEvent::failure(input, error), + }; + recorder.record(Box::new(event)); + } +} + +impl Drop for AdminFunctionRecordingGuard { + fn drop(&mut self) { + let Some(input) = self.input.take() else { + return; + }; + let Some(recorder) = self.recorder.upgrade() else { + return; + }; + recorder.record(Box::new(AdminFunctionEvent::cancelled(input))); + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + use std::fmt::{Debug, Formatter}; + use std::sync::{Arc, Mutex}; + use std::time::Duration; + + use common_event_recorder::event_table::jsonb_value; + use common_event_recorder::{ + Event, EventRecorder, EventRecorderRef, EventTypeFilter, EventTypeFilterRef, + }; + use common_query::{Output, OutputData}; + use datatypes::value::Value; + use serde_json::json; + use session::context::QueryContext; + use sql::dialect::GreptimeDbDialect; + use sql::parser::{ParseOptions, ParserContext}; + use sql::statements::statement::Statement; + use tokio::sync::Notify; + + use crate::error::Result; + use crate::statement::admin::layer::{AdminEventRecorderHandle, AdminFunctionRecordingLayer}; + use crate::statement::admin::{ + AdminFunctionLayer, AdminFunctionRequest, AdminFunctionResponse, AdminFunctionService, + AdminFunctionServiceRef, + }; + + #[derive(Clone)] + struct TraceLayer { + trace: Arc>>, + before: &'static str, + after: &'static str, + } + + impl AdminFunctionLayer for TraceLayer { + fn layer(&self, inner: AdminFunctionServiceRef) -> AdminFunctionServiceRef { + Arc::new(TraceService { + inner, + trace: self.trace.clone(), + before: self.before, + after: self.after, + }) + } + } + + struct TraceService { + inner: AdminFunctionServiceRef, + trace: Arc>>, + before: &'static str, + after: &'static str, + } + + #[async_trait::async_trait] + impl AdminFunctionService for TraceService { + async fn call(&self, request: AdminFunctionRequest) -> Result { + self.trace.lock().unwrap().push(self.before); + let result = self.inner.call(request).await; + self.trace.lock().unwrap().push(self.after); + result + } + } + + struct TestService { + trace: Option>>>, + immediate_result: Option, + } + + #[async_trait::async_trait] + impl AdminFunctionService for TestService { + async fn call(&self, _request: AdminFunctionRequest) -> Result { + if let Some(trace) = &self.trace { + trace.lock().unwrap().push("core"); + } + Ok(response(self.immediate_result.clone())) + } + } + + struct PendingService { + started: Arc, + } + + #[async_trait::async_trait] + impl AdminFunctionService for PendingService { + async fn call(&self, _request: AdminFunctionRequest) -> Result { + self.started.notify_one(); + std::future::pending().await + } + } + + struct TestRecorder { + events: Mutex>>, + filter: EventTypeFilterRef, + } + + impl TestRecorder { + fn new(filter: EventTypeFilter) -> Self { + Self { + events: Mutex::new(Vec::new()), + filter: Arc::new(filter), + } + } + } + + impl Debug for TestRecorder { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TestRecorder").finish() + } + } + + impl EventRecorder for TestRecorder { + fn record(&self, event: Box) { + self.events.lock().unwrap().push(event); + } + + fn event_type_filter(&self) -> EventTypeFilterRef { + self.filter.clone() + } + + fn close(&self) {} + } + + #[tokio::test] + async fn last_added_layer_runs_outermost() { + let trace = Arc::new(Mutex::new(Vec::new())); + let mut service: AdminFunctionServiceRef = Arc::new(TestService { + trace: Some(trace.clone()), + immediate_result: Some(Value::UInt64(0)), + }); + for (before, after) in [("a_before", "a_after"), ("b_before", "b_after")] { + service = TraceLayer { + trace: trace.clone(), + before, + after, + } + .layer(service); + } + + service.call(request()).await.unwrap(); + assert_eq!( + *trace.lock().unwrap(), + ["b_before", "a_before", "core", "a_after", "b_after"] + ); + } + + #[tokio::test] + async fn recording_respects_filter_and_weak_lifetime() { + let enabled = Arc::new(TestRecorder::new(EventTypeFilter::All)); + let enabled_ref: EventRecorderRef = enabled.clone(); + let handle = AdminEventRecorderHandle::default(); + handle.install(&enabled_ref); + recording_service(handle.clone()) + .call(request()) + .await + .unwrap(); + assert_eq!(enabled.events.lock().unwrap().len(), 1); + + let disabled = Arc::new(TestRecorder::new(EventTypeFilter::Only(HashSet::new()))); + let disabled_ref: EventRecorderRef = disabled.clone(); + handle.install(&disabled_ref); + recording_service(handle.clone()) + .call(request()) + .await + .unwrap(); + assert!(disabled.events.lock().unwrap().is_empty()); + + drop(disabled_ref); + drop(disabled); + assert!(handle.upgrade().is_none()); + recording_service(handle).call(request()).await.unwrap(); + + let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All)); + let recorder_ref: EventRecorderRef = recorder.clone(); + let handle = AdminEventRecorderHandle::default(); + handle.install(&recorder_ref); + let service = recording_service(handle.clone()); + let future = service.call(request()); + drop(recorder_ref); + drop(recorder); + assert!(handle.upgrade().is_none()); + future.await.unwrap(); + } + + #[tokio::test] + async fn recording_does_not_retain_recorder_while_pending() { + let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All)); + let recorder_ref: EventRecorderRef = recorder.clone(); + let handle = AdminEventRecorderHandle::default(); + handle.install(&recorder_ref); + let started = Arc::new(Notify::new()); + let service = + AdminFunctionRecordingLayer::new(handle.clone()).layer(Arc::new(PendingService { + started: started.clone(), + })); + + let task = tokio::spawn(async move { service.call(request()).await }); + started.notified().await; + drop(recorder_ref); + drop(recorder); + + assert!(handle.upgrade().is_none()); + task.abort(); + assert!(matches!(task.await, Err(error) if error.is_cancelled())); + } + + #[tokio::test] + async fn dropping_pending_call_records_failure() { + let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All)); + let recorder_ref: EventRecorderRef = recorder.clone(); + let handle = AdminEventRecorderHandle::default(); + handle.install(&recorder_ref); + let started = Arc::new(Notify::new()); + let service = AdminFunctionRecordingLayer::new(handle).layer(Arc::new(PendingService { + started: started.clone(), + })); + + let task = tokio::spawn(async move { service.call(request()).await }); + started.notified().await; + task.abort(); + assert!(matches!(task.await, Err(error) if error.is_cancelled())); + assert_eq!(recorder.events.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn timing_out_pending_call_records_failure() { + let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All)); + let recorder_ref: EventRecorderRef = recorder.clone(); + let handle = AdminEventRecorderHandle::default(); + handle.install(&recorder_ref); + let service = AdminFunctionRecordingLayer::new(handle).layer(Arc::new(PendingService { + started: Arc::new(Notify::new()), + })); + + assert!( + tokio::time::timeout(Duration::from_millis(10), service.call(request())) + .await + .is_err() + ); + assert_eq!(recorder.events.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn empty_immediate_result_preserves_success_output_and_records_no_result() { + let recorder = Arc::new(TestRecorder::new(EventTypeFilter::All)); + let recorder_ref: EventRecorderRef = recorder.clone(); + let handle = AdminEventRecorderHandle::default(); + handle.install(&recorder_ref); + let service = AdminFunctionRecordingLayer::new(handle).layer(Arc::new(TestService { + trace: None, + immediate_result: None, + })); + + let response = service.call(request()).await.unwrap(); + assert!(matches!(response.output.data, OutputData::AffectedRows(0))); + assert_eq!(response.immediate_result, None); + + let events = recorder.events.lock().unwrap(); + assert_eq!(events.len(), 1); + assert_eq!( + events[0].extra_rows().unwrap()[0].values[3], + jsonb_value(&json!({})) + ); + } + + fn recording_service(handle: AdminEventRecorderHandle) -> AdminFunctionServiceRef { + let core = Arc::new(TestService { + trace: None, + immediate_result: Some(Value::UInt64(0)), + }); + AdminFunctionRecordingLayer::new(handle).layer(core) + } + + fn response(immediate_result: Option) -> AdminFunctionResponse { + AdminFunctionResponse { + output: Output::new_with_affected_rows(0), + immediate_result, + } + } + + fn request() -> AdminFunctionRequest { + let Statement::Admin(statement) = ParserContext::create_with_dialect( + "ADMIN flush_table('demo')", + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap() + .remove(0) else { + panic!("expected ADMIN statement") + }; + AdminFunctionRequest { + statement, + query_ctx: QueryContext::arc(), + } + } +} diff --git a/tests-integration/tests/admin_function_event.rs b/tests-integration/tests/admin_function_event.rs new file mode 100644 index 00000000000..d669d2900ad --- /dev/null +++ b/tests-integration/tests/admin_function_event.rs @@ -0,0 +1,193 @@ +// Copyright 2023 Greptime Team +// +// 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. + +use std::sync::Arc; + +use common_test_util::temp_dir::create_temp_dir; +use frontend::instance::Instance; +use servers::query_handler::sql::SqlQueryHandler; +use session::context::QueryContext; +use sql::ast::FunctionArguments; +use sql::dialect::GreptimeDbDialect; +use sql::parser::{ParseOptions, ParserContext}; +use sql::statements::admin::Admin; +use sql::statements::statement::Statement; +use tests_integration::cluster::GreptimeDbClusterBuilder; +use tests_integration::standalone::GreptimeDbStandaloneBuilder; +use tests_integration::test_util::{StorageType, get_test_store_config}; + +use crate::event_recorder_test_util::assert_eventually_eq; + +const TABLE_NAME: &str = "admin_function_event_test"; + +#[tokio::test(flavor = "multi_thread")] +async fn test_standalone_admin_function_events() { + common_telemetry::init_default_ut_logging(); + let standalone = GreptimeDbStandaloneBuilder::new("test_standalone_admin_function_events") + .build() + .await; + assert_admin_function_events(standalone.fe_instance()).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_distributed_admin_function_events() { + let store_type = StorageType::File; + if !store_type.test_on() { + return; + } + + common_telemetry::init_default_ut_logging(); + let (store_config, _guard) = get_test_store_config(&store_type); + let home_dir = create_temp_dir("test_distributed_admin_function_events_data_home"); + let cluster = GreptimeDbClusterBuilder::new("test_distributed_admin_function_events") + .await + .with_datanodes(1) + .with_store_config(store_config) + .with_shared_home_dir(Arc::new(home_dir)) + .build(true) + .await; + assert_admin_function_events(cluster.fe_instance()).await; +} + +async fn assert_admin_function_events(instance: &Arc) { + execute( + instance, + &format!("CREATE TABLE {TABLE_NAME} (ts TIMESTAMP TIME INDEX)"), + ) + .await + .unwrap(); + execute(instance, &format!("ADMIN FLUSH_TABLE('{TABLE_NAME}')")) + .await + .unwrap(); + + assert_admin_event( + instance, + "flush_table", + "\ ++----------------+----------+---------------------+-----------+---------------------------------------------------------------------------------+--------------+ +| type | actor | admin_function_name | status | payload | output | ++----------------+----------+---------------------+-----------+---------------------------------------------------------------------------------+--------------+ +| admin_function | greptime | flush_table | Succeeded | {\"arguments\":{\"type\":\"list\",\"value\":[\"admin_function_event_test\"]},\"version\":1} | {\"result\":0} | ++----------------+----------+---------------------+-----------+---------------------------------------------------------------------------------+--------------+", + ) + .await; + + let failure = execute(instance, "ADMIN missing_admin_function('missing-input')").await; + assert!(failure.is_err()); + + assert_admin_event( + instance, + "missing_admin_function", + "\ ++----------------+----------+------------------------+--------+---------------------------------------------------------------------+-----------------------------------------------------------------+ +| type | actor | admin_function_name | status | payload | output | ++----------------+----------+------------------------+--------+---------------------------------------------------------------------+-----------------------------------------------------------------+ +| admin_function | greptime | missing_admin_function | Failed | {\"arguments\":{\"type\":\"list\",\"value\":[\"missing-input\"]},\"version\":1} | {\"error\":\"0: Admin function not found: missing_admin_function\"} | ++----------------+----------+------------------------+--------+---------------------------------------------------------------------+-----------------------------------------------------------------+", + ) + .await; + + let failure = + execute_with_arguments(instance, "missing_no_arguments", FunctionArguments::None).await; + assert!(failure.is_err()); + + assert_admin_event( + instance, + "missing_no_arguments", + "\ ++----------------+----------+----------------------+--------+---------------+---------------------------------------------------------------+ +| type | actor | admin_function_name | status | payload | output | ++----------------+----------+----------------------+--------+---------------+---------------------------------------------------------------+ +| admin_function | greptime | missing_no_arguments | Failed | {\"version\":1} | {\"error\":\"0: Admin function not found: missing_no_arguments\"} | ++----------------+----------+----------------------+--------+---------------+---------------------------------------------------------------+", + ) + .await; + + let Statement::Query(query) = ParserContext::create_with_dialect( + "SELECT 1", + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap() + .remove(0) else { + panic!("expected query statement") + }; + let failure = execute_with_arguments( + instance, + "missing_subquery", + FunctionArguments::Subquery(Box::new(query.inner)), + ) + .await; + assert!(failure.is_err()); + + assert_admin_event( + instance, + "missing_subquery", + "\ ++----------------+----------+---------------------+--------+------------------------------------------------------------------+-----------------------------------------------------------+ +| type | actor | admin_function_name | status | payload | output | ++----------------+----------+---------------------+--------+------------------------------------------------------------------+-----------------------------------------------------------+ +| admin_function | greptime | missing_subquery | Failed | {\"arguments\":{\"type\":\"subquery\",\"value\":\"SELECT 1\"},\"version\":1} | {\"error\":\"0: Admin function not found: missing_subquery\"} | ++----------------+----------+---------------------+--------+------------------------------------------------------------------+-----------------------------------------------------------+", + ) + .await; +} + +async fn assert_admin_event(instance: &Arc, function: &str, expected: &str) { + let query = format!( + r#"SELECT + type, + actor, + admin_function_name, + admin_function_status AS status, + json_to_string(payload) AS payload, + json_to_string(admin_function_output) AS output +FROM greptime_private.events +WHERE type = 'admin_function' + AND admin_function_name = '{function}' +ORDER BY timestamp DESC +LIMIT 1"# + ); + assert_eventually_eq(instance, &query, expected).await; +} + +async fn execute( + instance: &Arc, + sql: &str, +) -> servers::error::Result { + instance.do_query(sql, QueryContext::arc()).await.remove(0) +} + +async fn execute_with_arguments( + instance: &Arc, + function: &str, + arguments: FunctionArguments, +) -> operator::error::Result { + let Statement::Admin(mut statement) = ParserContext::create_with_dialect( + &format!("ADMIN {function}()"), + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap() + .remove(0) else { + panic!("expected ADMIN statement") + }; + let Admin::Func(admin_function) = &mut statement; + admin_function.args = arguments; + + instance + .statement_executor() + .execute_sql(Statement::Admin(statement), QueryContext::arc()) + .await +} diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index 1925b46d453..3112cb72c32 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -14,6 +14,7 @@ #![recursion_limit = "256"] +mod admin_function_event; mod database_ddl_event; mod event_recorder_test_util; #[macro_use] From 154f90b3652bf9b1b3e7eccc5f277f856a2f3079 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Wed, 12 Aug 2026 03:12:04 +0000 Subject: [PATCH 29/92] fix: harden permission checks and process visibility (#8852) Signed-off-by: shuiyisong --- src/auth/src/lib.rs | 4 +- src/auth/src/permission.rs | 2 + src/auth/src/user_info.rs | 75 +++++++----- src/auth/src/user_provider.rs | 10 +- .../src/user_provider/static_user_provider.rs | 4 + src/frontend/src/instance.rs | 112 +++++++++++++++++- src/frontend/src/instance/grpc.rs | 24 ++-- src/query/src/sql.rs | 6 +- 8 files changed, 195 insertions(+), 42 deletions(-) diff --git a/src/auth/src/lib.rs b/src/auth/src/lib.rs index d2a1c05b1cd..d107068940e 100644 --- a/src/auth/src/lib.rs +++ b/src/auth/src/lib.rs @@ -31,8 +31,8 @@ pub use common::{ }; pub use permission::{ ALL_ACTIONS, AccessMode, DASHBOARD_DELETE, DASHBOARD_QUERY, DASHBOARD_SAVE, - DefaultPermissionChecker, INFLUXDB_WRITE, JAEGER_QUERY, LOG_QUERY, LOG_WRITE, OPENTSDB_WRITE, - OTLP_WRITE, PIPELINE_DELETE, PIPELINE_INSERT, PIPELINE_QUERY, PROM_STORE_READ, + DefaultPermissionChecker, ICEBERG_READ, INFLUXDB_WRITE, JAEGER_QUERY, LOG_QUERY, LOG_WRITE, + OPENTSDB_WRITE, OTLP_WRITE, PIPELINE_DELETE, PIPELINE_INSERT, PIPELINE_QUERY, PROM_STORE_READ, PROM_STORE_WRITE, PROMQL_QUERY, PermissionAction, PermissionChecker, PermissionReq, PermissionResp, PermissionTableTarget, PermissionTableTargets, SEMANTIC_GRAPH_QUERY, }; diff --git a/src/auth/src/permission.rs b/src/auth/src/permission.rs index 901b935d6a2..984fe8ea4cf 100644 --- a/src/auth/src/permission.rs +++ b/src/auth/src/permission.rs @@ -139,6 +139,7 @@ pub const PROMQL_QUERY: PermissionAction = PermissionAction::read("promql.query" pub const LOG_QUERY: PermissionAction = PermissionAction::read("log.query"); pub const OPENTSDB_WRITE: PermissionAction = PermissionAction::write("opentsdb.write"); pub const INFLUXDB_WRITE: PermissionAction = PermissionAction::write("influxdb.write"); +pub const ICEBERG_READ: PermissionAction = PermissionAction::read("iceberg.read"); pub const PROM_STORE_WRITE: PermissionAction = PermissionAction::write("prom_store.write"); pub const PROM_STORE_READ: PermissionAction = PermissionAction::read("prom_store.read"); pub const OTLP_WRITE: PermissionAction = PermissionAction::write("otlp.write"); @@ -164,6 +165,7 @@ pub const ALL_ACTIONS: &[PermissionAction] = &[ LOG_QUERY, OPENTSDB_WRITE, INFLUXDB_WRITE, + ICEBERG_READ, PROM_STORE_WRITE, PROM_STORE_READ, OTLP_WRITE, diff --git a/src/auth/src/user_info.rs b/src/auth/src/user_info.rs index 6ce735c4e2f..ab7c7f32f46 100644 --- a/src/auth/src/user_info.rs +++ b/src/auth/src/user_info.rs @@ -21,6 +21,10 @@ use crate::UserInfoRef; pub trait UserInfo: Debug + Sync + Send { fn as_any(&self) -> &dyn Any; fn username(&self) -> &str; + + fn is_admin(&self) -> bool { + false + } } /// The user permission mode @@ -38,13 +42,14 @@ impl PermissionMode { /// - "rw", "readwrite", "read_write" => ReadWrite /// - "ro", "readonly", "read_only" => ReadOnly /// - "wo", "writeonly", "write_only" => WriteOnly - /// Returns None if the input string is not a valid permission mode. - pub fn from_str(s: &str) -> Self { + /// + /// Returns `None` if the input string is not a valid permission mode. + pub fn from_str(s: &str) -> Option { match s.to_lowercase().as_str() { - "readwrite" | "read_write" | "rw" => PermissionMode::ReadWrite, - "readonly" | "read_only" | "ro" => PermissionMode::ReadOnly, - "writeonly" | "write_only" | "wo" => PermissionMode::WriteOnly, - _ => PermissionMode::ReadWrite, + "readwrite" | "read_write" | "rw" => Some(PermissionMode::ReadWrite), + "readonly" | "read_only" | "ro" => Some(PermissionMode::ReadOnly), + "writeonly" | "write_only" | "wo" => Some(PermissionMode::WriteOnly), + _ => None, } } @@ -123,58 +128,72 @@ mod tests { // Test ReadWrite variants assert_eq!( PermissionMode::from_str("readwrite"), - PermissionMode::ReadWrite + Some(PermissionMode::ReadWrite) ); assert_eq!( PermissionMode::from_str("read_write"), - PermissionMode::ReadWrite + Some(PermissionMode::ReadWrite) + ); + assert_eq!( + PermissionMode::from_str("rw"), + Some(PermissionMode::ReadWrite) ); - assert_eq!(PermissionMode::from_str("rw"), PermissionMode::ReadWrite); assert_eq!( PermissionMode::from_str("ReadWrite"), - PermissionMode::ReadWrite + Some(PermissionMode::ReadWrite) + ); + assert_eq!( + PermissionMode::from_str("RW"), + Some(PermissionMode::ReadWrite) ); - assert_eq!(PermissionMode::from_str("RW"), PermissionMode::ReadWrite); // Test ReadOnly variants assert_eq!( PermissionMode::from_str("readonly"), - PermissionMode::ReadOnly + Some(PermissionMode::ReadOnly) ); assert_eq!( PermissionMode::from_str("read_only"), - PermissionMode::ReadOnly + Some(PermissionMode::ReadOnly) + ); + assert_eq!( + PermissionMode::from_str("ro"), + Some(PermissionMode::ReadOnly) ); - assert_eq!(PermissionMode::from_str("ro"), PermissionMode::ReadOnly); assert_eq!( PermissionMode::from_str("ReadOnly"), - PermissionMode::ReadOnly + Some(PermissionMode::ReadOnly) + ); + assert_eq!( + PermissionMode::from_str("RO"), + Some(PermissionMode::ReadOnly) ); - assert_eq!(PermissionMode::from_str("RO"), PermissionMode::ReadOnly); // Test WriteOnly variants assert_eq!( PermissionMode::from_str("writeonly"), - PermissionMode::WriteOnly + Some(PermissionMode::WriteOnly) ); assert_eq!( PermissionMode::from_str("write_only"), - PermissionMode::WriteOnly + Some(PermissionMode::WriteOnly) + ); + assert_eq!( + PermissionMode::from_str("wo"), + Some(PermissionMode::WriteOnly) ); - assert_eq!(PermissionMode::from_str("wo"), PermissionMode::WriteOnly); assert_eq!( PermissionMode::from_str("WriteOnly"), - PermissionMode::WriteOnly + Some(PermissionMode::WriteOnly) ); - assert_eq!(PermissionMode::from_str("WO"), PermissionMode::WriteOnly); - - // Test invalid inputs default to ReadWrite assert_eq!( - PermissionMode::from_str("invalid"), - PermissionMode::ReadWrite + PermissionMode::from_str("WO"), + Some(PermissionMode::WriteOnly) ); - assert_eq!(PermissionMode::from_str(""), PermissionMode::ReadWrite); - assert_eq!(PermissionMode::from_str("xyz"), PermissionMode::ReadWrite); + + for invalid in ["readonyl", "", "xyz"] { + assert_eq!(PermissionMode::from_str(invalid), None); + } } #[test] @@ -200,7 +219,7 @@ mod tests { for mode in modes { let str_repr = mode.as_str(); let parsed = PermissionMode::from_str(str_repr); - assert_eq!(mode, parsed); + assert_eq!(Some(mode), parsed); } } diff --git a/src/auth/src/user_provider.rs b/src/auth/src/user_provider.rs index 985e25978e8..bafbd4fffce 100644 --- a/src/auth/src/user_provider.rs +++ b/src/auth/src/user_provider.rs @@ -463,7 +463,7 @@ pub(crate) fn parse_credential_line( let (username_part, password) = (parts[0], parts[1]); let (username, permission_mode) = if let Some((user, perm)) = username_part.split_once(':') { - (user, PermissionMode::from_str(perm)) + (user, PermissionMode::from_str(perm)?) } else { (username_part, PermissionMode::default()) }; @@ -766,6 +766,14 @@ mod tests { let result = parse_credential_line("user=pass=word"); assert_eq!(result, None); + for line in [ + "user:readonyl=password", + "user:=password", + "user:arbitrary=password", + ] { + assert_eq!(parse_credential_line(line), None); + } + // Empty password let result = parse_credential_line("user="); assert_eq!( diff --git a/src/auth/src/user_provider/static_user_provider.rs b/src/auth/src/user_provider/static_user_provider.rs index 3e0dac8ee74..69d7fd63d06 100644 --- a/src/auth/src/user_provider/static_user_provider.rs +++ b/src/auth/src/user_provider/static_user_provider.rs @@ -147,6 +147,8 @@ pub mod test { let provider = StaticUserProvider::new("cmd:root=123456,admin=654321").unwrap(); test_authenticate(&provider, "root", "123456").await; test_authenticate(&provider, "admin", "654321").await; + + assert!(StaticUserProvider::new("cmd:user:readonyl=password").is_err()); } #[tokio::test] @@ -171,6 +173,7 @@ pub mod test { assert!( lw.write_all( b"root=123456 +invalid:readonyl=password admin=654321", ) .is_ok() @@ -182,5 +185,6 @@ admin=654321", let provider = StaticUserProvider::new(param.as_str()).unwrap(); test_authenticate(&provider, "root", "123456").await; test_authenticate(&provider, "admin", "654321").await; + test_authenticate_fails(&provider, "invalid", "password").await; } } diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index bdf4d66c7e9..0c300835e0e 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -1665,6 +1665,7 @@ fn should_track_plan_process(stmt: Option<&Statement>, plan: &LogicalPlan) -> bo #[cfg(test)] mod tests { + use std::any::Any; use std::collections::HashMap; use std::future::Future; use std::pin::Pin; @@ -1676,10 +1677,12 @@ mod tests { use api::prom_store::remote::{ Label, LabelMatcher, Query as RemoteQuery, ReadRequest, ReadResponse, Sample, }; + use api::v1::greptime_request::Request; use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse}; + use api::v1::query_request::Query; use auth::{ DASHBOARD_DELETE, DASHBOARD_QUERY, DASHBOARD_SAVE, JAEGER_QUERY, PIPELINE_DELETE, - PIPELINE_INSERT, PIPELINE_QUERY, PermissionAction, PermissionResp, UserInfoRef, + PIPELINE_INSERT, PIPELINE_QUERY, PermissionAction, PermissionResp, UserInfo, UserInfoRef, }; use catalog::process_manager::{ProcessManager, QueryStatement, SlowQueryTimer}; use common_base::Plugins; @@ -1942,6 +1945,23 @@ mod tests { ) } + #[derive(Debug)] + struct AdminUserInfo; + + impl UserInfo for AdminUserInfo { + fn as_any(&self) -> &dyn Any { + self + } + + fn username(&self) -> &str { + "admin" + } + + fn is_admin(&self) -> bool { + true + } + } + struct RejectUnresolvedPermissionChecker; impl PermissionChecker for RejectUnresolvedPermissionChecker { @@ -3102,6 +3122,45 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_write_only_grpc_sql_is_checked_after_parsing() -> TestResult<()> { + let plugins = Plugins::new(); + plugins.insert::(Arc::new(WriteOnlyPermissionChecker)); + let instance = test_instance_with_plugins( + test_table(1024, "source")?, + test_table(1025, "target")?, + plugins, + ) + .await?; + + let insert = Request::Query(api::v1::QueryRequest { + query: Some(Query::Sql( + "INSERT INTO target SELECT * FROM source".to_string(), + )), + }); + servers::query_handler::grpc::GrpcQueryHandler::do_query( + &instance, + insert, + QueryContext::arc(), + ) + .await + .unwrap(); + + let select = Request::Query(api::v1::QueryRequest { + query: Some(Query::Sql("SELECT * FROM source".to_string())), + }); + assert_permission_denied( + servers::query_handler::grpc::GrpcQueryHandler::do_query( + &instance, + select, + QueryContext::arc(), + ) + .await, + ); + + Ok(()) + } + #[tokio::test] async fn test_target_independent_checker_skips_target_resolution() -> TestResult<()> { let physical_table = "physical_metric"; @@ -3438,6 +3497,57 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_show_processlist_catalog_scope() -> TestResult<()> { + let instance = + test_instance_with_tables(test_table(1024, "source")?, test_table(1025, "target")?) + .await?; + let _current_catalog = instance.process_manager().register_query( + "greptime".to_string(), + vec!["public".to_string()], + "current_catalog_query".to_string(), + String::new(), + None, + None, + ); + let _other_catalog = instance.process_manager().register_query( + "other".to_string(), + vec!["public".to_string()], + "other_catalog_query".to_string(), + String::new(), + None, + None, + ); + + for sql in ["SHOW PROCESSLIST", "SHOW FULL PROCESSLIST"] { + let output = execute_one_sql(&instance, sql, test_query_ctx(43)).await?; + let process_list = output.data.pretty_print().await; + assert!( + process_list.contains("current_catalog_query"), + "{process_list}" + ); + assert!( + !process_list.contains("other_catalog_query"), + "{process_list}" + ); + + let admin_ctx = test_query_ctx(44); + admin_ctx.set_current_user(Arc::new(AdminUserInfo)); + let output = execute_one_sql(&instance, sql, admin_ctx).await?; + let process_list = output.data.pretty_print().await; + assert!( + process_list.contains("current_catalog_query"), + "{process_list}" + ); + assert!( + process_list.contains("other_catalog_query"), + "{process_list}" + ); + } + + Ok(()) + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn test_kill_query_cancels_insert_select() -> TestResult<()> { assert_kill_cancels_insert_select("KILL QUERY 4242").await diff --git a/src/frontend/src/instance/grpc.rs b/src/frontend/src/instance/grpc.rs index acb3a79e049..7d4f0ee7959 100644 --- a/src/frontend/src/instance/grpc.rs +++ b/src/frontend/src/instance/grpc.rs @@ -71,15 +71,21 @@ impl GrpcQueryHandler for Instance { let interceptor = interceptor_ref.as_ref(); interceptor.pre_execute(&request, ctx.clone())?; - self.plugins - .get::() - .as_ref() - .check_permission_with_context( - ctx.current_user(), - PermissionReq::GrpcRequest(&request), - Some(&ctx.current_schema()), - ) - .context(PermissionSnafu)?; + if !matches!( + &request, + Request::Query(query_request) + if matches!(&query_request.query, Some(Query::Sql(_))) + ) { + self.plugins + .get::() + .as_ref() + .check_permission_with_context( + ctx.current_user(), + PermissionReq::GrpcRequest(&request), + Some(&ctx.current_schema()), + ) + .context(PermissionSnafu)?; + } let output = match request { Request::Inserts(requests) => self.handle_inserts(requests, ctx.clone()).await?, diff --git a/src/query/src/sql.rs b/src/query/src/sql.rs index fc755deb8a7..b2d979c6664 100644 --- a/src/query/src/sql.rs +++ b/src/query/src/sql.rs @@ -1360,7 +1360,11 @@ pub async fn show_processlist( ] }; - let filters = vec![]; + let filters = if query_ctx.current_user().is_admin() { + vec![] + } else { + vec![col(process_list::CATALOG).eq(lit(query_ctx.current_catalog()))] + }; let like_field = None; let sort = vec![col("id").sort(true, true)]; query_from_information_schema_table( From e778a72829d5638a6ea7892f6c4fa32ecede86a0 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Wed, 12 Aug 2026 04:06:31 +0000 Subject: [PATCH 30/92] feat: complete the derived-edge vocabulary of the entity graph (#8836) * feat(operator): pair calls edges across trace tables and derive virtual-node edges Union the normalized client and server spans of all trace tables before the join, so a client span pairs with a server span stored in a different table. A client span with no matching server span becomes an edge to a virtual node named by span attributes (peer.service / db.name / server.address), with confidence < 1.0 and attributes.connection_type; a window's real pairs win over virtual candidates for the same edge key. Signed-off-by: Dennis Zhuang * feat(operator): derive same-row co-declared edges from the built-in vocabulary A table declaring both entity types of a vocabulary pair witnesses the edge on every row carrying both identities: runs_on / contains / part_of for any declaring table (provenance 'attribute'), agent uses model / agent invoked tool only for trace sources (span-structure observations, provenance 'trace'). Signed-off-by: Dennis Zhuang * feat(operator): derive parent_agent-calls-agent edges from span structure Trace tables declaring an agent entity pair each span with its child span across tables (no span-kind filter), keep pairs whose agent identities differ, and aggregate RED metrics per window, anchored on the parent span like the service derivation is anchored on the client. Signed-off-by: Dennis Zhuang * feat(frontend): feed co-declared and agent sources into the relationships scan scan_relationships now passes every declaring table (with its trace-ness) to the co-declared branch and the trace tables' agent declarations to the agent-calls derivation. enumerate validates the fixed trace-v1 columns and derives around a malformed trace table instead of failing the whole scan. Signed-off-by: Dennis Zhuang * test: cover cross-table pairing, virtual nodes, co-declared and agent edges sqlness exercises the new derivations end to end (including a malformed trace-model table being skipped); the integration authorization test now also pins that a pair split across tables derives no edge when the caller cannot read one side. Signed-off-by: Dennis Zhuang * chore: update the relationships module doc for the new branches Signed-off-by: Dennis Zhuang * chore: import shared derivation helpers via crate paths The fmt CI gate rejects module-level 'use super::' imports. Signed-off-by: Dennis Zhuang * fix: fold co-declared duplicates, decouple agent calls, verify the trace time index Review findings: the co-declared branch lacked a cross-source DISTINCT, so two tables witnessing the same edge in one window emitted duplicate rows; the agent-calls derivation was gated on a usable service declaration; the trace schema guard accepted a table whose time index is not the column the derivations bucket by. The empty-trace-table test asserted a union invariant with no information and is dropped. Signed-off-by: Dennis Zhuang * fix: rename the agent-tool edge to invokes and track current OTel peer attributes The vocabulary's other relation names are present tense; semconv 1.39/1.26 replaced peer.service and db.name with service.peer.name and db.namespace, so the virtual-node candidates now check the current names first and keep the deprecated ones for existing telemetry. Signed-off-by: Dennis Zhuang * fix: trust the trace-v1 table option instead of matching the fixed schema The option is only ever stamped by the ingest path, which guarantees the fixed span columns; matching column types here couples the graph to every trace schema evolution (e.g. #8816) for a case that cannot occur. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/frontend/src/instance/entity_graph.rs | 85 +- src/operator/src/statement/semantic_graph.rs | 78 +- .../statement/semantic_graph/relationships.rs | 1462 ++++++++++++++--- tests-integration/src/instance.rs | 61 +- .../common/system/semantic_graph.result | 201 +++ .../common/system/semantic_graph.sql | 137 ++ 6 files changed, 1748 insertions(+), 276 deletions(-) diff --git a/src/frontend/src/instance/entity_graph.rs b/src/frontend/src/instance/entity_graph.rs index e55e5f2b8d7..70c57846fd0 100644 --- a/src/frontend/src/instance/entity_graph.rs +++ b/src/frontend/src/instance/entity_graph.rs @@ -47,9 +47,9 @@ use datafusion::dataframe::DataFrame; use datafusion_expr::LogicalPlan; use futures::TryStreamExt; use operator::statement::semantic_graph::{ - CallsSource, DeclaredSource, EntityDeclaration, GraphQueryWindow, OBSERVED_AT_COLUMN, - RegistrySource, build_registry_plan, build_relationships_plan, - declared_relationships_schema_matches, + CallsSource, CoDeclaredSource, DeclaredSource, EntityDeclaration, GraphQueryWindow, + OBSERVED_AT_COLUMN, RegistrySource, RelationshipSources, build_registry_plan, + build_relationships_plan, declared_relationships_schema_matches, }; use query::QueryEngineRef; use session::context::{QueryContext, QueryContextBuilder, QueryContextRef}; @@ -73,6 +73,13 @@ pub struct EntityGraphProviderImpl { struct EntitySource { declarations: Vec, + is_trace: bool, + table: TableRef, +} + +struct TraceSource { + service: Option, + agent: Option, table: TableRef, } @@ -222,7 +229,7 @@ impl EntityGraphProviderImpl { &self, catalog: &str, query_ctx: Option<&QueryContext>, - ) -> Result<(Vec, Vec<(EntityDeclaration, TableRef)>), BoxedError> { + ) -> Result<(Vec, Vec), BoxedError> { let Some(catalog_manager) = self.catalog_manager.upgrade() else { return Ok((vec![], vec![])); }; @@ -288,26 +295,34 @@ impl EntityGraphProviderImpl { continue; } if is_trace { - // TODO(entity-graph): validate the complete fixed trace-v1 - // schema before adding the table to calls derivation, and - // skip malformed/stale tables instead of letting one poison - // the whole relationship scan. - match table_declarations - .iter() - .find(|d| d.entity_type == "service") - { - Some(service) => traces.push((service.clone(), table.clone())), + let find = |entity_type: &str| { + table_declarations + .iter() + .find(|d| d.entity_type == entity_type) + .cloned() + }; + let service = find("service"); + let agent = find("agent"); + if service.is_none() { // No usable service identity: the table cannot - // contribute calls edges (see declarations_for). - None => warn!( + // contribute service-calls edges (see declarations_for). + warn!( "Trace table `{}` has no usable service declaration; skipping calls derivation", table_info.name - ), + ); + } + if service.is_some() || agent.is_some() { + traces.push(TraceSource { + service, + agent, + table: table.clone(), + }); } } if !table_declarations.is_empty() { declarations.push(EntitySource { declarations: table_declarations, + is_trace, table, }); } @@ -480,21 +495,37 @@ impl EntityGraphProvider for EntityGraphProviderImpl { request: ScanRequest, query_ctx: Option, ) -> Result, BoxedError> { - let (_, traces) = self.enumerate(catalog, query_ctx.as_deref()).await?; - let mut scans = Vec::with_capacity(traces.len()); - for (service, trace) in traces { - scans.push(CallsSource { - service, - scan: self.read_table(trace)?, + let (sources, traces) = self.enumerate(catalog, query_ctx.as_deref()).await?; + let mut calls = Vec::with_capacity(traces.len()); + for trace in traces { + calls.push(CallsSource { + service: trace.service, + agent: trace.agent, + scan: self.read_table(trace.table)?, + }); + } + let mut co_declared = Vec::with_capacity(sources.len()); + for source in sources { + co_declared.push(CoDeclaredSource { + declarations: source.declarations, + is_trace: source.is_trace, + scan: self.read_table(source.table)?, }); } let declared = self.declared_source(catalog, query_ctx.as_deref()).await?; let Some(window) = Self::query_window(&request)? else { return Ok(None); }; - let Some(plan) = build_relationships_plan(scans, declared, &window) - .context(error::DataFusionSnafu) - .map_err(BoxedError::new)? + let Some(plan) = build_relationships_plan( + RelationshipSources { + traces: calls, + co_declared, + declared, + }, + &window, + ) + .context(error::DataFusionSnafu) + .map_err(BoxedError::new)? else { return Ok(None); }; @@ -529,6 +560,10 @@ mod tests { .iter() .map(|c| ColumnSchema::new(*c, ConcreteDataType::string_datatype(), true)), ); + assemble_table_info(column_schemas, extra) + } + + fn assemble_table_info(column_schemas: Vec, extra: &[(&str, &str)]) -> TableInfo { let schema = Arc::new( SchemaBuilder::try_from_columns(column_schemas) .unwrap() diff --git a/src/operator/src/statement/semantic_graph.rs b/src/operator/src/statement/semantic_graph.rs index af9108198d4..c14d9d0d04e 100644 --- a/src/operator/src/statement/semantic_graph.rs +++ b/src/operator/src/statement/semantic_graph.rs @@ -45,7 +45,9 @@ use datafusion::functions::{core as core_fns, datetime as datetime_fns, string a use datafusion::functions_nested::expr_fn::make_array; use datafusion_common::{Column, Result as DfResult, ScalarValue}; use datafusion_expr::{Expr, LogicalPlan, ScalarUDF, cast, ident, lit}; -pub use relationships::{CallsSource, DeclaredSource, build_relationships_plan}; +pub use relationships::{ + CallsSource, CoDeclaredSource, DeclaredSource, RelationshipSources, build_relationships_plan, +}; use store_api::mito_engine_options::{APPEND_MODE_KEY, MERGE_MODE_KEY, TTL_KEY}; /// Whether an existing declared-edge table still matches the canonical @@ -513,11 +515,44 @@ const REGISTRY_COLUMNS: [&str; 10] = [ const REGISTRY_VALID_COLUMN: &str = "__entity_valid"; -/// Projects all entity declarations of one source table with one source scan. -/// -/// Each output field is first built as an array whose entries correspond to the -/// table's declarations. Unnesting the arrays expands one source row into one -/// row per declared entity without cloning the table scan under `UNION ALL`. +/// Expands one source row into one output row per entry of `rows` with a +/// single source scan: each output field is first built as an array whose +/// entries correspond to the rows, then unnested. Every row is `[valid, +/// values...]` aligned with `columns`; rows whose `valid` expression is false +/// are dropped, and the distinct output columns are `columns`. +fn unnest_rows( + df: DataFrame, + window_predicate: Expr, + valid_column: &'static str, + columns: &[&'static str], + rows: Vec>, +) -> DfResult { + let mut arrays = vec![Vec::with_capacity(rows.len()); columns.len() + 1]; + for row in rows { + debug_assert_eq!(row.len(), arrays.len()); + for (array, value) in arrays.iter_mut().zip(row) { + array.push(value); + } + } + let array_names = std::iter::once(valid_column) + .chain(columns.iter().copied()) + .collect::>(); + let array_projection = array_names + .iter() + .zip(arrays) + .map(|(name, values)| make_array(values).alias(*name)) + .collect::>(); + + df.filter(window_predicate)? + .select(array_projection)? + .unnest_columns(&array_names)? + .filter(ident(valid_column))? + .select(columns.iter().map(|c| ident(*c)).collect::>())? + .distinct() +} + +/// Projects all entity declarations of one source table with one source scan +/// via [`unnest_rows`]. fn registry_source( first: &EntityDeclaration, rest: &[EntityDeclaration], @@ -531,7 +566,7 @@ fn registry_source( .gt_eq(window.source_start()) .and(ts.lt(window.source_end())); - let mut arrays = vec![Vec::new(); REGISTRY_COLUMNS.len() + 1]; + let mut rows = Vec::with_capacity(1 + rest.len()); for decl in std::iter::once(first).chain(rest) { // CAST even a single-column id: id columns must be tags but not // necessarily strings, and the computed table declares entity_id @@ -575,7 +610,7 @@ fn registry_source( predicate.and(ident(id).is_not_null()) }); - let row = [ + rows.push(vec![ valid, bin.clone(), bin.clone(), @@ -587,27 +622,16 @@ fn registry_source( scope, descriptive, source_tables, - ]; - for (array, value) in arrays.iter_mut().zip(row) { - array.push(value); - } + ]); } - let array_names = std::iter::once(REGISTRY_VALID_COLUMN) - .chain(REGISTRY_COLUMNS) - .collect::>(); - let array_projection = array_names - .iter() - .zip(arrays) - .map(|(name, values)| make_array(values).alias(*name)) - .collect::>(); - - df.filter(window_predicate)? - .select(array_projection)? - .unnest_columns(&array_names)? - .filter(ident(REGISTRY_VALID_COLUMN))? - .select(REGISTRY_COLUMNS.into_iter().map(ident).collect::>())? - .distinct() + unnest_rows( + df, + window_predicate, + REGISTRY_VALID_COLUMN, + ®ISTRY_COLUMNS, + rows, + ) } /// A declaring table's scan paired with the entity declarations it carries — diff --git a/src/operator/src/statement/semantic_graph/relationships.rs b/src/operator/src/statement/semantic_graph/relationships.rs index 7ebf5ae86fd..416c8dc10e1 100644 --- a/src/operator/src/statement/semantic_graph/relationships.rs +++ b/src/operator/src/statement/semantic_graph/relationships.rs @@ -13,9 +13,10 @@ // limitations under the License. //! The `semantic_relationships` derivation: the plan builders for every edge -//! branch (trace-derived `calls`, the declared-edge union) behind the computed -//! table. Shared expression helpers and the entity registry live in the parent -//! module. +//! branch behind the computed table — trace-derived service `calls` (with +//! virtual-node edges for unmatched clients), agent `calls` from span +//! structure, same-row co-declared edges, and the declared-edge union. Shared +//! expression helpers and the entity registry live in the parent module. use common_catalog::consts::{ DURATION_NANO_COLUMN, PARENT_SPAN_ID_COLUMN, SPAN_ID_COLUMN, SPAN_KIND_CLIENT, @@ -25,14 +26,15 @@ use common_catalog::consts::{ use datafusion::arrow::datatypes::DataType; use datafusion::dataframe::DataFrame; use datafusion::functions::core as core_fns; -use datafusion::functions_aggregate::expr_fn::{count, sum}; +use datafusion::functions_aggregate::expr_fn::{bool_or, count, min, sum}; use datafusion::functions_window::expr_fn::row_number; -use datafusion_common::Result as DfResult; -use datafusion_expr::{ExprFunctionExt, JoinType, LogicalPlan, cast, ident, lit}; +use datafusion_common::{Result as DfResult, ScalarValue}; +use datafusion_expr::{Expr, ExprFunctionExt, JoinType, LogicalPlan, cast, ident, lit, when}; use crate::statement::semantic_graph::{ DECLARED_EDGE_IDENTITY_COLUMNS, EntityDeclaration, GraphQueryWindow, OBSERVED_AT_COLUMN, - bin_interval, bin_ms, entity_id_expr, interval, null_json, qcol, union_all, + bin_interval, bin_ms, entity_id_expr, interval, null_json, parse_json_expr, qcol, union_all, + unnest_rows, }; /// The projected columns of `semantic_relationships`, in order. Every derived @@ -67,10 +69,23 @@ const RELATIONSHIP_COLUMNS: [&str; 16] = [ const CHILD_SPAN_EARLY_NANOS: i64 = 5 * 60 * 1_000_000_000; const CHILD_SPAN_LATE_NANOS: i64 = 60 * 60 * 1_000_000_000; -/// A trace table's scan paired with the `service` declaration it derives -/// `calls` edges for — a unit of `build_relationships_plan`. +/// A trace table's scan paired with the entity declarations its derivations +/// key on — a unit of `build_relationships_plan`. The `service` declaration +/// feeds service calls, the `agent` declaration agent calls; the two are +/// independent, so a table whose service declaration is unusable still +/// derives agent edges (and vice versa). pub struct CallsSource { - pub service: EntityDeclaration, + pub service: Option, + pub agent: Option, + pub scan: DataFrame, +} + +/// A declaring table's scan paired with its entity declarations, from which +/// the same-row co-declaration rules derive edges. `is_trace` gates the +/// agent-edge vocabulary, which the RFC ties to span structure. +pub struct CoDeclaredSource { + pub declarations: Vec, + pub is_trace: bool, pub scan: DataFrame, } @@ -80,17 +95,29 @@ pub struct DeclaredSource { pub scan: DataFrame, } -/// Builds the `semantic_relationships` plan: the trace-derived `calls` branch -/// unioned with the declared-edge branch, re-projected to the 16-column -/// contract. Returns `None` when there is neither a trace table nor a -/// declared-edge table, so the computed table streams empty. +/// Everything `build_relationships_plan` derives edges from. +pub struct RelationshipSources { + pub traces: Vec, + pub co_declared: Vec, + pub declared: Option, +} + +/// Builds the `semantic_relationships` plan: the service-calls, agent-calls, +/// co-declared, and declared-edge branches unioned and re-projected to the +/// 16-column contract. Returns `None` when no source can contribute edges, so +/// the computed table streams empty. pub fn build_relationships_plan( - traces: Vec, - declared: Option, + sources: RelationshipSources, window: &GraphQueryWindow, ) -> DfResult> { - let mut union_df = calls_branch(traces, window)?; - if let Some(declared) = declared { + let mut union_df = calls_branch(&sources.traces, window)?; + if let Some(agent_calls) = agent_calls_branch(&sources.traces, window)? { + union_df = union_all(union_df, agent_calls)?; + } + if let Some(co_declared) = co_declared_branch(sources.co_declared, window)? { + union_df = union_all(union_df, co_declared)?; + } + if let Some(declared) = sources.declared { union_df = union_all(union_df, declared_edges(declared.scan, window)?)?; } union_df @@ -102,6 +129,112 @@ pub fn build_relationships_plan( .transpose() } +/// The same-row co-declaration vocabulary open to any declaring table: a row +/// carrying both identities witnesses the edge, and the built-in direction is +/// `src -> dst`. Rows are `(src_type, dst_type, rel_type)`. +const ATTRIBUTE_EDGE_VOCABULARY: [(&str, &str, &str); 6] = [ + ("service.instance", "host", "runs_on"), + ("process", "host", "runs_on"), + ("k8s.pod", "k8s.node", "runs_on"), + ("k8s.pod", "k8s.container", "contains"), + ("service.instance", "service", "part_of"), + ("k8s.pod", "k8s.workload", "part_of"), +]; + +/// The agent-edge vocabulary, applied only to trace sources: the RFC derives +/// `agent uses model` / `agent invokes tool` from span structure (an LLM- or +/// tool-call span carries both identities), so a non-trace table co-declaring +/// these types must not fabricate invocations. +const AGENT_EDGE_VOCABULARY: [(&str, &str, &str); 2] = + [("agent", "model", "uses"), ("agent", "tool", "invokes")]; + +const CO_DECLARED_VALID_COLUMN: &str = "__edge_valid"; + +/// The same-row co-declared branch: for each source table, every vocabulary +/// pair whose two entity types the table declares yields an edge projection, +/// expanded from one scan via `unnest_rows` (the registry idiom). Attribute +/// edges carry `provenance = 'attribute'`; agent edges are span-structure +/// observations and carry `provenance = 'trace'`. +fn co_declared_branch( + sources: Vec, + window: &GraphQueryWindow, +) -> DfResult> { + let mut union_df: Option = None; + for source in sources { + let find = |ty: &str| source.declarations.iter().find(|d| d.entity_type == ty); + let mut pairs = Vec::new(); + for (src_type, dst_type, rel_type) in ATTRIBUTE_EDGE_VOCABULARY { + if let (Some(src), Some(dst)) = (find(src_type), find(dst_type)) { + pairs.push((src, dst, rel_type, "attribute")); + } + } + if source.is_trace { + for (src_type, dst_type, rel_type) in AGENT_EDGE_VOCABULARY { + if let (Some(src), Some(dst)) = (find(src_type), find(dst_type)) { + pairs.push((src, dst, rel_type, "trace")); + } + } + } + let Some((first, _, _, _)) = pairs.first() else { + continue; + }; + + let ts = ident(&first.time_index); + let bin = bin_ms(ts.clone()); + let window_predicate = ts + .clone() + .gt_eq(window.source_start()) + .and(ts.lt(window.source_end())); + let null_i64 = || lit(ScalarValue::Int64(None)); + let rows = pairs + .iter() + .map(|(src, dst, rel_type, provenance)| { + // Both endpoints must identify something on the row for the + // row to witness the edge. + let valid = src + .id_columns + .iter() + .chain(&dst.id_columns) + .fold(lit(true), |predicate, id| { + predicate.and(ident(id).is_not_null()) + }); + vec![ + valid, + bin.clone(), + bin.clone(), + bin.clone() + bin_interval(), + bin.clone() + bin_interval(), + lit(src.entity_type.as_str()), + entity_id_expr(&src.id_columns, &|c| ident(c)), + lit(dst.entity_type.as_str()), + entity_id_expr(&dst.id_columns, &|c| ident(c)), + lit(*rel_type), + lit(*provenance), + lit(1.0_f64), + null_i64(), + null_i64(), + lit(ScalarValue::Float64(None)), + null_i64(), + null_json(), + ] + }) + .collect(); + + let branch = unnest_rows( + source.scan, + window_predicate, + CO_DECLARED_VALID_COLUMN, + &RELATIONSHIP_COLUMNS, + rows, + )?; + union_df = union_all(union_df, branch)?; + } + // The per-source DISTINCT inside `unnest_rows` cannot see other sources: + // two tables witnessing the same edge in the same window must still fold + // into one row per (window, edge). + union_df.map(DataFrame::distinct).transpose() +} + /// Ranking column used to pick the latest declared-edge revision per edge key. const DECLARED_REVISION_COLUMN: &str = "__declared_revision"; @@ -179,27 +312,354 @@ fn declared_edges(scan: DataFrame, window: &GraphQueryWindow) -> DfResult, - window: &GraphQueryWindow, -) -> DfResult> { - let mut union_df: Option = None; +/// Span-attribute columns naming an uninstrumented peer, in precedence order, +/// each with the `connection_type` its match implies: current OTel names +/// first (`service.peer.name` and `db.namespace` replaced the deprecated +/// `peer.service` and `db.name` in semconv 1.39/1.26, kept for existing +/// telemetry), the bare network endpoint last. Attribute columns are dynamic +/// in `greptime_trace_v1` (created on first use), so only columns present in +/// a table's schema participate. +const VIRTUAL_DST_CANDIDATES: [(&str, &str); 5] = [ + ("span_attributes.service.peer.name", "virtual_node"), + ("span_attributes.peer.service", "virtual_node"), + ("span_attributes.db.namespace", "database"), + ("span_attributes.db.name", "database"), + ("span_attributes.server.address", "virtual_node"), +]; + +/// Confidence of a virtual-node edge: the peer was named by a client-side +/// attribute, not witnessed by a server span (the RFC requires `< 1.0`). +const VIRTUAL_NODE_CONFIDENCE: f64 = 0.5; + +/// The `calls` derivation: union the client spans and the server spans of all +/// trace tables (each projected to a normalized shape with its own table's +/// `service` identity), left-join clients to their child server spans on +/// `trace_id` + `parent_span_id`, and aggregate RED metrics per 60s window — +/// the plan form of the Tempo servicegraph connector. Unioning spans before +/// the join pairs a client with a server stored in a *different* trace table +/// (`x-greptime-trace-table-name` routing); with a single table it degenerates +/// to the previous self-join. Returns `None` when no trace table has a +/// usable `service` declaration. +/// +/// A client with no matching server span is an edge to a **virtual node** +/// named by [`VIRTUAL_DST_CANDIDATES`], with the client's own status/duration +/// and `confidence < 1.0`. Real pairs win: when a window's edge key holds any +/// pair, the edge reports only the pair population (the RFC pins RED metrics +/// to observed span pairs) and the unmatched clients are suppressed, so one +/// `(window, edge)` never yields two rows. +fn calls_branch(traces: &[CallsSource], window: &GraphQueryWindow) -> DfResult> { + let mut clients: Option = None; + let mut servers: Option = None; for trace in traces { - union_df = union_all(union_df, calls_pairs(&trace.service, trace.scan, window)?)?; + let Some(service) = &trace.service else { + continue; + }; + clients = union_all(clients, client_spans(service, &trace.scan, window)?)?; + servers = union_all(servers, server_spans(service, trace.scan.clone(), window)?)?; } - let Some(pairs) = union_df else { + let (Some(clients), Some(servers)) = (clients, servers) else { return Ok(None); }; - let df = pairs + let client = clients.alias("client")?; + let server = servers.alias("server")?; + let join_conditions = vec![ + qcol("client", TRACE_ID_COLUMN).eq(qcol("server", TRACE_ID_COLUMN)), + qcol("server", PARENT_SPAN_ID_COLUMN).eq(qcol("client", SPAN_ID_COLUMN)), + qcol("server", TRACE_TIMESTAMP_COLUMN) + .gt_eq(qcol("client", TRACE_TIMESTAMP_COLUMN) - interval(CHILD_SPAN_EARLY_NANOS)), + qcol("server", TRACE_TIMESTAMP_COLUMN) + .lt_eq(qcol("client", TRACE_TIMESTAMP_COLUMN) + interval(CHILD_SPAN_LATE_NANOS)), + ]; + + let observations = client + .join_on(server, JoinType::Left, join_conditions)? + .select(vec![ + bin_ms(qcol("client", TRACE_TIMESTAMP_COLUMN)).alias("observed_at"), + qcol("client", "src_id").alias("src_id"), + core_fns::coalesce() + .call(vec![ + qcol("server", "dst_id"), + qcol("client", "virtual_dst"), + ]) + .alias("dst_id"), + qcol("server", TRACE_ID_COLUMN) + .is_not_null() + .alias("paired"), + qcol("server", SPAN_STATUS_CODE_COLUMN).alias("server_status"), + qcol("server", DURATION_NANO_COLUMN).alias("server_duration_nano"), + qcol("client", SPAN_STATUS_CODE_COLUMN).alias("client_status"), + qcol("client", DURATION_NANO_COLUMN).alias("client_duration_nano"), + qcol("client", "virtual_conn").alias("virtual_conn"), + ])? + // No destination means no edge; a self-call is not an edge between + // two distinct entities. + .filter( + ident("dst_id") + .is_not_null() + .and(ident("src_id").not_eq(ident("dst_id"))), + )?; + + let paired = ident("paired"); + let df = observations + .aggregate( + vec![ident("observed_at"), ident("src_id"), ident("dst_id")], + vec![ + count(lit(1)) + .filter(paired.clone()) + .build()? + .alias("pair_count"), + count(lit(1)) + .filter( + paired + .clone() + .and(ident("server_status").eq(lit(SPAN_STATUS_ERROR))), + ) + .build()? + .alias("pair_errors"), + sum(ident("server_duration_nano")) + .filter(paired.clone()) + .build()? + .alias("pair_duration_nano"), + count(lit(1)) + .filter(!paired.clone()) + .build()? + .alias("unmatched_count"), + count(lit(1)) + .filter( + (!paired.clone()).and(ident("client_status").eq(lit(SPAN_STATUS_ERROR))), + ) + .build()? + .alias("unmatched_errors"), + sum(ident("client_duration_nano")) + .filter(!paired.clone()) + .build()? + .alias("unmatched_duration_nano"), + bool_or(paired.clone()).alias("has_pair"), + // `min` makes a mixed-provenance virtual edge deterministic + // (`database` sorts before `virtual_node`). + min(ident("virtual_conn")) + .filter(!paired) + .build()? + .alias("virtual_conn"), + ], + )? + .select(vec![ + ident("observed_at"), + ident("observed_at").alias("window_start"), + (ident("observed_at") + bin_interval()).alias("window_end"), + (ident("observed_at") + bin_interval()).alias("fresh_until"), + lit("service").alias("src_type"), + ident("src_id"), + lit("service").alias("dst_type"), + ident("dst_id"), + lit("calls").alias("rel_type"), + lit("trace").alias("provenance"), + real_wins(lit(1.0_f64), lit(VIRTUAL_NODE_CONFIDENCE))?.alias("confidence"), + real_wins(ident("pair_count"), ident("unmatched_count"))?.alias("request_count"), + real_wins(ident("pair_errors"), ident("unmatched_errors"))?.alias("error_count"), + // duration sums in nanoseconds; the contract column is seconds. + (cast( + real_wins( + ident("pair_duration_nano"), + ident("unmatched_duration_nano"), + )?, + DataType::Float64, + ) / lit(1e9_f64)) + .alias("duration_sum"), + real_wins(ident("pair_count"), ident("unmatched_count"))?.alias("duration_count"), + real_wins(null_json(), virtual_attrs_expr()?)?.alias("attributes"), + ])?; + Ok(Some(df)) +} + +/// `CASE WHEN has_pair THEN real ELSE virtual END` — the real-wins projection +/// over the mixed aggregate. +fn real_wins(real: Expr, r#virtual: Expr) -> DfResult { + when(ident("has_pair"), real).otherwise(r#virtual) +} + +/// The `attributes` JSON of a virtual edge, from the aggregated +/// `connection_type`. +fn virtual_attrs_expr() -> DfResult { + when( + ident("virtual_conn").eq(lit("database")), + parse_json_expr(lit(r#"{"connection_type":"database"}"#)), + ) + .otherwise(parse_json_expr(lit( + r#"{"connection_type":"virtual_node"}"#, + ))) +} + +/// The window + span-kind + identity predicate shared by both join sides. +/// `strict` bounds the scan to the source window; the non-strict side widens +/// by the join's time-proximity allowance (the join bounds reference the other +/// side's timestamp and cannot prune this side's scan on their own). +fn span_predicate(service: &EntityDeclaration, window: &GraphQueryWindow, strict: bool) -> Expr { + let ts = ident(TRACE_TIMESTAMP_COLUMN); + let mut predicate = if strict { + ts.clone() + .gt_eq(window.source_start()) + .and(ts.lt(window.source_end())) + } else { + ts.clone() + .gt_eq(window.source_start() - interval(CHILD_SPAN_EARLY_NANOS)) + .and(ts.lt(window.source_end() + interval(CHILD_SPAN_LATE_NANOS))) + }; + // A NULL identity component identifies nothing, on either endpoint. + for id in &service.id_columns { + predicate = predicate.and(ident(id).is_not_null()); + } + predicate +} + +/// One trace table's client spans, normalized to +/// `(timestamp, trace_id, span_id, src_id, status_code, duration_nano, +/// virtual_dst, virtual_conn)`. `src_id` is built from the table's `service` +/// declaration, so edges land on exactly the entity ids the registry emits (a +/// composite identity renders the same sorted `k=v` form); the per-table +/// projection is what lets tables with different declarations union. +fn client_spans( + service: &EntityDeclaration, + scan: &DataFrame, + window: &GraphQueryWindow, +) -> DfResult { + // NULLIF('') lets an empty value fall through to the next candidate. + let present: Vec<(Expr, &str)> = { + let schema = scan.schema(); + VIRTUAL_DST_CANDIDATES + .iter() + .filter(|(column, _)| schema.has_column_with_unqualified_name(column)) + .map(|(column, conn)| { + let value = + core_fns::nullif().call(vec![cast(ident(*column), DataType::Utf8), lit("")]); + (value, *conn) + }) + .collect() + }; + let null_utf8 = || lit(ScalarValue::Utf8(None)); + let virtual_dst = if present.is_empty() { + null_utf8() + } else { + core_fns::coalesce().call(present.iter().map(|(value, _)| value.clone()).collect()) + }; + let mut virtual_conn = null_utf8(); + for (value, conn) in present.into_iter().rev() { + virtual_conn = when(value.is_not_null(), lit(conn)).otherwise(virtual_conn)?; + } + + scan.clone() + .filter( + ident(SPAN_KIND_COLUMN) + .eq(lit(SPAN_KIND_CLIENT)) + .and(span_predicate(service, window, true)), + )? + .select(vec![ + ident(TRACE_TIMESTAMP_COLUMN), + ident(TRACE_ID_COLUMN), + ident(SPAN_ID_COLUMN), + // The cast inside entity_id_expr also normalizes tag columns, which + // come out of the storage engine dictionary-encoded. + entity_id_expr(&service.id_columns, &|c| ident(c)).alias("src_id"), + ident(SPAN_STATUS_CODE_COLUMN), + ident(DURATION_NANO_COLUMN), + virtual_dst.alias("virtual_dst"), + virtual_conn.alias("virtual_conn"), + ]) +} + +/// One trace table's server spans, normalized to +/// `(timestamp, trace_id, parent_span_id, dst_id, status_code, duration_nano)`. +fn server_spans( + service: &EntityDeclaration, + trace: DataFrame, + window: &GraphQueryWindow, +) -> DfResult { + trace + .filter( + ident(SPAN_KIND_COLUMN) + .eq(lit(SPAN_KIND_SERVER)) + .and(span_predicate(service, window, false)), + )? + .select(vec![ + ident(TRACE_TIMESTAMP_COLUMN), + ident(TRACE_ID_COLUMN), + ident(PARENT_SPAN_ID_COLUMN), + entity_id_expr(&service.id_columns, &|c| ident(c)).alias("dst_id"), + ident(SPAN_STATUS_CODE_COLUMN), + ident(DURATION_NANO_COLUMN), + ]) +} + +/// The `parent_agent calls agent` derivation over trace tables declaring an +/// `agent` entity: pair each span with its child span (no span-kind filter — +/// agent spans are typically INTERNAL) across all agent-declaring tables, keep +/// pairs whose agent identities differ, and aggregate RED metrics per 60s +/// window. Anchored like the service derivation on the caller: the parent side +/// is bounded to the source window and stamps `observed_at`, the child side +/// widens by the time-proximity allowance and supplies status/duration. +fn agent_calls_branch( + traces: &[CallsSource], + window: &GraphQueryWindow, +) -> DfResult> { + let mut parents: Option = None; + let mut children: Option = None; + for trace in traces { + let Some(agent) = &trace.agent else { + continue; + }; + let parent = trace + .scan + .clone() + .filter(span_predicate(agent, window, true))? + .select(vec![ + ident(TRACE_TIMESTAMP_COLUMN), + ident(TRACE_ID_COLUMN), + ident(SPAN_ID_COLUMN), + entity_id_expr(&agent.id_columns, &|c| ident(c)).alias("src_id"), + ])?; + let child = trace + .scan + .clone() + .filter(span_predicate(agent, window, false))? + .select(vec![ + ident(TRACE_TIMESTAMP_COLUMN), + ident(TRACE_ID_COLUMN), + ident(PARENT_SPAN_ID_COLUMN), + entity_id_expr(&agent.id_columns, &|c| ident(c)).alias("dst_id"), + ident(SPAN_STATUS_CODE_COLUMN), + ident(DURATION_NANO_COLUMN), + ])?; + parents = union_all(parents, parent)?; + children = union_all(children, child)?; + } + let (Some(parents), Some(children)) = (parents, children) else { + return Ok(None); + }; + + let parent = parents.alias("parent")?; + let child = children.alias("child")?; + let join_conditions = vec![ + qcol("parent", TRACE_ID_COLUMN).eq(qcol("child", TRACE_ID_COLUMN)), + qcol("child", PARENT_SPAN_ID_COLUMN).eq(qcol("parent", SPAN_ID_COLUMN)), + qcol("child", TRACE_TIMESTAMP_COLUMN) + .gt_eq(qcol("parent", TRACE_TIMESTAMP_COLUMN) - interval(CHILD_SPAN_EARLY_NANOS)), + qcol("child", TRACE_TIMESTAMP_COLUMN) + .lt_eq(qcol("parent", TRACE_TIMESTAMP_COLUMN) + interval(CHILD_SPAN_LATE_NANOS)), + ]; + + let df = parent + .join_on(child, JoinType::Inner, join_conditions)? + .select(vec![ + bin_ms(qcol("parent", TRACE_TIMESTAMP_COLUMN)).alias("observed_at"), + qcol("parent", "src_id").alias("src_id"), + qcol("child", "dst_id").alias("dst_id"), + qcol("child", SPAN_STATUS_CODE_COLUMN).alias("status_code"), + qcol("child", DURATION_NANO_COLUMN).alias("duration_nano"), + ])? + // A sub-span of the same agent is internal structure, not a call + // between two agents. + .filter(ident("src_id").not_eq(ident("dst_id")))? .aggregate( vec![ident("observed_at"), ident("src_id"), ident("dst_id")], vec![ @@ -216,16 +676,15 @@ fn calls_branch( ident("observed_at").alias("window_start"), (ident("observed_at") + bin_interval()).alias("window_end"), (ident("observed_at") + bin_interval()).alias("fresh_until"), - lit("service").alias("src_type"), + lit("agent").alias("src_type"), ident("src_id"), - lit("service").alias("dst_type"), + lit("agent").alias("dst_type"), ident("dst_id"), lit("calls").alias("rel_type"), lit("trace").alias("provenance"), lit(1.0_f64).alias("confidence"), ident("request_count"), ident("error_count"), - // duration_nano sums in nanoseconds; the contract column is seconds. (cast(ident("duration_nano_sum"), DataType::Float64) / lit(1e9_f64)) .alias("duration_sum"), ident("request_count").alias("duration_count"), @@ -234,74 +693,14 @@ fn calls_branch( Ok(Some(df)) } -/// One trace table's client/server span pairs, projected to the aggregation -/// inputs `(observed_at, src_id, dst_id, status_code, duration_nano)`. -/// Endpoint ids are built from the table's `service` entity declaration, so -/// edges land on exactly the entity ids the registry emits (a composite -/// service identity renders the same sorted `k=v` form on both sides). -fn calls_pairs( - service: &EntityDeclaration, - trace: DataFrame, - window: &GraphQueryWindow, -) -> DfResult { - let mut client_pred = ident(SPAN_KIND_COLUMN) - .eq(lit(SPAN_KIND_CLIENT)) - .and(ident(TRACE_TIMESTAMP_COLUMN).gt_eq(window.source_start())) - .and(ident(TRACE_TIMESTAMP_COLUMN).lt(window.source_end())); - let mut server_pred = ident(SPAN_KIND_COLUMN) - .eq(lit(SPAN_KIND_SERVER)) - // Static bounds implied by the window and the join's time-proximity - // conditions below; the join bounds reference client.timestamp and - // cannot prune the server-side scan. - .and( - ident(TRACE_TIMESTAMP_COLUMN) - .gt_eq(window.source_start() - interval(CHILD_SPAN_EARLY_NANOS)), - ) - .and( - ident(TRACE_TIMESTAMP_COLUMN).lt(window.source_end() + interval(CHILD_SPAN_LATE_NANOS)), - ); - // A NULL identity component identifies nothing, on either endpoint. - for id in &service.id_columns { - client_pred = client_pred.and(ident(id).is_not_null()); - server_pred = server_pred.and(ident(id).is_not_null()); - } - - let client = trace.clone().filter(client_pred)?.alias("client")?; - let server = trace.filter(server_pred)?.alias("server")?; - - let join_conditions = vec![ - qcol("client", TRACE_ID_COLUMN).eq(qcol("server", TRACE_ID_COLUMN)), - qcol("server", PARENT_SPAN_ID_COLUMN).eq(qcol("client", SPAN_ID_COLUMN)), - qcol("server", TRACE_TIMESTAMP_COLUMN) - .gt_eq(qcol("client", TRACE_TIMESTAMP_COLUMN) - interval(CHILD_SPAN_EARLY_NANOS)), - qcol("server", TRACE_TIMESTAMP_COLUMN) - .lt_eq(qcol("client", TRACE_TIMESTAMP_COLUMN) + interval(CHILD_SPAN_LATE_NANOS)), - ]; - - client - .join_on(server, JoinType::Inner, join_conditions)? - .select(vec![ - bin_ms(qcol("client", TRACE_TIMESTAMP_COLUMN)).alias("observed_at"), - // The cast inside entity_id_expr also normalizes tag columns, which - // come out of the storage engine dictionary-encoded. - entity_id_expr(&service.id_columns, &|c| qcol("client", c)).alias("src_id"), - entity_id_expr(&service.id_columns, &|c| qcol("server", c)).alias("dst_id"), - qcol("server", SPAN_STATUS_CODE_COLUMN).alias("status_code"), - qcol("server", DURATION_NANO_COLUMN).alias("duration_nano"), - ])? - // Exclude self-calls on the composed identity: an edge needs two - // distinct service entities. - .filter(ident("src_id").not_eq(ident("dst_id"))) -} - #[cfg(test)] mod tests { use std::sync::Arc; use common_catalog::consts::{SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SERVICE_NAME_COLUMN}; use datafusion::arrow::array::{ - ArrayRef, BinaryArray, Float64Array, Int64Array, StringArray, TimestampMillisecondArray, - TimestampNanosecondArray, UInt64Array, + Array, ArrayRef, BinaryArray, Float64Array, Int64Array, StringArray, + TimestampMillisecondArray, TimestampNanosecondArray, UInt64Array, }; use datafusion::arrow::datatypes::{Field, Schema, TimeUnit}; use datafusion::arrow::record_batch::RecordBatch; @@ -315,9 +714,29 @@ mod tests { GraphQueryWindow::from_observed(0, 10 * 60 * 1000) } - /// A trace-like table in the fixed `greptime_trace_v1` shape (ns timestamps). - fn trace_table_ctx() -> SessionContext { - let schema = Arc::new(Schema::new(vec![ + /// One span row of the fixed `greptime_trace_v1` shape: + /// `(ts_ms, trace_id, span_id, parent_span_id, span_kind, status_code, + /// service_name, duration_nano)`. + type Span<'a> = ( + i64, + &'a str, + &'a str, + Option<&'a str>, + &'a str, + &'a str, + &'a str, + u64, + ); + + /// Registers a trace-v1-shaped table (ns timestamps) plus optional dynamic + /// string columns (`extra`), one value per span. + fn register_trace_table( + ctx: &SessionContext, + name: &str, + extra: &[(&str, &[Option<&str>])], + spans: &[Span<'_>], + ) { + let mut fields = vec![ Field::new( TRACE_TIMESTAMP_COLUMN, DataType::Timestamp(TimeUnit::Nanosecond, None), @@ -329,84 +748,94 @@ mod tests { Field::new(SPAN_KIND_COLUMN, DataType::Utf8, false), Field::new(SPAN_STATUS_CODE_COLUMN, DataType::Utf8, false), Field::new(SERVICE_NAME_COLUMN, DataType::Utf8, false), - Field::new("service_namespace", DataType::Utf8, false), Field::new(DURATION_NANO_COLUMN, DataType::UInt64, false), - ])); + ]; + for (column, _) in extra { + fields.push(Field::new(*column, DataType::Utf8, true)); + } + let schema = Arc::new(Schema::new(fields)); const MS: i64 = 1_000_000; - // Two client->server pairs frontend->cart (one errored), one pair - // cart->cart (self-call, excluded), one unmatched client span. - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(TimestampNanosecondArray::from(vec![ - 1_000 * MS, // client frontend->cart - 1_010 * MS, // server cart - 2_000 * MS, // client frontend->cart (error) - 2_010 * MS, // server cart (error) - 3_000 * MS, // client cart->cart (self-call) - 3_010 * MS, // server cart - 4_000 * MS, // client with no matching server - ])) as ArrayRef, - Arc::new(StringArray::from(vec![ - "t1", "t1", "t2", "t2", "t3", "t3", "t4", - ])), - Arc::new(StringArray::from(vec![ - "c1", "s1", "c2", "s2", "c3", "s3", "c4", - ])), - Arc::new(StringArray::from(vec![ - None, - Some("c1"), - None, - Some("c2"), - None, - Some("c3"), - None, - ])), - Arc::new(StringArray::from(vec![ - "SPAN_KIND_CLIENT", - "SPAN_KIND_SERVER", - "SPAN_KIND_CLIENT", - "SPAN_KIND_SERVER", - "SPAN_KIND_CLIENT", - "SPAN_KIND_SERVER", - "SPAN_KIND_CLIENT", - ])), - Arc::new(StringArray::from(vec![ - "STATUS_CODE_UNSET", - "STATUS_CODE_UNSET", - "STATUS_CODE_UNSET", - "STATUS_CODE_ERROR", - "STATUS_CODE_UNSET", - "STATUS_CODE_UNSET", - "STATUS_CODE_UNSET", - ])), - Arc::new(StringArray::from(vec![ - "frontend", "cart", "frontend", "cart", "cart", "cart", "frontend", - ])), - Arc::new(StringArray::from(vec!["ns1"; 7])), - Arc::new(UInt64Array::from(vec![ - 0, - 500_000_000, // 0.5s - 0, - 1_500_000_000, // 1.5s - 0, - 100, - 0, - ])), - ], - ) - .unwrap(); - let ctx = SessionContext::new(); + let mut columns: Vec = vec![ + Arc::new(TimestampNanosecondArray::from( + spans.iter().map(|s| s.0 * MS).collect::>(), + )), + Arc::new(StringArray::from( + spans.iter().map(|s| s.1).collect::>(), + )), + Arc::new(StringArray::from( + spans.iter().map(|s| s.2).collect::>(), + )), + Arc::new(StringArray::from( + spans.iter().map(|s| s.3).collect::>(), + )), + Arc::new(StringArray::from( + spans.iter().map(|s| s.4).collect::>(), + )), + Arc::new(StringArray::from( + spans.iter().map(|s| s.5).collect::>(), + )), + Arc::new(StringArray::from( + spans.iter().map(|s| s.6).collect::>(), + )), + Arc::new(UInt64Array::from( + spans.iter().map(|s| s.7).collect::>(), + )), + ]; + for (_, values) in extra { + columns.push(Arc::new(StringArray::from(values.to_vec()))); + } + let batch = RecordBatch::try_new(schema.clone(), columns).unwrap(); ctx.register_table( - "opentelemetry_traces", - Arc::new(MemTable::try_new(schema.clone(), vec![vec![batch.clone()]]).unwrap()), - ) - .unwrap(); - ctx.register_table( - "opentelemetry_traces_2", + name, Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), ) .unwrap(); + } + + const UNSET: &str = "STATUS_CODE_UNSET"; + const ERROR: &str = "STATUS_CODE_ERROR"; + const CLIENT: &str = "SPAN_KIND_CLIENT"; + const SERVER: &str = "SPAN_KIND_SERVER"; + + /// Two client->server pairs frontend->cart (one errored), one pair + /// cart->cart (self-call, excluded), one unmatched client span. + const BASE_SPANS: [Span<'static>; 7] = [ + (1_000, "t1", "c1", None, CLIENT, UNSET, "frontend", 0), + ( + 1_010, + "t1", + "s1", + Some("c1"), + SERVER, + UNSET, + "cart", + 500_000_000, + ), + (2_000, "t2", "c2", None, CLIENT, UNSET, "frontend", 0), + ( + 2_010, + "t2", + "s2", + Some("c2"), + SERVER, + ERROR, + "cart", + 1_500_000_000, + ), + (3_000, "t3", "c3", None, CLIENT, UNSET, "cart", 0), + (3_010, "t3", "s3", Some("c3"), SERVER, UNSET, "cart", 100), + (4_000, "t4", "c4", None, CLIENT, UNSET, "frontend", 0), + ]; + + fn trace_table_ctx() -> SessionContext { + let ctx = SessionContext::new(); + let namespaces = [Some("ns1"); 7]; + register_trace_table( + &ctx, + "opentelemetry_traces", + &[("service_namespace", &namespaces)], + &BASE_SPANS, + ); ctx } @@ -428,7 +857,8 @@ mod tests { let trace = ctx.table("opentelemetry_traces").await.unwrap(); let plan = build_relationships_plan_for_test( vec![CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN]), + service: Some(trace_service_decl(&[SERVICE_NAME_COLUMN])), + agent: None, scan: trace, }], &test_window(), @@ -483,22 +913,272 @@ mod tests { assert!(json_texts(batch, 15)[0].is_none()); } + /// RED columns of one output row: `(request_count, error_count, + /// duration_sum, duration_count)`. + fn red_metrics(batch: &RecordBatch, row: usize) -> (i64, i64, f64, i64) { + let i64_at = |column: usize| { + batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap() + .value(row) + }; + let duration_sum = batch + .column(13) + .as_any() + .downcast_ref::() + .unwrap() + .value(row); + (i64_at(11), i64_at(12), duration_sum, i64_at(14)) + } + + fn confidence(batch: &RecordBatch, row: usize) -> f64 { + batch + .column(10) + .as_any() + .downcast_ref::() + .unwrap() + .value(row) + } + + fn sources(scans: &[DataFrame]) -> Vec { + scans + .iter() + .map(|scan| CallsSource { + service: Some(trace_service_decl(&[SERVICE_NAME_COLUMN])), + agent: None, + scan: scan.clone(), + }) + .collect() + } + #[tokio::test] async fn calls_plan_merges_edges_across_trace_tables() { - let ctx = trace_table_ctx(); - let t1 = ctx.table("opentelemetry_traces").await.unwrap(); - let t2 = ctx.table("opentelemetry_traces_2").await.unwrap(); - let plan = build_relationships_plan_for_test( - vec![ - CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN]), - scan: t1, - }, - CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN]), - scan: t2, - }, + let ctx = SessionContext::new(); + // Distinct traces observing the same frontend->cart edge, one per table. + register_trace_table( + &ctx, + "trace_a", + &[], + &[ + (1_000, "t1", "c1", None, CLIENT, UNSET, "frontend", 0), + ( + 1_010, + "t1", + "s1", + Some("c1"), + SERVER, + UNSET, + "cart", + 500_000_000, + ), ], + ); + register_trace_table( + &ctx, + "trace_b", + &[], + &[ + (2_000, "t2", "c2", None, CLIENT, UNSET, "frontend", 0), + ( + 2_010, + "t2", + "s2", + Some("c2"), + SERVER, + ERROR, + "cart", + 1_500_000_000, + ), + ], + ); + let a = ctx.table("trace_a").await.unwrap(); + let b = ctx.table("trace_b").await.unwrap(); + let plan = build_relationships_plan_for_test(sources(&[a, b]), &test_window()) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, 1); + assert_eq!(red_metrics(&batches[0], 0), (2, 1, 2.0, 2)); + } + + #[tokio::test] + async fn calls_pair_split_across_trace_tables() { + let ctx = SessionContext::new(); + // The client span and its child server span land in different tables + // (per-table trace routing): the union-then-join pairing must find it. + register_trace_table( + &ctx, + "trace_a", + &[], + &[(1_000, "t1", "c1", None, CLIENT, UNSET, "frontend", 0)], + ); + register_trace_table( + &ctx, + "trace_b", + &[], + &[( + 1_010, + "t1", + "s1", + Some("c1"), + SERVER, + UNSET, + "cart", + 500_000_000, + )], + ); + let a = ctx.table("trace_a").await.unwrap(); + let b = ctx.table("trace_b").await.unwrap(); + let plan = build_relationships_plan_for_test(sources(&[a, b]), &test_window()) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, 1); + let batch = &batches[0]; + assert_eq!(strings(batch, 5), vec!["frontend"]); + assert_eq!(strings(batch, 7), vec!["cart"]); + assert_eq!(confidence(batch, 0), 1.0); + assert_eq!(red_metrics(batch, 0), (1, 0, 0.5, 1)); + } + + #[tokio::test] + async fn virtual_node_edge_from_unmatched_client() { + let ctx = SessionContext::new(); + register_trace_table( + &ctx, + "trace_a", + &[( + "span_attributes.peer.service", + &[Some("redis"), None, Some("frontend")], + )], + &[ + // Unmatched client naming its peer: a virtual-node edge with + // the client's own status and duration. + ( + 1_000, + "t1", + "c1", + None, + CLIENT, + ERROR, + "frontend", + 250_000_000, + ), + // Unmatched client without a peer attribute: no edge. + (2_000, "t2", "c2", None, CLIENT, UNSET, "frontend", 0), + // Peer attribute naming the client's own service: excluded. + (3_000, "t3", "c3", None, CLIENT, UNSET, "frontend", 0), + ], + ); + let a = ctx.table("trace_a").await.unwrap(); + let plan = build_relationships_plan_for_test(sources(&[a]), &test_window()) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, 1); + let batch = &batches[0]; + assert_eq!(strings(batch, 5), vec!["frontend"]); + assert_eq!(strings(batch, 7), vec!["redis"]); + assert_eq!(strings(batch, 8), vec!["calls"]); + assert_eq!(strings(batch, 9), vec!["trace"]); + assert_eq!(confidence(batch, 0), VIRTUAL_NODE_CONFIDENCE); + assert_eq!(red_metrics(batch, 0), (1, 1, 0.25, 1)); + assert_eq!( + json_texts(batch, 15), + vec![Some(r#"{"connection_type":"virtual_node"}"#.to_string())] + ); + } + + #[tokio::test] + async fn virtual_endpoint_fallback_and_connection_type() { + let ctx = SessionContext::new(); + register_trace_table( + &ctx, + "trace_a", + &[ + // An empty high-precedence value must fall through to the next + // candidate, whose match maps connection_type to `database`. + ("span_attributes.service.peer.name", &[Some("")]), + ("span_attributes.db.namespace", &[Some("mysql")]), + ], + &[(1_000, "t1", "c1", None, CLIENT, UNSET, "frontend", 100)], + ); + let a = ctx.table("trace_a").await.unwrap(); + let plan = build_relationships_plan_for_test(sources(&[a]), &test_window()) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let batch = &batches[0]; + assert_eq!(strings(batch, 7), vec!["mysql"]); + assert_eq!( + json_texts(batch, 15), + vec![Some(r#"{"connection_type":"database"}"#.to_string())] + ); + } + + #[tokio::test] + async fn agent_calls_from_parent_child_spans() { + let ctx = SessionContext::new(); + const INTERNAL: &str = "SPAN_KIND_INTERNAL"; + register_trace_table( + &ctx, + "agent_traces", + &[( + "agent_id", + &[ + Some("orchestrator"), + Some("researcher"), + Some("researcher"), + Some("researcher"), + ], + )], + &[ + // orchestrator delegates to researcher twice (one errored); + // the researcher's own sub-span is same-agent structure. + (1_000, "t1", "p1", None, INTERNAL, UNSET, "app", 0), + ( + 1_010, + "t1", + "a1", + Some("p1"), + INTERNAL, + ERROR, + "app", + 2_000_000_000, + ), + ( + 2_000, + "t1", + "a2", + Some("p1"), + INTERNAL, + UNSET, + "app", + 1_000_000_000, + ), + (1_020, "t1", "a3", Some("a1"), INTERNAL, UNSET, "app", 100), + ], + ); + let scan = ctx.table("agent_traces").await.unwrap(); + let plan = build_relationships_plan( + RelationshipSources { + traces: vec![CallsSource { + service: None, + agent: Some(co_decl("agent", &["agent_id"])), + scan, + }], + co_declared: vec![], + declared: None, + }, &test_window(), ) .unwrap() @@ -506,28 +1186,343 @@ mod tests { let batches = collect(&ctx, plan).await; let total: usize = batches.iter().map(|b| b.num_rows()).sum(); - // The same frontend->cart edge from both tables folds into one row with - // summed RED metrics. + // The same-agent child (a3 under a1) is excluded; agent edges need + // no service declaration. assert_eq!(total, 1); let batch = &batches[0]; - let request_count = batch - .column(11) - .as_any() - .downcast_ref::() + assert_eq!(strings(batch, 4), vec!["agent"]); + assert_eq!(strings(batch, 5), vec!["orchestrator"]); + assert_eq!(strings(batch, 6), vec!["agent"]); + assert_eq!(strings(batch, 7), vec!["researcher"]); + assert_eq!(strings(batch, 8), vec!["calls"]); + assert_eq!(strings(batch, 9), vec!["trace"]); + assert_eq!(confidence(batch, 0), 1.0); + assert_eq!(red_metrics(batch, 0), (2, 1, 3.0, 2)); + } + + #[tokio::test] + async fn agent_calls_anchor_on_the_parent_span() { + let ctx = SessionContext::new(); + const INTERNAL: &str = "SPAN_KIND_INTERNAL"; + // The parent is inside the queried window; the child starts after its + // end but within the proximity allowance and must still pair. + register_trace_table( + &ctx, + "agent_traces", + &[("agent_id", &[Some("orchestrator"), Some("researcher")])], + &[ + (599_000, "t1", "p1", None, INTERNAL, UNSET, "app", 0), + ( + 601_000, + "t1", + "a1", + Some("p1"), + INTERNAL, + UNSET, + "app", + 500_000_000, + ), + ], + ); + let scan = ctx.table("agent_traces").await.unwrap(); + let plan = build_relationships_plan( + RelationshipSources { + traces: vec![CallsSource { + service: None, + agent: Some(co_decl("agent", &["agent_id"])), + scan, + }], + co_declared: vec![], + declared: None, + }, + &test_window(), + ) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, 1); + let batch = &batches[0]; + // observed_at is the parent's bin. + assert_eq!(ts_values(batch, 0), vec![540_000]); + assert_eq!(red_metrics(batch, 0), (1, 0, 0.5, 1)); + } + + /// A metric-like declaring table: ms timestamps, one row with a NULL + /// `instance` (witnessing nothing for instance edges), two rows duplicating + /// the same identities inside one bin (collapsed by DISTINCT). + fn co_declared_ctx() -> SessionContext { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("instance", DataType::Utf8, true), + Field::new("host", DataType::Utf8, false), + Field::new("service", DataType::Utf8, false), + Field::new("agent_id", DataType::Utf8, false), + Field::new("model_name", DataType::Utf8, false), + Field::new("tool_name", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000, 2_000, 3_000])) as ArrayRef, + Arc::new(StringArray::from(vec![Some("i-1"), Some("i-1"), None])), + Arc::new(StringArray::from(vec!["h-1", "h-1", "h-2"])), + Arc::new(StringArray::from(vec!["svc", "svc", "svc"])), + Arc::new(StringArray::from(vec!["agent-1"; 3])), + Arc::new(StringArray::from(vec!["gpt"; 3])), + Arc::new(StringArray::from(vec!["search"; 3])), + ], + ) + .unwrap(); + let ctx = SessionContext::new(); + for name in ["co_metrics", "co_metrics_2"] { + ctx.register_table( + name, + Arc::new(MemTable::try_new(schema.clone(), vec![vec![batch.clone()]]).unwrap()), + ) .unwrap(); - assert_eq!(request_count.value(0), 4); - let error_count = batch - .column(12) - .as_any() - .downcast_ref::() + } + ctx + } + + fn co_decl(entity_type: &str, id_columns: &[&str]) -> EntityDeclaration { + EntityDeclaration { + schema: "public".to_string(), + table: "co_metrics".to_string(), + time_index: "ts".to_string(), + entity_type: entity_type.to_string(), + id_columns: id_columns.iter().map(|s| s.to_string()).collect(), + descriptive_columns: vec![], + scope_columns: vec![], + } + } + + async fn co_declared_edges( + ctx: &SessionContext, + declarations: Vec, + is_trace: bool, + ) -> Option> { + let scan = ctx.table("co_metrics").await.unwrap(); + let plan = build_relationships_plan( + RelationshipSources { + traces: vec![], + co_declared: vec![CoDeclaredSource { + declarations, + is_trace, + scan, + }], + declared: None, + }, + &test_window(), + ) + .unwrap()?; + assert_eq!( + plan.schema() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>(), + RELATIONSHIP_COLUMNS + ); + let batches = collect(ctx, plan).await; + let mut rows = vec![]; + for batch in &batches { + for i in 0..batch.num_rows() { + rows.push(( + strings(batch, 4)[i].clone(), + strings(batch, 5)[i].clone(), + strings(batch, 6)[i].clone(), + strings(batch, 7)[i].clone(), + strings(batch, 8)[i].clone(), + strings(batch, 9)[i].clone(), + )); + assert_eq!(confidence(batch, i), 1.0); + // Co-declared edges carry no RED metrics or attributes. + for column in 11..=15 { + assert!(batch.column(column).is_null(i)); + } + } + } + rows.sort(); + Some(rows) + } + + #[tokio::test] + async fn co_declared_direction_follows_vocabulary() { + let ctx = co_declared_ctx(); + // Declaration order must not matter: the vocabulary fixes direction. + let rows = co_declared_edges( + &ctx, + vec![ + co_decl("host", &["host"]), + co_decl("service", &["service"]), + co_decl("service.instance", &["instance"]), + ], + false, + ) + .await + .unwrap(); + // The two non-NULL-instance rows share one 60s bin, so DISTINCT folds + // them into one row per edge; the NULL-instance row witnesses nothing. + assert_eq!( + rows, + vec![ + ( + "service.instance".to_string(), + "i-1".to_string(), + "host".to_string(), + "h-1".to_string(), + "runs_on".to_string(), + "attribute".to_string(), + ), + ( + "service.instance".to_string(), + "i-1".to_string(), + "service".to_string(), + "svc".to_string(), + "part_of".to_string(), + "attribute".to_string(), + ), + ] + ); + } + + #[tokio::test] + async fn co_declared_edge_folds_across_sources() { + let ctx = co_declared_ctx(); + let declarations = || { + vec![ + co_decl("service.instance", &["instance"]), + co_decl("host", &["host"]), + ] + }; + let source = |scan| CoDeclaredSource { + declarations: declarations(), + is_trace: false, + scan, + }; + let a = ctx.table("co_metrics").await.unwrap(); + let b = ctx.table("co_metrics_2").await.unwrap(); + let plan = build_relationships_plan( + RelationshipSources { + traces: vec![], + co_declared: vec![source(a), source(b)], + declared: None, + }, + &test_window(), + ) + .unwrap() + .unwrap(); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + // Both tables witness the same runs_on edge in the same window: one + // row per (window, edge), not one per source. + assert_eq!(total, 1); + } + + #[tokio::test] + async fn co_declared_ignores_types_outside_the_vocabulary() { + let ctx = co_declared_ctx(); + // (service, host) is not a vocabulary pair: no edge, no plan. + let rows = co_declared_edges( + &ctx, + vec![co_decl("service", &["service"]), co_decl("host", &["host"])], + false, + ) + .await; + assert!(rows.is_none()); + } + + #[tokio::test] + async fn agent_edges_require_a_trace_source() { + let ctx = co_declared_ctx(); + let declarations = || { + vec![ + co_decl("agent", &["agent_id"]), + co_decl("model", &["model_name"]), + co_decl("tool", &["tool_name"]), + ] + }; + // A non-trace table co-declaring agent/model/tool must not fabricate + // invocations. + assert!( + co_declared_edges(&ctx, declarations(), false) + .await + .is_none() + ); + + let rows = co_declared_edges(&ctx, declarations(), true).await.unwrap(); + assert_eq!( + rows, + vec![ + ( + "agent".to_string(), + "agent-1".to_string(), + "model".to_string(), + "gpt".to_string(), + "uses".to_string(), + "trace".to_string(), + ), + ( + "agent".to_string(), + "agent-1".to_string(), + "tool".to_string(), + "search".to_string(), + "invokes".to_string(), + "trace".to_string(), + ), + ] + ); + } + + #[tokio::test] + async fn real_pair_suppresses_virtual_edge() { + let ctx = SessionContext::new(); + register_trace_table( + &ctx, + "trace_a", + &[( + "span_attributes.peer.service", + &[Some("cart"), None, Some("cart")], + )], + &[ + // A sampled-out server: the client names the same peer an + // actual pair witnesses in the same window. + (1_000, "t1", "c1", None, CLIENT, ERROR, "frontend", 999), + ( + 2_010, + "t2", + "s2", + Some("c2"), + SERVER, + UNSET, + "cart", + 500_000_000, + ), + (2_000, "t2", "c2", None, CLIENT, UNSET, "frontend", 0), + ], + ); + let a = ctx.table("trace_a").await.unwrap(); + let plan = build_relationships_plan_for_test(sources(&[a]), &test_window()) + .unwrap() .unwrap(); - assert_eq!(error_count.value(0), 2); - let duration_sum = batch - .column(13) - .as_any() - .downcast_ref::() - .unwrap(); - assert!((duration_sum.value(0) - 4.0).abs() < 1e-9); + + let batches = collect(&ctx, plan).await; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + // One (window, edge) row: the pair population wins; the unmatched + // client neither adds a second row nor inflates the pair RED metrics. + assert_eq!(total, 1); + let batch = &batches[0]; + assert_eq!(strings(batch, 7), vec!["cart"]); + assert_eq!(confidence(batch, 0), 1.0); + assert_eq!(red_metrics(batch, 0), (1, 0, 0.5, 1)); + assert!(json_texts(batch, 15)[0].is_none()); } #[tokio::test] @@ -536,7 +1531,11 @@ mod tests { let trace = ctx.table("opentelemetry_traces").await.unwrap(); let plan = build_relationships_plan_for_test( vec![CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN, "service_namespace"]), + service: Some(trace_service_decl(&[ + SERVICE_NAME_COLUMN, + "service_namespace", + ])), + agent: None, scan: trace, }], &test_window(), @@ -564,7 +1563,14 @@ mod tests { traces: Vec, window: &GraphQueryWindow, ) -> DfResult> { - build_relationships_plan(traces, None, window) + build_relationships_plan( + RelationshipSources { + traces, + co_declared: vec![], + declared: None, + }, + window, + ) } /// Registers a declared-edge table holding: @@ -720,9 +1726,16 @@ mod tests { // Queried window [120s, 600s). let window = GraphQueryWindow::from_observed(120_000, 600_000); - let plan = build_relationships_plan(vec![], Some(DeclaredSource { scan }), &window) - .unwrap() - .unwrap(); + let plan = build_relationships_plan( + RelationshipSources { + traces: vec![], + co_declared: vec![], + declared: Some(DeclaredSource { scan }), + }, + &window, + ) + .unwrap() + .unwrap(); let names = plan .schema() .fields() @@ -836,9 +1849,16 @@ mod tests { // the first revision was the edge's state then and must not be hidden // by the newer one. let window = GraphQueryWindow::from_observed(0, 150_000); - let plan = build_relationships_plan(vec![], Some(DeclaredSource { scan }), &window) - .unwrap() - .unwrap(); + let plan = build_relationships_plan( + RelationshipSources { + traces: vec![], + co_declared: vec![], + declared: Some(DeclaredSource { scan }), + }, + &window, + ) + .unwrap() + .unwrap(); let batches = collect(&ctx, plan).await; let mut rows: Vec<(String, f64)> = vec![]; @@ -879,11 +1899,15 @@ mod tests { .unwrap(); let plan = build_relationships_plan( - vec![CallsSource { - service: trace_service_decl(&[SERVICE_NAME_COLUMN]), - scan: trace, - }], - Some(DeclaredSource { scan: declared }), + RelationshipSources { + traces: vec![CallsSource { + service: Some(trace_service_decl(&[SERVICE_NAME_COLUMN])), + agent: None, + scan: trace, + }], + co_declared: vec![], + declared: Some(DeclaredSource { scan: declared }), + }, &test_window(), ) .unwrap() diff --git a/tests-integration/src/instance.rs b/tests-integration/src/instance.rs index 547f64436f6..08d2a1aff15 100644 --- a/tests-integration/src/instance.rs +++ b/tests-integration/src/instance.rs @@ -563,7 +563,9 @@ mod tests { ) -> auth::error::Result { if matches!(req, auth::PermissionReq::Action(auth::SEMANTIC_GRAPH_QUERY)) && let auth::PermissionTableTargets::Resolved(targets) = targets - && targets.iter().any(|target| target.table == "secret_traces") + && targets + .iter() + .any(|target| target.table.starts_with("secret_")) { return Ok(auth::PermissionResp::Reject); } @@ -580,7 +582,7 @@ mod tests { .await; let instance = standalone.fe_instance().clone(); - for table in ["open_traces", "secret_traces"] { + for (table, trace_id) in [("open_traces", "t0"), ("secret_traces", "t1")] { let create = format!( r#"create table {table} ( "timestamp" timestamp(9) time index, @@ -597,8 +599,8 @@ mod tests { create_table(&instance, &create).await; let insert = format!( "insert into {table} values \ - (now(), 't0', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'client-{table}', 0), \ - (now(), 't0', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'server-{table}', 100)" + (now(), '{trace_id}', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'client-{table}', 0), \ + (now(), '{trace_id}', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'server-{table}', 100)" ); let output = query(&instance, &insert).await; let OutputData::AffectedRows(x) = output.data else { @@ -607,6 +609,51 @@ mod tests { assert_eq!(x, 2); } + // A pair split across tables (client here, server in a denied table): + // a join-derived edge needs read access to all of its input tables. + create_table( + &instance, + r#"create table cross_clients ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + primary key (service_name) + ) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true')"#, + ) + .await; + create_table( + &instance, + r#"create table secret_cross_servers ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + primary key (service_name) + ) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true')"#, + ) + .await; + for insert in [ + "insert into cross_clients values \ + (now(), 't9', 'c9', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'client-cross', 0)", + "insert into secret_cross_servers values \ + (now(), 't9', 's9', 'c9', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'server-cross', 100)", + ] { + let output = query(&instance, insert).await; + let OutputData::AffectedRows(x) = output.data else { + unreachable!() + }; + assert_eq!(x, 1); + } + let sql = "select src_id, dst_id from greptime_private.semantic_relationships \ order by src_id"; let output = query(&instance, sql).await; @@ -623,6 +670,10 @@ mod tests { !pretty_print.contains("secret_traces"), "denied source leaked into edges:\n{pretty_print}" ); + assert!( + !pretty_print.contains("client-cross"), + "edge with a denied join side leaked:\n{pretty_print}" + ); let sql = "select entity_id from greptime_private.semantic_entities order by entity_id"; let output = query(&instance, sql).await; @@ -636,7 +687,7 @@ mod tests { "allowed entity missing:\n{pretty_print}" ); assert!( - !pretty_print.contains("secret_traces"), + !pretty_print.contains("secret_traces") && !pretty_print.contains("server-cross"), "denied source leaked into entities:\n{pretty_print}" ); } diff --git a/tests/cases/standalone/common/system/semantic_graph.result b/tests/cases/standalone/common/system/semantic_graph.result index 73f00f57c6f..6617d9f839d 100644 --- a/tests/cases/standalone/common/system/semantic_graph.result +++ b/tests/cases/standalone/common/system/semantic_graph.result @@ -277,3 +277,204 @@ drop table greptime_private.semantic_relationships_declared; Affected Rows: 0 +-- Same-row co-declaration: a table declaring both entity types of a built-in +-- vocabulary pair witnesses the edge on every row carrying both identities; +-- the vocabulary fixes the direction. +create table graph_pod_metrics ( + ts timestamp time index, + instance string, + host string, + "service" string, + cpu double, + primary key (instance, host, "service") +) with ( + 'greptime.semantic.entity.service.instance.id' = 'instance', + 'greptime.semantic.entity.host.id' = 'host', + 'greptime.semantic.entity.service.id' = 'service' +); + +Affected Rows: 0 + +insert into graph_pod_metrics values (now(), 'cart-0', 'node-1', 'cart', 0.5); + +Affected Rows: 1 + +select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence +from greptime_private.semantic_relationships +order by rel_type, src_id; + ++------------------+--------+----------+--------+----------+------------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | confidence | ++------------------+--------+----------+--------+----------+------------+------------+ +| service.instance | cart-0 | service | cart | part_of | attribute | 1.0 | +| service.instance | cart-0 | host | node-1 | runs_on | attribute | 1.0 | ++------------------+--------+----------+--------+----------+------------+------------+ + +drop table graph_pod_metrics; + +Affected Rows: 0 + +-- Cross-table calls pairing: the client span and its child server span land in +-- different trace tables and must still pair into one edge. A trace-model +-- table that lost the fixed span columns is skipped instead of failing the +-- scan. +create table graph_traces_a ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +Affected Rows: 0 + +create table graph_traces_b ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +Affected Rows: 0 + +create table graph_traces_malformed ( + ts timestamp time index, + note string +) with ('table_data_model' = 'greptime_trace_v1'); + +Affected Rows: 0 + +insert into graph_traces_a values + (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0); + +Affected Rows: 1 + +insert into graph_traces_b values + (now(), 't1', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_ERROR', 'cart', 1500000000); + +Affected Rows: 1 + +insert into graph_traces_malformed values (now(), 'not a trace'); + +Affected Rows: 1 + +select src_id, dst_id, rel_type, provenance, confidence, request_count, error_count +from greptime_private.semantic_relationships +order by src_id; + ++----------+--------+----------+------------+------------+---------------+-------------+ +| src_id | dst_id | rel_type | provenance | confidence | request_count | error_count | ++----------+--------+----------+------------+------------+---------------+-------------+ +| frontend | cart | calls | trace | 1.0 | 1 | 1 | ++----------+--------+----------+------------+------------+---------------+-------------+ + +drop table graph_traces_a; + +Affected Rows: 0 + +drop table graph_traces_b; + +Affected Rows: 0 + +drop table graph_traces_malformed; + +Affected Rows: 0 + +-- Virtual nodes: a client span with no matching server span names its peer +-- through span attributes; the edge carries confidence < 1.0 and the +-- connection type. +create table graph_traces_virtual ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + "span_attributes.peer.service" string, + "span_attributes.db.name" string, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +Affected Rows: 0 + +insert into graph_traces_virtual values + (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 250000000, 'redis', NULL), + (now(), 't2', 'c2', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 100000000, NULL, 'orders-db'); + +Affected Rows: 2 + +-- SQLNESS PROTOCOL MYSQL +select src_id, dst_id, rel_type, confidence, request_count, attributes +from greptime_private.semantic_relationships +order by dst_id; + ++----------+-----------+----------+------------+---------------+------------------------------------+ +| src_id | dst_id | rel_type | confidence | request_count | attributes | ++----------+-----------+----------+------------+---------------+------------------------------------+ +| frontend | orders-db | calls | 0.5 | 1 | {"connection_type":"database"} | +| frontend | redis | calls | 0.5 | 1 | {"connection_type":"virtual_node"} | ++----------+-----------+----------+------------+---------------+------------------------------------+ + +drop table graph_traces_virtual; + +Affected Rows: 0 + +-- Agent edges: span structure derives parent_agent calls agent, and span rows +-- co-declaring agent+model / agent+tool witness uses / invokes. +create table graph_agent_traces ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + agent_id string, + model_name string, + tool_name string, + primary key (service_name, agent_id, model_name, tool_name) +) with ( + 'table_data_model' = 'greptime_trace_v1', + 'append_mode' = 'true', + 'greptime.semantic.entity.agent.id' = 'agent_id', + 'greptime.semantic.entity.model.id' = 'model_name', + 'greptime.semantic.entity.tool.id' = 'tool_name' +); + +Affected Rows: 0 + +insert into graph_agent_traces values + (now(), 't1', 'p1', NULL, 'SPAN_KIND_INTERNAL', 'STATUS_CODE_UNSET', 'app', 0, 'orchestrator', NULL, NULL), + (now(), 't1', 'a1', 'p1', 'SPAN_KIND_INTERNAL', 'STATUS_CODE_UNSET', 'app', 2000000000, 'researcher', 'gpt-5', NULL), + (now(), 't1', 'a2', 'a1', 'SPAN_KIND_INTERNAL', 'STATUS_CODE_UNSET', 'app', 500000000, 'researcher', NULL, 'web_search'); + +Affected Rows: 3 + +select src_type, src_id, dst_type, dst_id, rel_type, provenance +from greptime_private.semantic_relationships +order by rel_type, dst_id; + ++----------+--------------+----------+------------+----------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | ++----------+--------------+----------+------------+----------+------------+ +| agent | orchestrator | agent | researcher | calls | trace | +| agent | researcher | tool | web_search | invokes | trace | +| agent | researcher | model | gpt-5 | uses | trace | ++----------+--------------+----------+------------+----------+------------+ + +drop table graph_agent_traces; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/system/semantic_graph.sql b/tests/cases/standalone/common/system/semantic_graph.sql index c5a93d6c6b7..9acb003331a 100644 --- a/tests/cases/standalone/common/system/semantic_graph.sql +++ b/tests/cases/standalone/common/system/semantic_graph.sql @@ -161,3 +161,140 @@ values (now(), 'service', 'reborn', 'depends_on', 'service', 'db', 'declared', ' select src_id from greptime_private.semantic_relationships order by src_id; drop table greptime_private.semantic_relationships_declared; + +-- Same-row co-declaration: a table declaring both entity types of a built-in +-- vocabulary pair witnesses the edge on every row carrying both identities; +-- the vocabulary fixes the direction. +create table graph_pod_metrics ( + ts timestamp time index, + instance string, + host string, + "service" string, + cpu double, + primary key (instance, host, "service") +) with ( + 'greptime.semantic.entity.service.instance.id' = 'instance', + 'greptime.semantic.entity.host.id' = 'host', + 'greptime.semantic.entity.service.id' = 'service' +); + +insert into graph_pod_metrics values (now(), 'cart-0', 'node-1', 'cart', 0.5); + +select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence +from greptime_private.semantic_relationships +order by rel_type, src_id; + +drop table graph_pod_metrics; + +-- Cross-table calls pairing: the client span and its child server span land in +-- different trace tables and must still pair into one edge. A trace-model +-- table that lost the fixed span columns is skipped instead of failing the +-- scan. +create table graph_traces_a ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +create table graph_traces_b ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +create table graph_traces_malformed ( + ts timestamp time index, + note string +) with ('table_data_model' = 'greptime_trace_v1'); + +insert into graph_traces_a values + (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0); + +insert into graph_traces_b values + (now(), 't1', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_ERROR', 'cart', 1500000000); + +insert into graph_traces_malformed values (now(), 'not a trace'); + +select src_id, dst_id, rel_type, provenance, confidence, request_count, error_count +from greptime_private.semantic_relationships +order by src_id; + +drop table graph_traces_a; + +drop table graph_traces_b; + +drop table graph_traces_malformed; + +-- Virtual nodes: a client span with no matching server span names its peer +-- through span attributes; the edge carries confidence < 1.0 and the +-- connection type. +create table graph_traces_virtual ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + "span_attributes.peer.service" string, + "span_attributes.db.name" string, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +insert into graph_traces_virtual values + (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 250000000, 'redis', NULL), + (now(), 't2', 'c2', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 100000000, NULL, 'orders-db'); + +-- SQLNESS PROTOCOL MYSQL +select src_id, dst_id, rel_type, confidence, request_count, attributes +from greptime_private.semantic_relationships +order by dst_id; + +drop table graph_traces_virtual; + +-- Agent edges: span structure derives parent_agent calls agent, and span rows +-- co-declaring agent+model / agent+tool witness uses / invokes. +create table graph_agent_traces ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + agent_id string, + model_name string, + tool_name string, + primary key (service_name, agent_id, model_name, tool_name) +) with ( + 'table_data_model' = 'greptime_trace_v1', + 'append_mode' = 'true', + 'greptime.semantic.entity.agent.id' = 'agent_id', + 'greptime.semantic.entity.model.id' = 'model_name', + 'greptime.semantic.entity.tool.id' = 'tool_name' +); + +insert into graph_agent_traces values + (now(), 't1', 'p1', NULL, 'SPAN_KIND_INTERNAL', 'STATUS_CODE_UNSET', 'app', 0, 'orchestrator', NULL, NULL), + (now(), 't1', 'a1', 'p1', 'SPAN_KIND_INTERNAL', 'STATUS_CODE_UNSET', 'app', 2000000000, 'researcher', 'gpt-5', NULL), + (now(), 't1', 'a2', 'a1', 'SPAN_KIND_INTERNAL', 'STATUS_CODE_UNSET', 'app', 500000000, 'researcher', NULL, 'web_search'); + +select src_type, src_id, dst_type, dst_id, rel_type, provenance +from greptime_private.semantic_relationships +order by rel_type, dst_id; + +drop table graph_agent_traces; From 3b3a9032e06faff55819701e2836251ceb90206f Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Wed, 12 Aug 2026 04:59:10 +0000 Subject: [PATCH 31/92] feat(servers): expose native histograms over Prometheus HTTP (#8850) * feat(servers): expose native histograms over Prometheus HTTP Signed-off-by: shuiyisong * fix(servers): refine Prometheus HTTP metadata handling Signed-off-by: shuiyisong * test(servers): expand Prometheus metadata coverage Signed-off-by: shuiyisong * fix(servers): return OpenMetrics units from metadata API Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/servers/src/grpc/prom_query_gateway.rs | 9 +- src/servers/src/http.rs | 5 +- src/servers/src/http/prometheus.rs | 490 +++++++++++++++++- .../src/http/result/prometheus_resp.rs | 79 ++- src/servers/src/otlp/metrics.rs | 1 + src/servers/src/otlp/metrics/translator.rs | 14 + tests-integration/tests/grpc.rs | 1 + tests-integration/tests/http.rs | 218 ++++++++ 8 files changed, 795 insertions(+), 22 deletions(-) diff --git a/src/servers/src/grpc/prom_query_gateway.rs b/src/servers/src/grpc/prom_query_gateway.rs index 6f4151ce683..99940ad9a24 100644 --- a/src/servers/src/grpc/prom_query_gateway.rs +++ b/src/servers/src/grpc/prom_query_gateway.rs @@ -134,12 +134,19 @@ impl PrometheusGatewayService { } }; let (metric_name, mut result_type) = retrieve_metric_name_and_result_type(query.expr()); + let query_id = ctx.remote_query_id().map(str::to_string); let result = self.handler.do_query_parsed(query, ctx).await; // range query only returns matrix if is_range_query { result_type = ValueType::Matrix; }; - PrometheusJsonResponse::from_query_result(result, metric_name, result_type).await + PrometheusJsonResponse::from_query_result( + result, + metric_name, + result_type, + query_id.as_deref(), + ) + .await } } diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index d52bb307012..bf7478eacfd 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -60,8 +60,8 @@ use crate::http::influxdb::{influxdb_health, influxdb_ping, influxdb_write_v1, i use crate::http::otlp::OtlpState; use crate::http::prom_store::PromStoreState; use crate::http::prometheus::{ - build_info_query, format_query, instant_query, label_values_query, labels_query, parse_query, - range_query, series_query, + build_info_query, format_query, instant_query, label_values_query, labels_query, + metadata_query, parse_query, range_query, series_query, }; use crate::http::result::arrow_result::ArrowResponse; use crate::http::result::csv_result::CsvResponse; @@ -1341,6 +1341,7 @@ impl HttpServer { .route("/query", routing::post(instant_query).get(instant_query)) .route("/query_range", routing::post(range_query).get(range_query)) .route("/labels", routing::post(labels_query).get(labels_query)) + .route("/metadata", routing::get(metadata_query)) .route("/series", routing::post(series_query).get(series_query)) .route("/parse_query", routing::post(parse_query).get(parse_query)) .route( diff --git a/src/servers/src/http/prometheus.rs b/src/servers/src/http/prometheus.rs index 708eba652dc..218a5ae68f5 100644 --- a/src/servers/src/http/prometheus.rs +++ b/src/servers/src/http/prometheus.rs @@ -34,6 +34,7 @@ use common_catalog::parse_catalog_and_schema_from_db_string; use common_decimal::Decimal128; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; +use common_query::native_histogram::is_native_histogram_value_type; use common_query::{Output, OutputData}; use common_recordbatch::{RecordBatch, RecordBatches}; use common_telemetry::{debug, tracing}; @@ -64,6 +65,8 @@ use store_api::metric_engine_consts::{ DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY, }; use table::TableRef; +use table::metadata::TableInfo; +use table::requests::{SEMANTIC_METRIC_TEMPORALITY, SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT}; pub use super::result::prometheus_resp::PrometheusJsonResponse; use crate::error::{ @@ -71,6 +74,7 @@ use crate::error::{ NotSupportedSnafu, Result, TableNotFoundSnafu, UnexpectedResultSnafu, }; use crate::http::header::collect_plan_metrics; +use crate::otlp::metrics::ucum_to_openmetrics_unit; use crate::prom_store::{FIELD_NAME_LABEL, METRIC_NAME_LABEL, is_database_selection_label}; use crate::prometheus_handler::{ ParsedPromQuery, PrometheusHandlerRef, resolve_schema_from_matchers, @@ -136,6 +140,15 @@ pub struct PromData { pub result: PromQueryResult, } +/// Metadata for a Prometheus metric family. +#[derive(Debug, Default, Serialize, Deserialize, PartialEq, Eq)] +pub struct PromMetadata { + #[serde(rename = "type")] + pub metric_type: String, + pub unit: String, + pub help: String, +} + /// A "holder" for the reference([Arc]) to a column name, /// to help avoiding cloning [String]s when used as a [HashMap] key. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] @@ -153,6 +166,7 @@ pub enum PrometheusResponse { PromData(PromData), Labels(Vec), Series(Vec>), + Metadata(BTreeMap>), LabelValues(Vec), FormatQuery(String), BuildInfo(OwnedBuildInfo), @@ -166,7 +180,7 @@ impl PrometheusResponse { /// Append the other [`PrometheusResponse]`. /// # NOTE /// Only append matrix and vector results, otherwise just ignore the other response. - fn append(&mut self, other: PrometheusResponse) { + pub(super) fn append(&mut self, other: PrometheusResponse) { match (self, other) { ( PrometheusResponse::PromData(PromData { @@ -209,6 +223,14 @@ pub struct FormatQuery { query: Option, } +/// Query parameters for the Prometheus metric metadata endpoint. +#[derive(Debug, Default, Serialize, Deserialize)] +pub struct MetadataQuery { + db: Option, + limit: Option, + metric: Option, +} + #[axum_macros::debug_handler] #[tracing::instrument( skip_all, @@ -246,6 +268,67 @@ pub async fn build_info_query() -> PrometheusJsonResponse { PrometheusJsonResponse::success(PrometheusResponse::BuildInfo(build_info.into())) } +#[axum_macros::debug_handler] +#[tracing::instrument( + skip_all, + fields(protocol = "prometheus", request_type = "metadata_query") +)] +pub async fn metadata_query( + State(handler): State, + Query(params): Query, + Extension(mut query_ctx): Extension, +) -> PrometheusJsonResponse { + let (catalog, schema) = get_catalog_schema(¶ms.db, &query_ctx); + try_update_catalog_schema(&mut query_ctx, &catalog, &schema); + let query_ctx = Arc::new(query_ctx); + + if let Err(err) = handler.check_query_permission(&[], &query_ctx).await { + return PrometheusJsonResponse::error(err.status_code(), err.to_string()); + } + if let Some(metric) = ¶ms.metric + && let Err(err) = handler + .check_query_target_permission( + current_schema_metric_targets(&query_ctx, std::slice::from_ref(metric)), + &query_ctx, + ) + .await + { + return PrometheusJsonResponse::error(err.status_code(), err.to_string()); + } + + let mut metadata = + match retrieve_metric_metadata(&query_ctx, handler.catalog_manager(), ¶ms).await { + Ok(metadata) => metadata, + Err(err) => { + return PrometheusJsonResponse::error( + StatusCode::InvalidArguments, + err.to_string(), + ); + } + }; + let allowed_metric_names = match handler + .filter_metadata_metric_names( + metadata.keys().cloned().collect(), + query_ctx.current_schema().as_str(), + &query_ctx, + ) + .await + { + Ok(metric_names) => metric_names, + Err(err) => { + return PrometheusJsonResponse::error(err.status_code(), err.to_string()); + } + } + .into_iter() + .collect::>(); + metadata.retain(|metric, _| allowed_metric_names.contains(metric)); + if let Some(limit) = params.limit { + metadata = metadata.into_iter().take(limit).collect(); + } + + PrometheusJsonResponse::success(PrometheusResponse::Metadata(metadata)) +} + #[derive(Debug, Default, Serialize, Deserialize)] pub struct InstantQuery { query: Option, @@ -391,7 +474,7 @@ pub async fn instant_query( responses .into_iter() .reduce(|mut acc, resp| { - acc.data.append(resp.data); + acc.append_query_response(resp); acc }) .unwrap() @@ -407,8 +490,10 @@ async fn do_instant_query( query_ctx: QueryContextRef, ) -> PrometheusJsonResponse { let (metric_name, result_type) = retrieve_metric_name_and_result_type(prom_query.expr()); + let query_id = query_ctx.remote_query_id().map(str::to_string); let result = handler.do_query_parsed(prom_query, query_ctx).await; - PrometheusJsonResponse::from_query_result(result, metric_name, result_type).await + PrometheusJsonResponse::from_query_result(result, metric_name, result_type, query_id.as_deref()) + .await } #[derive(Debug, Default, Serialize, Deserialize)] @@ -519,7 +604,7 @@ pub async fn range_query( responses .into_iter() .reduce(|mut acc, resp| { - acc.data.append(resp.data); + acc.append_query_response(resp); acc }) .unwrap() @@ -535,8 +620,15 @@ async fn do_range_query( query_ctx: QueryContextRef, ) -> PrometheusJsonResponse { let (metric_name, _) = retrieve_metric_name_and_result_type(prom_query.expr()); + let query_id = query_ctx.remote_query_id().map(str::to_string); let result = handler.do_query_parsed(prom_query, query_ctx).await; - PrometheusJsonResponse::from_query_result(result, metric_name, ValueType::Matrix).await + PrometheusJsonResponse::from_query_result( + result, + metric_name, + ValueType::Matrix, + query_id.as_deref(), + ) + .await } #[derive(Debug, Default, Serialize)] @@ -1165,15 +1257,15 @@ fn record_batches_to_labels_name( labels: &mut HashSet, ) -> Result<()> { let mut column_indices = Vec::new(); - let mut field_column_indices = Vec::new(); + let mut value_column_indices = Vec::new(); for (i, column) in batches.schema().column_schemas().iter().enumerate() { - if let ConcreteDataType::Float64(_) = column.data_type { - field_column_indices.push(i); + if is_prometheus_value_column(&column.data_type) { + value_column_indices.push(i); } column_indices.push(i); } - if field_column_indices.is_empty() { + if value_column_indices.is_empty() { return Err(Error::Internal { err_msg: "no value column found".to_string(), }); @@ -1185,21 +1277,18 @@ fn record_batches_to_labels_name( .map(|c| batches.schema().column_name_by_index(*c).to_string()) .collect::>(); - let field_columns = field_column_indices + let value_columns = value_column_indices .iter() - .map(|i| { - let column = batch.column(*i); - column.as_primitive::() - }) + .map(|i| batch.column(*i)) .collect::>(); for row_index in 0..batch.num_rows() { - // if all field columns are null, skip this row - if field_columns.iter().all(|c| c.is_null(row_index)) { + // if all value columns are null, skip this row + if value_columns.iter().all(|c| c.is_null(row_index)) { continue; } - // if a field is not null, record the tag name and return + // if a value is not null, record the tag name and return names.iter().for_each(|name| { let _ = labels.insert(name.clone()); }); @@ -1209,6 +1298,10 @@ fn record_batches_to_labels_name( Ok(()) } +fn is_prometheus_value_column(data_type: &ConcreteDataType) -> bool { + matches!(data_type, ConcreteDataType::Float64(_)) || is_native_histogram_value_type(data_type) +} + pub(crate) fn retrieve_metric_name_and_result_type( promql_expr: &PromqlExpr, ) -> (Option, ValueType) { @@ -1869,6 +1962,107 @@ async fn retrieve_table_names( Ok(table_names) } +async fn retrieve_metric_metadata( + query_ctx: &QueryContext, + manager: CatalogManagerRef, + params: &MetadataQuery, +) -> Result>> { + let mut metadata = BTreeMap::new(); + if params.limit == Some(0) { + return Ok(metadata); + } + + let catalog = query_ctx.current_catalog(); + let schema = query_ctx.current_schema(); + + if let Some(metric) = ¶ms.metric { + let Some(table) = manager + .table(catalog, &schema, metric, Some(query_ctx)) + .await? + else { + return Ok(metadata); + }; + let table_info = table.table_info(); + if is_prometheus_metric_table(table_info.as_ref()) { + metadata.insert( + table_info.name.clone(), + vec![prometheus_metadata_from_table(table_info.as_ref())], + ); + } + return Ok(metadata); + } + + let mut tables_stream = manager.tables(catalog, &schema, Some(query_ctx)); + + while let Some(table) = tables_stream.next().await { + let table = table?; + let table_info = table.table_info(); + if !is_prometheus_metric_table(table_info.as_ref()) { + continue; + } + + metadata.insert( + table_info.name.clone(), + vec![prometheus_metadata_from_table(table_info.as_ref())], + ); + } + + Ok(metadata) +} + +fn is_prometheus_metric_table(table_info: &TableInfo) -> bool { + table_info + .meta + .options + .extra_options + .contains_key(LOGICAL_TABLE_METADATA_KEY) +} + +fn prometheus_metadata_from_table(table_info: &TableInfo) -> PromMetadata { + let options = &table_info.meta.options.extra_options; + let metric_type = match options.get(SEMANTIC_METRIC_TYPE) { + Some(metric_type) + if options + .get(SEMANTIC_METRIC_TEMPORALITY) + .is_some_and(|temporality| temporality == "delta") + && matches!( + metric_type.as_str(), + "counter" | "histogram" | "updown_counter" + ) => + { + "unknown".to_string() + } + Some(metric_type) => match metric_type.as_str() { + "updown_counter" => "gauge".to_string(), + "gauge_histogram" => "gaugehistogram".to_string(), + "mixed" => "unknown".to_string(), + metric_type => metric_type.to_string(), + }, + None if table_has_native_histogram_value(table_info) => "histogram".to_string(), + None => String::new(), + }; + let unit = options + .get(SEMANTIC_METRIC_UNIT) + .map(|unit| ucum_to_openmetrics_unit(unit)) + .unwrap_or_default(); + + PromMetadata { + metric_type, + unit, + // TODO: Persist and return Prometheus help text and OTLP metric descriptions. + help: String::new(), + } +} + +fn table_has_native_histogram_value(table_info: &TableInfo) -> bool { + table_info + .meta + .schema + .column_schemas() + .iter() + .any(|column| is_native_histogram_value_type(&column.data_type)) +} + async fn retrieve_field_names( query_ctx: &QueryContext, manager: CatalogManagerRef, @@ -2170,11 +2364,18 @@ mod tests { use catalog::memory::MemoryCatalogManager; use catalog::{RegisterSchemaRequest, RegisterTableRequest}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_query::native_histogram::{ + CounterResetHint, NativeHistogram, build_histogram_array, native_histogram_value_type, + }; + use common_query::prelude::greptime_native_histogram; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; + use datatypes::vectors::{StringVector, StructVector}; use promql_parser::parser::value::ValueType; use table::metadata::{TableInfoBuilder, TableMetaBuilder, TableType, TableVersion}; - use table::requests::TableOptions; + use table::requests::{ + SEMANTIC_METRIC_TEMPORALITY, SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, TableOptions, + }; use table::test_util::EmptyTable; use table::test_util::table_info::test_table_info; @@ -2191,6 +2392,7 @@ mod tests { struct TestPrometheusHandler { catalog_manager: CatalogManagerRef, + deny_operation: bool, denied_table: Option<&'static str>, metric_names: Vec, queries: Mutex>, @@ -2215,6 +2417,11 @@ mod tests { } async fn check_query_permission(&self, _: &[PromQuery], _: &QueryContextRef) -> Result<()> { + if self.deny_operation { + return auth::error::PermissionDeniedSnafu + .fail() + .context(crate::error::AuthSnafu); + } Ok(()) } @@ -2320,6 +2527,7 @@ mod tests { async fn test_promql_timer_records_parse_errors() { let handler: PrometheusHandlerRef = Arc::new(TestPrometheusHandler { catalog_manager: MemoryCatalogManager::new(), + deny_operation: false, denied_table: None, metric_names: Vec::new(), queries: Mutex::new(Vec::new()), @@ -2402,6 +2610,7 @@ mod tests { let response = label_values_query( State(Arc::new(TestPrometheusHandler { catalog_manager: manager, + deny_operation: false, denied_table: Some("denied"), metric_names: Vec::new(), queries: Mutex::new(Vec::new()), @@ -2444,6 +2653,7 @@ mod tests { let handler = Arc::new(TestPrometheusHandler { catalog_manager: manager, + deny_operation: false, denied_table: None, metric_names: vec!["cpu_user".to_string(), "cpu_system".to_string()], queries: Mutex::new(Vec::new()), @@ -3146,4 +3356,248 @@ mod tests { (PermissionTableTargets::resolved(Vec::new()), 2) ); } + + #[test] + fn test_record_batches_to_labels_name_accepts_native_histogram_value() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new("host", ConcreteDataType::string_datatype(), false), + ColumnSchema::new( + greptime_native_histogram(), + native_histogram_value_type().clone(), + true, + ), + ])); + let batch = RecordBatch::new_empty(schema.clone()); + let batches = RecordBatches::try_new(schema.clone(), vec![batch]).unwrap(); + let mut labels = HashSet::new(); + + record_batches_to_labels_name(batches, &mut labels).unwrap(); + + assert!(labels.is_empty()); + + let histogram = NativeHistogram { + schema: 0, + zero_threshold: 0.001, + sum: 1.0, + reset_hint: CounterResetHint::Unknown, + start_timestamp: None, + custom_values: vec![], + positive_spans: vec![], + negative_spans: vec![], + count: 1.0, + zero_count: 1.0, + positive_buckets: vec![], + negative_buckets: vec![], + }; + let histogram_array = build_histogram_array(&[Some(histogram)]); + let histogram_array = histogram_array + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + let ConcreteDataType::Struct(histogram_type) = native_histogram_value_type().clone() else { + unreachable!("native histogram type must be a struct") + }; + let batch = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(StringVector::from(vec![Some("localhost")])) as _, + Arc::new(StructVector::try_new(histogram_type, histogram_array).unwrap()) as _, + ], + ) + .unwrap(); + let batches = RecordBatches::try_new(schema, vec![batch]).unwrap(); + + record_batches_to_labels_name(batches, &mut labels).unwrap(); + + assert_eq!( + labels, + HashSet::from(["host".to_string(), greptime_native_histogram().to_string(),]) + ); + } + + fn prometheus_metric_table_info(table_id: u32, name: &str) -> TableInfo { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "greptime_timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), false), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true), + ])); + let mut options = TableOptions::default(); + options.extra_options.insert( + LOGICAL_TABLE_METADATA_KEY.to_string(), + "greptime_physical_table".to_string(), + ); + options + .extra_options + .insert(SEMANTIC_METRIC_TYPE.to_string(), "counter".to_string()); + options + .extra_options + .insert(SEMANTIC_METRIC_UNIT.to_string(), "By".to_string()); + let meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![1]) + .engine("metric".to_string()) + .next_column_id(3) + .options(options) + .build() + .unwrap(); + + TableInfoBuilder::default() + .table_id(table_id) + .table_version(0 as TableVersion) + .name(name) + .catalog_name(DEFAULT_CATALOG_NAME) + .schema_name(DEFAULT_SCHEMA_NAME) + .table_type(TableType::Base) + .meta(meta) + .build() + .unwrap() + } + + #[tokio::test] + async fn test_retrieve_metric_metadata_uses_semantic_options() { + let table_info = prometheus_metric_table_info(1025, "http_requests_total"); + let manager: CatalogManagerRef = + MemoryCatalogManager::new_with_table(EmptyTable::from_table_info(&table_info)); + let query_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME); + + let metadata = retrieve_metric_metadata(&query_ctx, manager, &MetadataQuery::default()) + .await + .unwrap(); + + assert_eq!( + metadata.get("http_requests_total"), + Some(&vec![PromMetadata { + metric_type: "counter".to_string(), + unit: "bytes".to_string(), + help: String::new(), + }]) + ); + + for (ucum, openmetrics) in [ + ("1", "ratios"), + ("ms", "milliseconds"), + ("m/s", "meters_per_second"), + ("USD", "USD"), + ] { + let mut table_info = table_info.clone(); + table_info + .meta + .options + .extra_options + .insert(SEMANTIC_METRIC_UNIT.to_string(), ucum.to_string()); + assert_eq!( + prometheus_metadata_from_table(&table_info).unit, + openmetrics + ); + } + + for (metric_type, temporality, expected) in [ + ("updown_counter", None, "gauge"), + ("gauge_histogram", None, "gaugehistogram"), + ("summary", None, "summary"), + ("mixed", None, "unknown"), + ("counter", Some("delta"), "unknown"), + ("histogram", Some("delta"), "unknown"), + ] { + let mut table_info = table_info.clone(); + table_info + .meta + .options + .extra_options + .insert(SEMANTIC_METRIC_TYPE.to_string(), metric_type.to_string()); + if let Some(temporality) = temporality { + table_info.meta.options.extra_options.insert( + SEMANTIC_METRIC_TEMPORALITY.to_string(), + temporality.to_string(), + ); + } + assert_eq!( + prometheus_metadata_from_table(&table_info).metric_type, + expected + ); + } + } + + #[tokio::test] + async fn test_metadata_query_checks_permissions_and_filters_tables() { + let allowed = prometheus_metric_table_info(1025, "allowed"); + let denied = prometheus_metric_table_info(1026, "denied"); + let manager = MemoryCatalogManager::new_with_table(EmptyTable::from_table_info(&allowed)); + manager + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: denied.name.clone(), + table_id: denied.table_id(), + table: EmptyTable::from_table_info(&denied), + }) + .unwrap(); + let query_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME); + + let metadata = retrieve_metric_metadata( + &query_ctx, + manager.clone(), + &MetadataQuery { + metric: Some("allowed".to_string()), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(metadata.len(), 1); + assert!(metadata.contains_key("allowed")); + + let response = metadata_query( + State(Arc::new(TestPrometheusHandler { + catalog_manager: manager.clone(), + deny_operation: true, + denied_table: None, + metric_names: Vec::new(), + queries: Mutex::new(Vec::new()), + })), + Query(MetadataQuery::default()), + Extension(query_ctx.clone()), + ) + .await; + assert_eq!(Some(StatusCode::PermissionDenied), response.status_code); + + let handler: PrometheusHandlerRef = Arc::new(TestPrometheusHandler { + catalog_manager: manager, + deny_operation: false, + denied_table: Some("denied"), + metric_names: Vec::new(), + queries: Mutex::new(Vec::new()), + }); + let response = metadata_query( + State(handler.clone()), + Query(MetadataQuery::default()), + Extension(query_ctx.clone()), + ) + .await; + assert!(response.status_code.is_none()); + let PrometheusResponse::Metadata(metadata) = response.data else { + panic!("expected metadata response"); + }; + assert_eq!( + metadata.keys().cloned().collect::>(), + vec!["allowed".to_string()] + ); + + let response = metadata_query( + State(handler), + Query(MetadataQuery { + metric: Some("denied".to_string()), + ..Default::default() + }), + Extension(query_ctx), + ) + .await; + assert_eq!(Some(StatusCode::PermissionDenied), response.status_code); + } } diff --git a/src/servers/src/http/result/prometheus_resp.rs b/src/servers/src/http/result/prometheus_resp.rs index 368d925c30c..da32103b79e 100644 --- a/src/servers/src/http/result/prometheus_resp.rs +++ b/src/servers/src/http/result/prometheus_resp.rs @@ -28,6 +28,9 @@ use common_query::native_histogram::{ NativeHistogram, is_native_histogram_value_type, read_histogram, }; use common_query::prometheus::{format_prometheus_float, is_prometheus_stale_nan}; +use common_query::promql_annotations::{ + PromqlAnnotationCollector, get_promql_annotation_collector, +}; use common_query::{Output, OutputData}; use common_recordbatch::RecordBatches; use datatypes::arrow_array::string_array_value_at_index; @@ -80,6 +83,8 @@ pub struct PrometheusJsonResponse { pub error_type: Option, #[serde(skip_serializing_if = "Option::is_none")] pub warnings: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub infos: Option>, #[serde(skip)] pub status_code: Option, @@ -124,6 +129,7 @@ impl PrometheusJsonResponse { error: Some(reason.into()), error_type: Some(error_type.to_string()), warnings: None, + infos: None, resp_metrics: Default::default(), status_code: Some(error_type), } @@ -136,17 +142,37 @@ impl PrometheusJsonResponse { error: None, error_type: None, warnings: None, + infos: None, resp_metrics: Default::default(), status_code: None, } } + /// Adds collected PromQL warnings and infos to the response. + fn append_promql_annotations(&mut self, collector: &PromqlAnnotationCollector) { + let mut warnings = self.warnings.take().unwrap_or_default(); + let mut infos = self.infos.take().unwrap_or_default(); + collector.append_to(&mut warnings, &mut infos); + self.warnings = (!warnings.is_empty()).then_some(warnings); + self.infos = (!infos.is_empty()).then_some(infos); + } + + /// Merges data and annotations from another expanded PromQL query response. + pub(crate) fn append_query_response(&mut self, mut other: Self) { + self.data.append(other.data); + merge_annotations(&mut self.warnings, other.warnings.take()); + merge_annotations(&mut self.infos, other.infos.take()); + } + /// Convert from `Result` pub async fn from_query_result( result: Result, metric_name: Option, result_type: ValueType, + query_id: Option<&str>, ) -> Self { + // Hold the collector while a streaming result is consumed. + let collector = query_id.and_then(get_promql_annotation_collector); let response: Result = try { let result = result?; let mut resp = @@ -187,7 +213,7 @@ impl PrometheusJsonResponse { let result_type_string = result_type.to_string(); - match response { + let mut response = match response { Ok(resp) => resp, Err(err) => { // Prometheus won't report error if querying nonexist label and metric @@ -202,7 +228,11 @@ impl PrometheusJsonResponse { Self::error(err.status_code(), err.output_msg()) } } + }; + if let Some(collector) = collector { + response.append_promql_annotations(&collector); } + response } /// Convert [RecordBatches] to [PromData] @@ -435,6 +465,16 @@ impl PrometheusJsonResponse { } } +fn merge_annotations(target: &mut Option>, source: Option>) { + let Some(source) = source else { + return; + }; + let target = target.get_or_insert_default(); + target.extend(source); + target.sort(); + target.dedup(); +} + #[cfg(test)] mod tests { use std::sync::Arc; @@ -444,6 +484,9 @@ mod tests { native_histogram_value_type, }; use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; + use common_query::promql_annotations::{ + PromqlAnnotationCollector, promql_annotation_collector, + }; use common_recordbatch::{RecordBatch, RecordBatches}; use datatypes::data_type::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; @@ -453,6 +496,40 @@ mod tests { use super::*; + #[tokio::test] + async fn query_response_preserves_and_merges_promql_annotations() { + let query_id = "query_response_preserves_and_merges_promql_annotations"; + let left = promql_annotation_collector(query_id); + left.record_warning("shared warning"); + left.record_info("left info"); + let mut response = PrometheusJsonResponse::from_query_result( + Ok(Output::new_with_record_batches(RecordBatches::empty())), + None, + ValueType::Vector, + Some(query_id), + ) + .await; + + let right = PromqlAnnotationCollector::default(); + right.record_warning("shared warning"); + right.record_info("right info"); + let mut other = PrometheusJsonResponse::success(PrometheusResponse::None); + other.append_promql_annotations(&right); + response.append_query_response(other); + + assert_eq!(response.warnings, Some(vec!["shared warning".to_string()])); + assert_eq!( + response.infos, + Some(vec!["left info".to_string(), "right info".to_string()]) + ); + let json = serde_json::to_value(response).unwrap(); + assert_eq!(json["warnings"], serde_json::json!(["shared warning"])); + assert_eq!( + json["infos"], + serde_json::json!(["left info", "right info"]) + ); + } + fn sample_histogram() -> NativeHistogram { NativeHistogram { schema: 0, diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index b43cc73338f..6d94b3a139f 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -38,6 +38,7 @@ use crate::semantic::{ mod translator; pub use translator::legacy_normalize_otlp_name; +pub(crate) use translator::ucum_to_openmetrics_unit; use translator::{translate_label_name, translate_metric_name}; /// the default column count for table writer diff --git a/src/servers/src/otlp/metrics/translator.rs b/src/servers/src/otlp/metrics/translator.rs index bb929cd6705..8b3b45d3274 100644 --- a/src/servers/src/otlp/metrics/translator.rs +++ b/src/servers/src/otlp/metrics/translator.rs @@ -219,6 +219,20 @@ fn build_clean_unit_suffix(unit: &str) -> (Option, Option) { ) } +/// Converts a UCUM unit to its Prometheus/OpenMetrics unit word. +pub(crate) fn ucum_to_openmetrics_unit(unit: &str) -> String { + if unit == "1" { + return "ratios".to_string(); + } + + match build_clean_unit_suffix(unit) { + (Some(main), Some(per)) => format!("{main}_per_{per}"), + (Some(main), None) => main, + (None, Some(per)) => format!("per_{per}"), + (None, None) => String::new(), + } +} + fn build_unit_suffixes(unit: &str) -> (String, String) { let (main, per) = unit.split_once('/').unwrap_or((unit, "")); let main_unit_suffix = unit_suffix(main, &UNIT_MAP); diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index 792a2ce965d..23862870dba 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -980,6 +980,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) { error: None, error_type: None, warnings: None, + infos: None, resp_metrics: Default::default(), status_code: None, }; diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index cd8af8768a4..042be385a6d 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2741,6 +2741,28 @@ pub async fn test_prometheus_remote_write_v2(store_type: StorageType) { ) .await; + let res = client + .get("/v1/prometheus/api/v1/metadata?metric=remote_write_v2_typed_total") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!( + serde_json::to_string_pretty(&body).unwrap(), + r#"{ + "status": "success", + "data": { + "remote_write_v2_typed_total": [ + { + "type": "counter", + "unit": "seconds", + "help": "" + } + ] + } +}"# + ); + guard.remove_all().await; } @@ -2882,6 +2904,202 @@ pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: Storag ) .await; + let res = client + .get("/v1/prometheus/api/v1/metadata?metric=remote_write_v2_latency_seconds") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!( + serde_json::to_string_pretty(&body).unwrap(), + r#"{ + "status": "success", + "data": { + "remote_write_v2_latency_seconds": [ + { + "type": "histogram", + "unit": "", + "help": "" + } + ] + } +}"# + ); + + let res = client + .get("/v1/prometheus/api/v1/query?query=histogram_count(remote_write_v2_latency_seconds)&time=4") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!( + serde_json::to_string_pretty(&body).unwrap(), + r#"{ + "status": "success", + "data": { + "resultType": "vector", + "result": [ + { + "metric": { + "__name__": "remote_write_v2_latency_seconds", + "instance": "localhost:9090", + "job": "api" + }, + "value": [ + 4.0, + "6" + ] + } + ] + } +}"# + ); + + let res = client + .get("/v1/prometheus/api/v1/query?query=remote_write_v2_latency_seconds&time=4") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!( + serde_json::to_string_pretty(&body).unwrap(), + r#"{ + "status": "success", + "data": { + "resultType": "vector", + "result": [ + { + "metric": { + "__name__": "remote_write_v2_latency_seconds", + "instance": "localhost:9090", + "job": "api" + }, + "histogram": [ + 4.0, + { + "count": "6", + "sum": "20", + "buckets": [ + [ + 3, + "-0.002", + "0.002", + "0.5" + ], + [ + 0, + "1.4142135623730951", + "1.681792830507429", + "2" + ], + [ + 0, + "1.681792830507429", + "2", + "3.5" + ] + ] + } + ] + } + ] + } +}"# + ); + + let res = client + .get("/v1/prometheus/api/v1/query_range?query=remote_write_v2_latency_seconds&start=3&end=4&step=1") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + let body = res.json::().await; + assert_eq!( + serde_json::to_string_pretty(&body).unwrap(), + r#"{ + "status": "success", + "data": { + "resultType": "matrix", + "result": [ + { + "metric": { + "__name__": "remote_write_v2_latency_seconds", + "instance": "localhost:9090", + "job": "api" + }, + "histograms": [ + [ + 3.0, + { + "count": "8", + "sum": "10", + "buckets": [ + [ + 1, + "-0.5", + "-0.3535533905932738", + "1" + ], + [ + 3, + "-0.001", + "0.001", + "1" + ], + [ + 0, + "0.7071067811865476", + "1", + "1" + ], + [ + 0, + "1", + "1.4142135623730951", + "3" + ], + [ + 0, + "1.4142135623730951", + "2", + "2" + ] + ] + } + ], + [ + 4.0, + { + "count": "6", + "sum": "20", + "buckets": [ + [ + 3, + "-0.002", + "0.002", + "0.5" + ], + [ + 0, + "1.4142135623730951", + "1.681792830507429", + "2" + ], + [ + 0, + "1.681792830507429", + "2", + "3.5" + ] + ] + } + ] + ] + } + ] + } +}"# + ); + validate_data( "prometheus_remote_write_v2_native_histogram_table_created", &client, From 9241e96fa2ea8ed3132e38dea93c655337d24040 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 12 Aug 2026 07:49:46 +0000 Subject: [PATCH 32/92] fix(metric-engine): handle Utf8View tag/label columns without panicking (#8772) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(metric-engine): handle Utf8View tag/label columns without panicking label_replace (planned as DataFusion regexp_replace) coerces to Utf8View, so label columns materialize as StringViewArray; build_tag_arrays' StringArray downcast then panicked ('tag column must be utf8') — e.g. for OTLP/json2 ingest. TSID computation, sparse-PK encoding and tag extraction now accept generic ArrayRef tag columns (Utf8/LargeUtf8/ Utf8View/Dictionary) via string_array_value_at_index, and build_tag_arrays errors instead of panicking on non-string columns. The mito2 time-series memtable string-field paths are hardened the same way. Adds label_replace_with_utf8view_labels_does_not_panic (issue #8732). Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * refactor(metric-engine): add is_string_null_at helper for tag null checks Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(datatypes): use is_none_or to satisfy clippy unnecessary-map-or Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix: reject oversized string batches before memtable append Distinguish a full active string builder from a batch that cannot fit an empty Arrow string builder at all. Scan every string field so a later intrinsically oversized field cannot be skipped after an earlier field requests a freeze. Return InvalidBatch instead of reaching Arrow's offset overflow panic. Also cover Utf8View tags with nulls through the metric-engine tag, TSID, and sparse-primary-key path. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/datatypes/src/arrow_array.rs | 37 ++++ src/metric-engine/src/batch_modifier.rs | 207 ++++++++++++++++-- src/mito2/src/memtable/time_series.rs | 202 ++++++++++++++--- .../src/http/result/prometheus_resp.rs | 52 +++++ 4 files changed, 447 insertions(+), 51 deletions(-) diff --git a/src/datatypes/src/arrow_array.rs b/src/datatypes/src/arrow_array.rs index e5ef7add827..848cc7c91f7 100644 --- a/src/datatypes/src/arrow_array.rs +++ b/src/datatypes/src/arrow_array.rs @@ -168,6 +168,43 @@ pub fn string_array_value_at_index(array: &ArrayRef, i: usize) -> Option<&str> { } } +/// Check whether the string value at index `i` is null for string or +/// dictionary-encoded string arrays. +/// +/// Returns `true` when the value is null or the array type is not a string +/// type, which corresponds to [`string_array_value_at_index`] returning `None`. +/// +/// # Panics +/// +/// If index `i` is out of bounds. +pub fn is_string_null_at(array: &ArrayRef, i: usize) -> bool { + match array.data_type() { + DataType::Utf8 => { + let array = array.as_string::(); + !array.is_valid(i) + } + DataType::LargeUtf8 => { + let array = array.as_string::(); + !array.is_valid(i) + } + DataType::Utf8View => { + let array = array.as_string_view(); + !array.is_valid(i) + } + DataType::Dictionary(key_type, value_type) + if key_type.is_integer() && value_type.is_string() => + { + downcast_dictionary_array! { + array => array + .key(i) + .is_none_or(|key| is_string_null_at(array.values(), key)), + _ => true, + } + } + _ => true, + } +} + /// Get the string value at index `i` for `Utf8`, `LargeUtf8`, or `Utf8View` arrays. /// /// Note: This method does not check for nulls and the value is arbitrary diff --git a/src/metric-engine/src/batch_modifier.rs b/src/metric-engine/src/batch_modifier.rs index 2162215d1b1..575dffff1a3 100644 --- a/src/metric-engine/src/batch_modifier.rs +++ b/src/metric-engine/src/batch_modifier.rs @@ -15,9 +15,10 @@ use std::hash::Hasher; use std::sync::Arc; -use datatypes::arrow::array::{Array, BinaryBuilder, StringArray, UInt64Array}; +use datatypes::arrow::array::{Array, ArrayRef, BinaryBuilder, UInt64Array}; use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema}; use datatypes::arrow::record_batch::RecordBatch; +use datatypes::arrow_array::{is_string_null_at, string_array_value_at_index}; use fxhash::FxHasher; use mito_codec::row_converter::SparsePrimaryKeyCodec; use snafu::ResultExt; @@ -49,7 +50,7 @@ pub struct TagColumnInfo { pub fn compute_tsid_array( batch: &RecordBatch, sorted_tag_columns: &[TagColumnInfo], - tag_arrays: &[&StringArray], + tag_arrays: &[&ArrayRef], ) -> UInt64Array { let num_rows = batch.num_rows(); @@ -64,20 +65,22 @@ pub fn compute_tsid_array( let mut tsid_values = Vec::with_capacity(num_rows); for row in 0..num_rows { - let has_null = tag_arrays.iter().any(|arr| arr.is_null(row)); + let has_null = tag_arrays.iter().any(|arr| is_string_null_at(arr, row)); let tsid = if !has_null { let mut hasher = FxHasher::default(); hasher.write_u64(label_name_hash); for arr in tag_arrays { - hasher.write(arr.value(row).as_bytes()); - hasher.write_u8(0xff); + if let Some(value) = string_array_value_at_index(arr, row) { + hasher.write(value.as_bytes()); + hasher.write_u8(0xff); + } } hasher.finish() } else { let mut name_hasher = FxHasher::default(); for (tc, arr) in sorted_tag_columns.iter().zip(tag_arrays.iter()) { - if !arr.is_null(row) { + if !is_string_null_at(arr, row) { name_hasher.write(tc.name.as_bytes()); name_hasher.write_u8(0xff); } @@ -87,8 +90,8 @@ pub fn compute_tsid_array( let mut val_hasher = FxHasher::default(); val_hasher.write_u64(row_label_hash); for arr in tag_arrays { - if !arr.is_null(row) { - val_hasher.write(arr.value(row).as_bytes()); + if let Some(value) = string_array_value_at_index(arr, row) { + val_hasher.write(value.as_bytes()); val_hasher.write_u8(0xff); } } @@ -104,15 +107,26 @@ pub fn compute_tsid_array( fn build_tag_arrays<'a>( batch: &'a RecordBatch, sorted_tag_columns: &[TagColumnInfo], -) -> Vec<&'a StringArray> { +) -> Result> { sorted_tag_columns .iter() .map(|tc| { - batch - .column(tc.index) - .as_any() - .downcast_ref::() - .expect("tag column must be utf8") + let array = batch.column(tc.index); + match array.data_type() { + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Ok(array), + DataType::Dictionary(key_type, value_type) + if key_type.is_integer() && value_type.is_string() => + { + Ok(array) + } + data_type => UnexpectedRequestSnafu { + reason: format!( + "Tag column '{}' must be a string, given: {data_type}", + tc.name + ), + } + .fail(), + } }) .collect() } @@ -144,7 +158,7 @@ pub fn modify_batch_sparse( ) -> Result { let num_rows = batch.num_rows(); let codec = SparsePrimaryKeyCodec::schemaless(); - let tag_arrays: Vec<&StringArray> = build_tag_arrays(&batch, sorted_tag_columns); + let tag_arrays = build_tag_arrays(&batch, sorted_tag_columns)?; let tsid_array = compute_tsid_array(&batch, sorted_tag_columns, &tag_arrays); let mut pk_builder = BinaryBuilder::with_capacity(num_rows, 0); @@ -158,8 +172,9 @@ pub fn modify_batch_sparse( let tags = sorted_tag_columns .iter() .zip(tag_arrays.iter()) - .filter(|(_, arr)| !arr.is_null(row)) - .map(|(tc, arr)| (tc.column_id, arr.value(row).as_bytes())); + .filter_map(|(tc, arr)| { + string_array_value_at_index(arr, row).map(|value| (tc.column_id, value.as_bytes())) + }); codec .encode_raw_tag_value(tags, &mut buffer) .context(EncodePrimaryKeySnafu)?; @@ -268,7 +283,7 @@ mod tests { column_id: 1, }, ]; - let tag_arrays = build_tag_arrays(&batch, &tag_columns); + let tag_arrays = build_tag_arrays(&batch, &tag_columns).unwrap(); let tsid_array = compute_tsid_array(&batch, &tag_columns, &tag_arrays); assert_eq!(tsid_array.value(0), 2721566936019240841); @@ -300,7 +315,7 @@ mod tests { column_id: 2, }, ]; - let tag_arrays_2 = build_tag_arrays(&batch_no_null, &tag_cols_2); + let tag_arrays_2 = build_tag_arrays(&batch_no_null, &tag_cols_2).unwrap(); let tsid_no_null = compute_tsid_array(&batch_no_null, &tag_cols_2, &tag_arrays_2); let schema3 = Arc::new(ArrowSchema::new(vec![ @@ -334,7 +349,7 @@ mod tests { column_id: 3, }, ]; - let tag_arrays_3 = build_tag_arrays(&batch_with_null, &tag_cols_3); + let tag_arrays_3 = build_tag_arrays(&batch_with_null, &tag_cols_3).unwrap(); let tsid_with_null = compute_tsid_array(&batch_with_null, &tag_cols_3, &tag_arrays_3); assert_eq!(tsid_no_null.value(0), tsid_with_null.value(0)); @@ -452,4 +467,156 @@ mod tests { .unwrap(); assert_eq!(actual_array.value(0), expected_pk.as_slice()); } + + #[test] + fn label_replace_with_utf8view_labels_does_not_panic() { + // Tag (label) columns may arrive as any string representation (`Utf8`, + // `LargeUtf8`, `Utf8View`, or dictionary-encoded). `modify_batch_sparse` + // must not assume the tag columns are plain `StringArray`s. + let tag_arrays: Vec = vec![ + Arc::new(StringArray::from(vec!["greptimedb"])), + Arc::new(datatypes::arrow::array::LargeStringArray::from(vec![ + "greptimedb", + ])), + Arc::new(datatypes::arrow::array::StringViewArray::from(vec![ + "greptimedb", + ])), + Arc::new(datatypes::arrow::array::DictionaryArray::< + datatypes::arrow::datatypes::UInt32Type, + >::new( + datatypes::arrow::array::UInt32Array::from(vec![0]), + Arc::new(StringArray::from(vec!["greptimedb"])), + )), + ]; + let tag_columns = vec![TagColumnInfo { + name: "namespace".to_string(), + index: 2, + column_id: 2, + }]; + let non_tag_indices = vec![0, 1]; + + let primary_keys = tag_arrays + .into_iter() + .map(|tag_array| { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("greptime_timestamp", DataType::Int64, false), + Field::new("greptime_value", DataType::Float64, true), + Field::new("namespace", tag_array.data_type().clone(), true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![1000])), + Arc::new(datatypes::arrow::array::Float64Array::from(vec![42.0])), + tag_array, + ], + ) + .unwrap(); + let modified = + modify_batch_sparse(batch, 1025, &tag_columns, &non_tag_indices).unwrap(); + modified + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + .to_vec() + }) + .collect::>(); + + assert!(primary_keys.windows(2).all(|keys| keys[0] == keys[1])); + } + + #[test] + fn test_compute_tsid_utf8view_with_nulls_matches_utf8() { + // Regression test for Utf8View tag/label columns: the whole + // `build_tag_arrays` -> `compute_tsid_array` -> sparse primary key path must handle + // `Utf8View` columns exactly like `Utf8`, including rows with null tags, instead of + // panicking or producing different TSIDs. + let namespace_values: Vec> = + vec![Some("ns-a"), Some("ns-b"), None, Some("ns-c")]; + let host_values: Vec> = vec![Some("host-1"), None, Some("host-2"), None]; + + let build_batch = |data_type: DataType| { + let (namespace, host): (ArrayRef, ArrayRef) = match &data_type { + DataType::Utf8 => ( + Arc::new(StringArray::from(namespace_values.clone())), + Arc::new(StringArray::from(host_values.clone())), + ), + DataType::Utf8View => ( + Arc::new(datatypes::arrow::array::StringViewArray::from( + namespace_values.clone(), + )), + Arc::new(datatypes::arrow::array::StringViewArray::from( + host_values.clone(), + )), + ), + _ => unreachable!(), + }; + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("greptime_timestamp", DataType::Int64, false), + Field::new("greptime_value", DataType::Float64, true), + Field::new("namespace", data_type.clone(), true), + Field::new("host", data_type, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![1000, 1001, 1002, 1003])), + Arc::new(datatypes::arrow::array::Float64Array::from(vec![ + 1.0, 2.0, 3.0, 4.0, + ])), + namespace, + host, + ], + ) + .unwrap() + }; + + let tag_columns = vec![ + TagColumnInfo { + name: "host".to_string(), + index: 3, + column_id: 3, + }, + TagColumnInfo { + name: "namespace".to_string(), + index: 2, + column_id: 2, + }, + ]; + + let utf8_batch = build_batch(DataType::Utf8); + let utf8view_batch = build_batch(DataType::Utf8View); + + // Both representations must be accepted by `build_tag_arrays`. + let utf8_tag_arrays = build_tag_arrays(&utf8_batch, &tag_columns).unwrap(); + let utf8view_tag_arrays = build_tag_arrays(&utf8view_batch, &tag_columns).unwrap(); + + // TSIDs must be identical across representations, including rows with null tags. + let utf8_tsids = compute_tsid_array(&utf8_batch, &tag_columns, &utf8_tag_arrays); + let utf8view_tsids = + compute_tsid_array(&utf8view_batch, &tag_columns, &utf8view_tag_arrays); + assert_eq!(utf8_tsids, utf8view_tsids); + + // The full sparse primary key (TSID + tag values) must also be identical. + let non_tag_indices = vec![0, 1]; + let modified_utf8 = + modify_batch_sparse(utf8_batch, 1025, &tag_columns, &non_tag_indices).unwrap(); + let modified_utf8view = + modify_batch_sparse(utf8view_batch, 1025, &tag_columns, &non_tag_indices).unwrap(); + let utf8_pks = modified_utf8 + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let utf8view_pks = modified_utf8view + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..utf8_pks.len() { + assert_eq!(utf8_pks.value(row), utf8view_pks.value(row)); + } + } } diff --git a/src/mito2/src/memtable/time_series.rs b/src/mito2/src/memtable/time_series.rs index f5b772b076f..7e89d569eeb 100644 --- a/src/mito2/src/memtable/time_series.rs +++ b/src/mito2/src/memtable/time_series.rs @@ -32,8 +32,8 @@ use datatypes::prelude::{ScalarVector, Vector, VectorRef}; use datatypes::types::TimestampType; use datatypes::value::{Value, ValueRef}; use datatypes::vectors::{ - Helper, TimestampMicrosecondVector, TimestampMillisecondVector, TimestampNanosecondVector, - TimestampSecondVector, UInt8Vector, UInt64Vector, + Helper, StringVector, TimestampMicrosecondVector, TimestampMillisecondVector, + TimestampNanosecondVector, TimestampSecondVector, UInt8Vector, UInt64Vector, }; use mito_codec::key_values::KeyValue; use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt}; @@ -72,6 +72,93 @@ const INITIAL_BUILDER_CAPACITY: usize = 4; /// Vector builder capacity. const BUILDER_CAPACITY: usize = 512; +fn checked_string_values_size(mut lengths: impl Iterator) -> Option { + lengths.try_fold(0_i32, |size, len| { + size.checked_add(i32::try_from(len).ok()?) + }) +} + +/// Checks whether a string field builder that currently holds `current_len` bytes can +/// additionally accommodate `space_needed` bytes of string data, given that the offset of a +/// single Arrow string array is limited to `limit` (`i32::MAX` in production). +/// +/// Returns: +/// - `Ok(true)` if `current_len + space_needed <= limit`; +/// - `Ok(false)` if `space_needed <= limit` but `current_len + space_needed > limit`, i.e. +/// the current builder is too full but an empty builder could accommodate the batch (the +/// caller may freeze the current builder and replay the batch onto an empty one); +/// - `Err(InvalidBatchSnafu)` if the batch's string data alone (`space_needed`, or the fact +/// that its total bytes cannot even be represented as `i32`) exceeds `limit`, i.e. the +/// batch can never be accommodated by any builder. +fn check_string_capacity(current_len: i32, space_needed: Option, limit: i32) -> Result { + let Some(space_needed) = space_needed.filter(|&space| space <= limit) else { + return error::InvalidBatchSnafu { + reason: format!( + "String data of the batch exceeds the Arrow string array offset limit ({limit}) and cannot be stored in a single column" + ), + } + .fail(); + }; + Ok(current_len + .checked_add(space_needed) + .is_some_and(|v| v <= limit)) +} + +/// Scans all string fields of a batch and checks whether they can be accommodated. +/// +/// Every string field is checked (no early return on `Ok(false)`) so that an intrinsic +/// oversize in *any* field surfaces as an error even when an earlier field only overflows +/// the current builder. +/// +/// Returns: +/// - `Ok(true)` if every string field fits into the current builders; +/// - `Ok(false)` if no field is intrinsically oversized but at least one field only fits +/// into an empty builder (the caller may freeze the current builder and replay the batch +/// onto an empty one); +/// - `Err(InvalidBatchSnafu)` if any single field's string data alone (`space_needed`, or +/// the fact that its total bytes cannot even be represented as `i32`) exceeds `limit`, +/// i.e. that field can never be accommodated by any builder. +fn scan_string_capacity( + fields: &[VectorRef], + field_builders: &[Option], + field_types: &[ConcreteDataType], + limit: i32, +) -> Result { + let mut can_fit_current = true; + for ((field_src, field_dest), field_type) in fields + .iter() + .zip(field_builders.iter()) + .zip(field_types.iter()) + { + if !matches!(field_type, ConcreteDataType::String(_)) { + continue; + } + let current_size = match field_dest { + Some(FieldBuilder::String(builder)) => builder.next_offset(), + None => 0, + Some(FieldBuilder::Other(_)) => unreachable!(), + }; + let array = field_src.to_arrow_array(); + let space_needed = if let Some(string_array) = array.as_any().downcast_ref::() + { + i32::try_from(string_array.value_data().len()).ok() + } else { + let string_vector = field_src + .as_any() + .downcast_ref::() + .with_context(|| error::InvalidBatchSnafu { + reason: format!( + "Field type mismatch, expecting String, given: {}", + field_src.data_type() + ), + })?; + checked_string_values_size(string_vector.iter_data().flatten().map(str::len)) + }; + can_fit_current &= check_string_capacity(current_size, space_needed, limit)?; + } + Ok(can_fit_current) +} + /// Builder to build [TimeSeriesMemtable]. #[derive(Debug, Default)] pub struct TimeSeriesMemtableBuilder { @@ -920,32 +1007,17 @@ impl ValueBuilder { } /// Checks if current value builder have sufficient space to accommodate `fields`. - /// Returns false if there is no space to accommodate fields due to offset overflow. + /// + /// Returns `Ok(false)` if the current builder lacks the remaining space to accommodate + /// the fields due to offset overflow, but an empty builder would be able to accommodate + /// the batch (the caller may freeze the current builder and replay the batch onto an + /// empty one). + /// + /// Returns `Err(InvalidBatchSnafu)` if the string data of a single batch itself exceeds + /// the Arrow string array offset limit and thus can never be accommodated, not even by an + /// empty builder. pub(crate) fn can_accommodate(&self, fields: &[VectorRef]) -> Result { - for (field_src, field_dest) in fields.iter().zip(self.fields.iter()) { - let Some(builder) = field_dest else { - continue; - }; - let FieldBuilder::String(builder) = builder else { - continue; - }; - let array = field_src.to_arrow_array(); - let string_array = array - .as_any() - .downcast_ref::() - .with_context(|| error::InvalidBatchSnafu { - reason: format!( - "Field type mismatch, expecting String, given: {}", - field_src.data_type() - ), - })?; - let space_needed = string_array.value_data().len() as i32; - // offset may overflow - if builder.next_offset().checked_add(space_needed).is_none() { - return Ok(false); - } - } - Ok(true) + scan_string_capacity(fields, &self.fields, &self.field_types, i32::MAX) } pub(crate) fn extend( @@ -1017,17 +1089,26 @@ impl ValueBuilder { match builder { FieldBuilder::String(builder) => { let array = field_src.to_arrow_array(); - let string_array = - array + if let Some(string_array) = array.as_any().downcast_ref::() { + builder.append_array(string_array); + } else { + let string_vector = field_src .as_any() - .downcast_ref::() + .downcast_ref::() .with_context(|| error::InvalidBatchSnafu { reason: format!( "Field type mismatch, expecting String, given: {}", field_src.data_type() ), })?; - builder.append_array(string_array); + for value in string_vector.iter_data() { + if let Some(value) = value { + builder.append(value); + } else { + builder.append_null(); + } + } + } } FieldBuilder::Other(builder) => { let len = field_src.len(); @@ -2296,4 +2377,63 @@ mod tests { assert!(iter.next().is_none()); } + + #[test] + fn test_can_accommodate_string_offset_overflow() { + // A batch whose string data alone exceeds the Arrow string array offset limit can + // never be accommodated, not even by an empty builder. `can_accommodate` must return + // a structured `InvalidBatch` error instead of `Ok(false)`: the latter would make + // `Series::extend` freeze the current builder and replay the same oversized batch + // onto an empty builder, which then panics on offset overflow in `StringBuilder`. + assert!(check_string_capacity(0, Some(101), 100).is_err()); + assert!(check_string_capacity(0, None, i32::MAX).is_err()); + + // A batch that fits in an empty builder but not in the current one reports + // `Ok(false)` so the caller can freeze and replay onto an empty builder. + assert!(!check_string_capacity(95, Some(10), 100).unwrap()); + assert!(!check_string_capacity(91, Some(10), 100).unwrap()); + + // The current builder can accommodate the batch when the combined size fits, + // including the case where it exactly reaches the limit. + assert!(check_string_capacity(0, Some(100), 100).unwrap()); + assert!(check_string_capacity(90, Some(10), 100).unwrap()); + } + + #[test] + fn test_can_accommodate_checks_all_string_fields() { + // Two string fields: the first only overflows the *current* builder (an empty + // builder would fit it), the second is intrinsically oversized. The scan must not + // early-return `Ok(false)` on the first field: it has to check every string field so + // the intrinsic oversize of the second field still surfaces as `InvalidBatch`. + let limit = 10; + // 8 bytes already in the current builder: 8 + 4 > 10, but 4 <= 10, so this field + // alone fits an empty builder. + let mut current = StringBuilder::with_capacity(1, 8); + current.append("12345678"); + let field_builders = [Some(FieldBuilder::String(current)), None]; + let field_types = [ + ConcreteDataType::string_datatype(), + ConcreteDataType::string_datatype(), + ]; + + let first = Arc::new(StringVector::from(StringArray::from(vec!["abcd"]))) as VectorRef; + // 15 bytes > limit: intrinsically oversized, can never fit any builder. + let second = Arc::new(StringVector::from(StringArray::from(vec![ + "123456789012345", + ]))) as VectorRef; + + let err = scan_string_capacity(&[first, second], &field_builders, &field_types, limit) + .unwrap_err(); + assert!( + matches!(err, error::Error::InvalidBatch { .. }), + "expected InvalidBatch, got {err:?}" + ); + + // Sanity: without the oversized field the result is just `Ok(false)` (current full, + // empty builder fits), which lets the caller freeze and replay onto an empty builder. + let first_only = Arc::new(StringVector::from(StringArray::from(vec!["abcd"]))) as VectorRef; + assert!( + !scan_string_capacity(&[first_only], &field_builders, &field_types, limit).unwrap() + ); + } } diff --git a/src/servers/src/http/result/prometheus_resp.rs b/src/servers/src/http/result/prometheus_resp.rs index da32103b79e..e79628fac0d 100644 --- a/src/servers/src/http/result/prometheus_resp.rs +++ b/src/servers/src/http/result/prometheus_resp.rs @@ -479,6 +479,7 @@ fn merge_annotations(target: &mut Option>, source: Option Date: Wed, 12 Aug 2026 07:50:54 +0000 Subject: [PATCH 33/92] feat(promql): support native histogram aggregations (#8848) * feat(promql): support native histogram aggregations Signed-off-by: shuiyisong * fix(promql): correct mixed native histogram aggregations Signed-off-by: shuiyisong * fix(promql): format mixed count_values labels consistently Signed-off-by: shuiyisong * fix(promql): preserve reset hint warnings for incompatible histograms Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/common/query/src/native_histogram.rs | 59 +- src/promql/src/functions.rs | 2 +- src/promql/src/functions/native_histogram.rs | 448 +++++++- src/query/src/promql/planner.rs | 1005 +++++++++++++++-- .../src/query_engine/default_serializer.rs | 39 +- 5 files changed, 1406 insertions(+), 147 deletions(-) diff --git a/src/common/query/src/native_histogram.rs b/src/common/query/src/native_histogram.rs index 2db873487bd..b090ac33fcd 100644 --- a/src/common/query/src/native_histogram.rs +++ b/src/common/query/src/native_histogram.rs @@ -247,6 +247,25 @@ pub struct NativeHistogram { pub negative_buckets: Vec, } +/// Formats a histogram component like Prometheus's default `%g` formatter. +fn format_promql_histogram_float(value: f64) -> String { + let abs = value.abs(); + if !value.is_finite() || value == 0.0 || (1e-4..1e6).contains(&abs) { + return format_prometheus_float(value); + } + + let scientific = format!("{value:e}"); + let Some((mantissa, exponent)) = scientific.rsplit_once('e') else { + return scientific; + }; + let (sign, exponent) = if let Some(exponent) = exponent.strip_prefix('-') { + ('-', exponent) + } else { + ('+', exponent.strip_prefix('+').unwrap_or(exponent)) + }; + format!("{mantissa}e{sign}{exponent:0>2}") +} + impl NativeHistogram { fn uses_custom_buckets(&self) -> bool { self.schema == CUSTOM_BUCKETS_SCHEMA @@ -378,7 +397,10 @@ impl NativeHistogram { /// Formats this histogram using PromQL native histogram sample notation. pub fn promql_string(&self) -> String { - let mut parts = vec![format!("count:{}", self.count), format!("sum:{}", self.sum)]; + let mut parts = vec![ + format!("count:{}", format_promql_histogram_float(self.count)), + format!("sum:{}", format_promql_histogram_float(self.sum)), + ]; if let Some(buckets) = self.all_buckets() { parts.extend( buckets @@ -392,7 +414,11 @@ impl NativeHistogram { }; format!( "{}{},{}{}:{}", - left, bucket.lower, bucket.upper, right, bucket.count + left, + format_promql_histogram_float(bucket.lower), + format_promql_histogram_float(bucket.upper), + right, + format_promql_histogram_float(bucket.count) ) }), ); @@ -1624,6 +1650,35 @@ mod tests { } } + #[test] + fn promql_string_matches_prometheus_float_format() { + assert_eq!(format_promql_histogram_float(0.0001), "0.0001"); + assert_eq!(format_promql_histogram_float(1_000_000.0), "1e+06"); + + let histogram = NativeHistogram { + schema: CUSTOM_BUCKETS_SCHEMA, + zero_threshold: 0.0, + sum: 2_349_209.324, + reset_hint: CounterResetHint::Unknown, + start_timestamp: None, + custom_values: vec![0.00001], + positive_spans: vec![Span { + offset: 0, + length: 2, + }], + negative_spans: Vec::new(), + count: 3.0, + zero_count: 0.0, + positive_buckets: vec![1.0, 2.0], + negative_buckets: Vec::new(), + }; + + assert_eq!( + histogram.promql_string(), + "{count:3, sum:2.349209324e+06, [-Inf,1e-05]:1, (1e-05,+Inf]:2}" + ); + } + #[test] fn span_rebuild_checks_offsets_and_preserves_empty_input() { assert_eq!( diff --git a/src/promql/src/functions.rs b/src/promql/src/functions.rs index b34c21c5569..22cd64b9147 100644 --- a/src/promql/src/functions.rs +++ b/src/promql/src/functions.rs @@ -54,7 +54,7 @@ pub use native_histogram::{ NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile, NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, - NativeHistogramToString, + NativeHistogramToString, PromqlFloatToString, }; pub use predict_linear::PredictLinear; pub use quantile::QuantileOverTime; diff --git a/src/promql/src/functions/native_histogram.rs b/src/promql/src/functions/native_histogram.rs index 5b633defa96..18363c8f771 100644 --- a/src/promql/src/functions/native_histogram.rs +++ b/src/promql/src/functions/native_histogram.rs @@ -20,6 +20,7 @@ use std::mem::size_of; use std::sync::Arc; use common_query::native_histogram::*; +use common_query::prometheus::format_prometheus_float; use common_query::promql_annotations::PromqlAnnotationCollector; use datafusion::arrow::array::{ Array, ArrayRef, BooleanArray, Float64Array, Float64Builder, Int64Array, StringBuilder, @@ -92,20 +93,28 @@ fn read_scalar_f64_arg( #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] enum AnnotationReturn { FloatNull, + BooleanTrue, BooleanFalse, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +enum AnnotationLevel { + Info, + Warning, +} + impl AnnotationReturn { fn data_type(self) -> DataType { match self { Self::FloatNull => DataType::Float64, - Self::BooleanFalse => DataType::Boolean, + Self::BooleanTrue | Self::BooleanFalse => DataType::Boolean, } } fn scalar_value(self) -> ScalarValue { match self { Self::FloatNull => ScalarValue::Float64(None), + Self::BooleanTrue => ScalarValue::Boolean(Some(true)), Self::BooleanFalse => ScalarValue::Boolean(Some(false)), } } @@ -116,6 +125,7 @@ struct NativeHistogramAnnotationUdf { name: &'static str, signature: Signature, return_kind: AnnotationReturn, + level: AnnotationLevel, message: String, collector: Option, } @@ -124,6 +134,7 @@ impl NativeHistogramAnnotationUdf { fn new( name: &'static str, return_kind: AnnotationReturn, + level: AnnotationLevel, message: String, collector: Option, ) -> Self { @@ -131,6 +142,7 @@ impl NativeHistogramAnnotationUdf { name, signature: Signature::variadic_any(Volatility::Volatile), return_kind, + level, message, collector, } @@ -141,6 +153,7 @@ impl PartialEq for NativeHistogramAnnotationUdf { fn eq(&self, other: &Self) -> bool { self.name == other.name && self.return_kind == other.return_kind + && self.level == other.level && self.message == other.message } } @@ -151,6 +164,7 @@ impl Hash for NativeHistogramAnnotationUdf { fn hash(&self, state: &mut H) { self.name.hash(state); self.return_kind.hash(state); + self.level.hash(state); self.message.hash(state); } } @@ -188,7 +202,10 @@ impl ScalarUDFImpl for NativeHistogramAnnotationUdf { .cloned() .or_else(|| self.collector.clone()) { - collector.record_info(self.message.clone()); + match self.level { + AnnotationLevel::Info => collector.record_info(self.message.clone()), + AnnotationLevel::Warning => collector.record_warning(self.message.clone()), + } } Ok(ColumnarValue::Scalar(self.return_kind.scalar_value())) } @@ -205,6 +222,10 @@ impl NativeHistogramDrop { "prom_native_histogram_drop_bool" } + const fn bool_true_name() -> &'static str { + "prom_native_histogram_keep_bool" + } + pub fn float_null_udf( message: String, collector: Option, @@ -212,6 +233,7 @@ impl NativeHistogramDrop { ScalarUDF::new_from_impl(NativeHistogramAnnotationUdf::new( Self::float_null_name(), AnnotationReturn::FloatNull, + AnnotationLevel::Info, message, collector, )) @@ -224,14 +246,39 @@ impl NativeHistogramDrop { ScalarUDF::new_from_impl(NativeHistogramAnnotationUdf::new( Self::bool_false_name(), AnnotationReturn::BooleanFalse, + AnnotationLevel::Info, + message, + collector, + )) + } + + pub fn bool_true_udf( + message: String, + collector: Option, + ) -> ScalarUDF { + ScalarUDF::new_from_impl(NativeHistogramAnnotationUdf::new( + Self::bool_true_name(), + AnnotationReturn::BooleanTrue, + AnnotationLevel::Info, + message, + collector, + )) + } + + pub fn warning_bool_false_udf( + message: String, + collector: Option, + ) -> ScalarUDF { + ScalarUDF::new_from_impl(NativeHistogramAnnotationUdf::new( + Self::bool_false_name(), + AnnotationReturn::BooleanFalse, + AnnotationLevel::Warning, message, collector, )) } } -// Synthetic drop UDFs only cover info-level invalid PromQL combinations. -// Warning annotations stay in the histogram UDFs that actually drop samples. fn record_info(collector: &Option, message: impl Into) { if let Some(collector) = collector { collector.record_info(message); @@ -358,10 +405,8 @@ fn histogram_pair_udf_with_collector( match (read_histogram(lhs, row)?, read_histogram(rhs, row)?) { (Some(lhs), Some(rhs)) => { record_custom_reconciliation(&collector, name, &lhs, &rhs); + record_counter_reset_contradiction(&collector, name, &lhs, &rhs); let result = op(&lhs, &rhs); - if result.is_some() { - record_counter_reset_contradiction(&collector, name, &lhs, &rhs); - } if result.is_none() { record_warning( &collector, @@ -697,6 +742,39 @@ impl NativeHistogramStdvar { } } +/// Formats float samples as PromQL label values. +pub struct PromqlFloatToString; + +impl PromqlFloatToString { + pub const fn name() -> &'static str { + "prom_float_to_string" + } + + pub fn scalar_udf() -> ScalarUDF { + create_udf( + Self::name(), + vec![DataType::Float64], + DataType::Utf8, + Volatility::Volatile, + Arc::new(|input: &[ColumnarValue]| { + let values = extract_array(&input[0])?; + let values = values + .as_any() + .downcast_ref::() + .expect("validated Float64 input"); + let mut result = StringBuilder::new(); + for value in values.iter() { + match value { + Some(value) => result.append_value(format_prometheus_float(value)), + None => result.append_null(), + } + } + Ok(ColumnarValue::Array(Arc::new(result.finish()))) + }), + ) + } +} + pub struct NativeHistogramToString; impl NativeHistogramToString { @@ -845,19 +923,38 @@ impl NativeHistogramAggregateAccumulator { } } - fn push_histogram(&mut self, histogram: NativeHistogram, count: u64) { - if self.dropped_incompatible { - return; + fn push_histogram(&mut self, histogram: NativeHistogram, count: u64) -> DfResult<()> { + if self.kind.needs_count() && count == 0 { + return Ok(()); } self.observe_reset_hints( histogram.reset_hint == COUNTER_RESET_HINT, histogram.reset_hint == NOT_COUNTER_RESET_HINT, ); + if self.dropped_incompatible { + return Ok(()); + } + let combined_count = if self.kind.needs_count() { + self.count.checked_add(count).ok_or_else(|| { + DataFusionError::Execution(format!( + "{}: native histogram sample count overflow", + self.kind.name() + )) + })? + } else { + self.count + }; let value = match self.value.take() { Some(value) => { record_custom_reconciliation(&self.collector, self.kind.name(), &value, &histogram); - match value.add(&histogram) { + let combined = match self.kind { + NativeHistogramAggregateKind::Sum => value.add(&histogram), + NativeHistogramAggregateKind::Avg => { + weighted_histogram_mean(value, self.count, histogram, count, combined_count) + } + }; + match combined { Some(value) => Some(value), None => { self.record_incompatible(); @@ -869,8 +966,9 @@ impl NativeHistogramAggregateAccumulator { }; if !self.dropped_incompatible { self.value = value; - self.count += count; + self.count = combined_count; } + Ok(()) } fn mark_incompatible(&mut self) { @@ -891,24 +989,48 @@ impl NativeHistogramAggregateAccumulator { } } -fn range_sum_histograms( +fn weighted_histogram_mean( + left: NativeHistogram, + left_count: u64, + right: NativeHistogram, + right_count: u64, + total_count: u64, +) -> Option { + let total_count = total_count as f64; + left.scale(left_count as f64 / total_count) + .add(&right.scale(right_count as f64 / total_count)) +} + +fn range_fold_histograms( samples: Vec, + kind: NativeHistogramAggregateKind, name: &'static str, collector: &Option, ) -> Option { + if samples + .iter() + .any(|histogram| histogram.reset_hint == COUNTER_RESET_HINT) + && samples + .iter() + .any(|histogram| histogram.reset_hint == NOT_COUNTER_RESET_HINT) + { + record_counter_reset_contradiction_warning(collector, name); + } + let mut value = None; - let mut counter_reset_seen = false; - let mut not_counter_reset_seen = false; + let mut count = 0u64; for histogram in samples { - counter_reset_seen |= histogram.reset_hint == COUNTER_RESET_HINT; - not_counter_reset_seen |= histogram.reset_hint == NOT_COUNTER_RESET_HINT; - if counter_reset_seen && not_counter_reset_seen { - record_counter_reset_contradiction_warning(collector, name); - } value = match value { Some(value) => { record_custom_reconciliation(collector, name, &value, &histogram); - match value.add(&histogram) { + let next_count = count.checked_add(1)?; + let combined = match kind { + NativeHistogramAggregateKind::Sum => value.add(&histogram), + NativeHistogramAggregateKind::Avg => { + weighted_histogram_mean(value, count, histogram, 1, next_count) + } + }; + match combined { Some(value) => Some(value), None => { record_warning( @@ -923,6 +1045,7 @@ fn range_sum_histograms( } None => Some(histogram), }; + count = count.checked_add(1)?; } value } @@ -1016,14 +1139,18 @@ fn native_histogram_range_histogram( continue; }; let histogram = match kind { - NativeHistogramRangeHistogramKind::Sum => { - range_sum_histograms(samples, func_name, &collector) - } - NativeHistogramRangeHistogramKind::Avg => { - let count = samples.len(); - range_sum_histograms(samples, func_name, &collector) - .map(|histogram| histogram.divide_by(count as f64)) - } + NativeHistogramRangeHistogramKind::Sum => range_fold_histograms( + samples, + NativeHistogramAggregateKind::Sum, + func_name, + &collector, + ), + NativeHistogramRangeHistogramKind::Avg => range_fold_histograms( + samples, + NativeHistogramAggregateKind::Avg, + func_name, + &collector, + ), NativeHistogramRangeHistogramKind::Last => samples.last().cloned(), }; result.push(histogram); @@ -2113,7 +2240,7 @@ impl DfAccumulator for NativeHistogramAggregateAccumulator { let Some(histogram) = read_histogram(histograms, row)? else { continue; }; - self.push_histogram(histogram, 1); + self.push_histogram(histogram, 1)?; } Ok(()) @@ -2125,7 +2252,7 @@ impl DfAccumulator for NativeHistogramAggregateAccumulator { (_, false, None) => None, (NativeHistogramAggregateKind::Sum, false, value) => value, (NativeHistogramAggregateKind::Avg, false, Some(value)) if self.count > 0 => { - Some(value.divide_by(self.count as f64)) + Some(value) } (NativeHistogramAggregateKind::Avg, _, _) => None, }; @@ -2214,22 +2341,21 @@ impl DfAccumulator for NativeHistogramAggregateAccumulator { })?; for row in 0..histograms.len() { - if dropped.value(row) { - self.mark_incompatible(); - continue; - } - if self.dropped_incompatible { - continue; - } self.observe_reset_hints( counter_reset_seen.value(row), not_counter_reset_seen.value(row), ); + if dropped.value(row) { + self.mark_incompatible(); + } + if self.dropped_incompatible { + continue; + } let Some(histogram) = read_histogram(histograms, row)? else { continue; }; let count = counts.map(|counts| counts.value(row)).unwrap_or(1); - self.push_histogram(histogram, count); + self.push_histogram(histogram, count)?; } Ok(()) @@ -2799,6 +2925,15 @@ mod tests { read_histogram(&array, 0).unwrap().unwrap() } + fn evaluated_histogram( + accumulator: &mut NativeHistogramAggregateAccumulator, + ) -> NativeHistogram { + let ScalarValue::Struct(array) = accumulator.evaluate().unwrap() else { + panic!("native histogram accumulator returned a non-struct value"); + }; + read_histogram(&array, 0).unwrap().unwrap() + } + fn histogram_range_input(values: Vec>) -> Vec { let timestamps = Arc::new(TimestampMillisecondArray::from_iter( (0..values.len()).map(|idx| Some((idx as i64 + 1) * 1000)), @@ -3195,6 +3330,80 @@ mod tests { assert_eq!(result.positive_buckets, vec![2.0]); } + #[test] + fn histogram_averages_avoid_sum_overflow() { + let large = sample_histogram(1.0e308, 1.0e308, vec![1.0e308]); + + let range_average = run_histogram_range_udf( + NativeHistogramAvgOverTime::scalar_udf(), + vec![large.clone(), large.clone()], + ); + assert_eq!(range_average.count, 1.0e308); + assert_eq!(range_average.sum, 1.0e308); + assert_eq!(range_average.positive_buckets, vec![1.0e308]); + + let mut aggregate = + NativeHistogramAggregateAccumulator::new(NativeHistogramAggregateKind::Avg, None); + aggregate.push_histogram(large.clone(), 1).unwrap(); + aggregate.push_histogram(large, 1).unwrap(); + let aggregate_average = evaluated_histogram(&mut aggregate); + assert_eq!(aggregate_average.count, 1.0e308); + assert_eq!(aggregate_average.sum, 1.0e308); + assert_eq!(aggregate_average.positive_buckets, vec![1.0e308]); + } + + #[test] + fn histogram_average_partial_states_are_weighted() { + let mut first = + NativeHistogramAggregateAccumulator::new(NativeHistogramAggregateKind::Avg, None); + first + .push_histogram(sample_histogram(1.0, 1.0, vec![1.0]), 1) + .unwrap(); + first + .push_histogram(sample_histogram(3.0, 3.0, vec![3.0]), 1) + .unwrap(); + let first_state = first + .state() + .unwrap() + .into_iter() + .map(|value| value.to_array_of_size(1).unwrap()) + .collect::>(); + + let mut second = + NativeHistogramAggregateAccumulator::new(NativeHistogramAggregateKind::Avg, None); + second + .push_histogram(sample_histogram(8.0, 8.0, vec![8.0]), 1) + .unwrap(); + let second_state = second + .state() + .unwrap() + .into_iter() + .map(|value| value.to_array_of_size(1).unwrap()) + .collect::>(); + + let mut merged = + NativeHistogramAggregateAccumulator::new(NativeHistogramAggregateKind::Avg, None); + merged.merge_batch(&first_state).unwrap(); + merged.merge_batch(&second_state).unwrap(); + let average = evaluated_histogram(&mut merged); + assert_eq!(average.count, 4.0); + assert_eq!(average.sum, 4.0); + assert_eq!(average.positive_buckets, vec![4.0]); + } + + #[test] + fn histogram_average_rejects_sample_count_overflow() { + let mut aggregate = + NativeHistogramAggregateAccumulator::new(NativeHistogramAggregateKind::Avg, None); + aggregate.value = Some(sample_histogram(1.0, 1.0, vec![1.0])); + aggregate.count = u64::MAX; + + let error = aggregate + .push_histogram(sample_histogram(1.0, 1.0, vec![1.0]), 1) + .unwrap_err(); + assert!(error.to_string().contains("sample count overflow")); + } + #[test] fn presence_only_range_functions_preserve_null_semantics() { let first = sample_histogram(1.0, 1.0, vec![1.0]); @@ -3311,25 +3520,32 @@ mod tests { } #[test] - fn subtraction_records_reset_hint_contradictions() { - let mut left = sample_histogram(2.0, 2.0, vec![2.0]); - left.reset_hint = COUNTER_RESET_HINT; - let mut right = sample_histogram(1.0, 1.0, vec![1.0]); - right.reset_hint = NOT_COUNTER_RESET_HINT; - let collector = PromqlAnnotationCollector::default(); + fn subtraction_records_reset_hint_contradictions_for_incompatible_histograms() { + for incompatible in [false, true] { + let mut left = sample_histogram(2.0, 2.0, vec![2.0]); + left.reset_hint = COUNTER_RESET_HINT; + let mut right = sample_histogram(1.0, 1.0, vec![1.0]); + right.reset_hint = NOT_COUNTER_RESET_HINT; + if incompatible { + right.schema = CUSTOM_BUCKETS_SCHEMA; + right.custom_values = vec![1.0]; + } + let collector = PromqlAnnotationCollector::default(); - run_histogram_udf( - NativeHistogramSub::scalar_udf_with_collector(Some(collector.clone())), - vec![ - ColumnarValue::Array(build_histogram_array(&[Some(left)])), - ColumnarValue::Array(build_histogram_array(&[Some(right)])), - ], - ); + run_udf( + NativeHistogramSub::scalar_udf_with_collector(Some(collector.clone())), + vec![ + ColumnarValue::Array(build_histogram_array(&[Some(left)])), + ColumnarValue::Array(build_histogram_array(&[Some(right)])), + ], + native_histogram_arrow_type(), + ); - assert!(collected_warnings(&collector).contains(&format!( - "{}: native histogram counter reset hints contradict", - NativeHistogramSub::name() - ))); + assert!(collected_warnings(&collector).contains(&format!( + "{}: native histogram counter reset hints contradict", + NativeHistogramSub::name() + ))); + } } #[test] @@ -3343,8 +3559,9 @@ mod tests { let range_collector = PromqlAnnotationCollector::default(); assert!( - range_sum_histograms( + range_fold_histograms( vec![reset.clone(), unknown.clone(), not_reset.clone()], + NativeHistogramAggregateKind::Sum, NativeHistogramSumOverTime::name(), &Some(range_collector.clone()), ) @@ -3360,8 +3577,8 @@ mod tests { NativeHistogramAggregateKind::Avg, ] { let mut first_partial = NativeHistogramAggregateAccumulator::new(kind, None); - first_partial.push_histogram(reset.clone(), 1); - first_partial.push_histogram(unknown.clone(), 1); + first_partial.push_histogram(reset.clone(), 1).unwrap(); + first_partial.push_histogram(unknown.clone(), 1).unwrap(); let first_states = first_partial .state() .unwrap() @@ -3370,7 +3587,7 @@ mod tests { .collect::>(); let mut second_partial = NativeHistogramAggregateAccumulator::new(kind, None); - second_partial.push_histogram(not_reset.clone(), 1); + second_partial.push_histogram(not_reset.clone(), 1).unwrap(); let second_states = second_partial .state() .unwrap() @@ -3390,6 +3607,115 @@ mod tests { } } + #[test] + fn incompatible_aggregates_do_not_hide_reset_hint_contradictions() { + let mut reset = sample_histogram(1.0, 1.0, vec![1.0]); + reset.reset_hint = COUNTER_RESET_HINT; + let mut incompatible_reset = sample_histogram(2.0, 2.0, vec![2.0]); + incompatible_reset.schema = CUSTOM_BUCKETS_SCHEMA; + incompatible_reset.custom_values = vec![1.0]; + incompatible_reset.reset_hint = COUNTER_RESET_HINT; + let mut not_reset = sample_histogram(3.0, 3.0, vec![3.0]); + not_reset.reset_hint = NOT_COUNTER_RESET_HINT; + + for kind in [ + NativeHistogramAggregateKind::Sum, + NativeHistogramAggregateKind::Avg, + ] { + for histograms in [ + [reset.clone(), incompatible_reset.clone(), not_reset.clone()], + [reset.clone(), not_reset.clone(), incompatible_reset.clone()], + ] { + let collector = PromqlAnnotationCollector::default(); + let mut aggregate = + NativeHistogramAggregateAccumulator::new(kind, Some(collector.clone())); + for histogram in histograms { + aggregate.push_histogram(histogram, 1).unwrap(); + } + + assert!(aggregate.dropped_incompatible); + let warnings = collected_warnings(&collector); + assert!(warnings.contains(&format!( + "{}: dropped native histogram aggregate with incompatible schemas", + kind.name() + ))); + assert!(warnings.contains(&format!( + "{}: native histogram counter reset hints contradict", + kind.name() + ))); + } + + let mut dropped_partial = NativeHistogramAggregateAccumulator::new(kind, None); + dropped_partial.push_histogram(reset.clone(), 1).unwrap(); + dropped_partial + .push_histogram(incompatible_reset.clone(), 1) + .unwrap(); + let dropped_states = dropped_partial + .state() + .unwrap() + .into_iter() + .map(|value| value.to_array_of_size(1).unwrap()) + .collect::>(); + + let mut opposing_partial = NativeHistogramAggregateAccumulator::new(kind, None); + opposing_partial + .push_histogram(not_reset.clone(), 1) + .unwrap(); + let opposing_states = opposing_partial + .state() + .unwrap() + .into_iter() + .map(|value| value.to_array_of_size(1).unwrap()) + .collect::>(); + + for (first, second) in [ + (&dropped_states, &opposing_states), + (&opposing_states, &dropped_states), + ] { + let collector = PromqlAnnotationCollector::default(); + let mut merged = + NativeHistogramAggregateAccumulator::new(kind, Some(collector.clone())); + merged.merge_batch(first).unwrap(); + merged.merge_batch(second).unwrap(); + + assert!(merged.dropped_incompatible); + assert!(collected_warnings(&collector).contains(&format!( + "{}: native histogram counter reset hints contradict", + kind.name() + ))); + } + } + + for (kind, name) in [ + ( + NativeHistogramAggregateKind::Sum, + NativeHistogramSumOverTime::name(), + ), + ( + NativeHistogramAggregateKind::Avg, + NativeHistogramAvgOverTime::name(), + ), + ] { + let collector = PromqlAnnotationCollector::default(); + assert!( + range_fold_histograms( + vec![reset.clone(), incompatible_reset.clone(), not_reset.clone()], + kind, + name, + &Some(collector.clone()), + ) + .is_none() + ); + let warnings = collected_warnings(&collector); + assert!(warnings.contains(&format!( + "{name}: dropped native histogram range with incompatible schemas" + ))); + assert!(warnings.contains(&format!( + "{name}: native histogram counter reset hints contradict" + ))); + } + } + #[test] fn histogram_aggregate_accumulator_accounts_for_heap_allocations() { let histogram = sample_histogram(3.0, 3.0, vec![1.0, 2.0]); @@ -3402,7 +3728,7 @@ mod tests { NativeHistogramAggregateAccumulator::new(NativeHistogramAggregateKind::Sum, None); let empty_size = accumulator.size(); - accumulator.push_histogram(histogram, 1); + accumulator.push_histogram(histogram, 1).unwrap(); assert!(heap_size > 0); assert_eq!(accumulator.size(), empty_size + heap_size); diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 0933bbcaca4..5dfb1a537a8 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -48,6 +48,7 @@ use datafusion::sql::TableReference; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; use datafusion_common::{DFSchema, NullEquality}; use datafusion_expr::expr::WindowFunctionParams; +use datafusion_expr::expr_fn::when; use datafusion_expr::utils::{conjunction, disjunction}; use datafusion_expr::{ ExprSchemable, Literal, Projection, SortExpr, TableScan, TableSource, col, lit, @@ -64,17 +65,17 @@ use promql::extension_plan::{ use promql::functions::{ AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, MixedRange, - NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAvg, - NativeHistogramAvgOverTime, NativeHistogramChanges, NativeHistogramCount, - NativeHistogramCountOverTime, NativeHistogramDelta, NativeHistogramDivScalar, - NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, NativeHistogramIDelta, - NativeHistogramIRate, NativeHistogramIncrease, NativeHistogramLastOverTime, - NativeHistogramMulScalar, NativeHistogramNeg, NativeHistogramNotEq, - NativeHistogramPresentOverTime, NativeHistogramQuantile, NativeHistogramRate, - NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar, - NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, PredictLinear, - PresentOverTime, QuantileOverTime, Rate, Resets, Round, StddevOverTime, StdvarOverTime, - SumOverTime, quantile_udaf, + NativeHistogramAbsentOverTime, NativeHistogramAdd, NativeHistogramAggAvg, + NativeHistogramAggSum, NativeHistogramAvg, NativeHistogramAvgOverTime, NativeHistogramChanges, + NativeHistogramCount, NativeHistogramCountOverTime, NativeHistogramDelta, + NativeHistogramDivScalar, NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, + NativeHistogramIDelta, NativeHistogramIRate, NativeHistogramIncrease, + NativeHistogramLastOverTime, NativeHistogramMulScalar, NativeHistogramNeg, + NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile, + NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, + NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, + NativeHistogramToString, PredictLinear, PresentOverTime, PromqlFloatToString, QuantileOverTime, + Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf, }; use promql_parser::label::{METRIC_NAME, MatchOp, Matcher, Matchers}; use promql_parser::parser::token::TokenType; @@ -654,9 +655,21 @@ impl PromPlanner { // calculate columns to group by // Need to append time index column into group by columns let mut group_exprs = self.agg_modifier_to_col(input.schema(), modifier, true)?; + let mixed_sample_columns = + Self::alternative_sample_columns(input.schema(), &self.ctx.field_columns) + .map(|(float, histogram)| (float.to_string(), histogram.to_string())); + // Aggregates over native histogram inputs may drop every sample in a group + // (e.g. `min` over histogram-only samples, or `sum` over histograms with + // incompatible schemas) and leave a NULL-valued group row behind. Compute this + // before `create_aggregate_exprs` mutates `ctx.field_columns`. + let preserve_any_value = mixed_sample_columns.is_some(); + let has_native_histogram = preserve_any_value + || self.all_field_columns_are_native_histograms(input.schema()); // convert op and value columns to aggregate exprs let (mut aggr_exprs, prev_field_exprs) = self.create_aggregate_exprs(*op, param, &input)?; + let prev_field_exprs = + normalize_cols(prev_field_exprs, &input).context(DataFusionPlanningSnafu)?; let keep_tsid = op.id() != token::T_COUNT_VALUES && input_has_tsid @@ -680,16 +693,28 @@ impl PromPlanner { let label = Self::get_param_value_as_str(*op, param)?; // `count_values` must be grouped by fields, // and project the fields to the new label. - group_exprs.extend(prev_field_exprs.clone()); + let count_value_exprs = prev_field_exprs.iter().map(|expr| { + match expr { + DfExpr::Column(column) => DfExpr::Column(column.clone()), + _ => DfExpr::Column(Column::from_name(expr.schema_name().to_string())), + } + .alias(label) + }); + let aggregate_group_exprs = group_exprs + .iter() + .cloned() + .chain(prev_field_exprs.clone()) + .collect::>(); + group_exprs.push(col(label)); let project_fields = self .create_field_column_exprs()? .into_iter() .chain(self.create_tag_column_exprs()?) .chain(Some(self.create_time_index_column_expr()?)) - .chain(prev_field_exprs.into_iter().map(|expr| expr.alias(label))); + .chain(count_value_exprs); builder - .aggregate(group_exprs.clone(), aggr_exprs) + .aggregate(aggregate_group_exprs, aggr_exprs) .context(DataFusionPlanningSnafu)? .project(project_fields) .context(DataFusionPlanningSnafu)? @@ -699,6 +724,59 @@ impl PromPlanner { .context(DataFusionPlanningSnafu)? }; + let builder = if let Some((float, histogram)) = mixed_sample_columns { + let builder = match op.id() { + token::T_SUM | token::T_AVG => builder + .filter(self.mixed_aggregate_filter_expr(*op, &float, &histogram)?) + .context(DataFusionPlanningSnafu)?, + token::T_MIN + | token::T_MAX + | token::T_STDDEV + | token::T_STDVAR + | token::T_QUANTILE => builder + .filter(self.mixed_ignored_histogram_filter_expr(*op, &histogram)?) + .context(DataFusionPlanningSnafu)?, + _ => builder, + }; + + match op.id() { + token::T_SUM + | token::T_AVG + | token::T_MIN + | token::T_MAX + | token::T_STDDEV + | token::T_STDVAR + | token::T_QUANTILE => { + let project_fields = self + .create_field_column_exprs()? + .into_iter() + .chain(self.create_tag_column_exprs()?) + .chain(self.ctx.use_tsid.then_some(DfExpr::Column( + Column::from_name(DATA_SCHEMA_TSID_COLUMN_NAME), + ))) + .chain(Some(self.create_time_index_column_expr()?)); + builder + .project(project_fields) + .context(DataFusionPlanningSnafu)? + } + _ => builder, + } + } else { + builder + }; + + // Drop group rows whose every aggregated sample was discarded (NULL), so that + // e.g. `group(min(native_histogram))` doesn't resurrect groups Prometheus + // considers unseen. For alternative float/histogram fields keep the row if any + // field survived. + let builder = if has_native_histogram { + builder + .filter(self.create_empty_values_filter_expr(preserve_any_value)?) + .context(DataFusionPlanningSnafu)? + } else { + builder + }; + let sort_expr = group_exprs.into_iter().map(|expr| expr.sort(true, false)); builder @@ -731,6 +809,32 @@ impl PromPlanner { let group_exprs = self.agg_modifier_to_col(input.schema(), modifier, false)?; + let mut input = input; + if let Some((float_column, histogram_column)) = + Self::alternative_sample_columns(input.schema(), &self.ctx.field_columns) + .map(|(float, histogram)| (float.to_string(), histogram.to_string())) + { + let drop_histogram = DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramDrop::bool_false_udf( + format!( + "{}: dropped native histogram samples because this aggregation is not supported for native histograms", + op + ), + self.promql_annotations.clone(), + )), + args: vec![col(&histogram_column)], + }); + let keep_float = when(col(&histogram_column).is_not_null(), drop_histogram) + .otherwise(col(&float_column).is_not_null()) + .context(DataFusionPlanningSnafu)?; + input = LogicalPlanBuilder::from(input) + .filter(keep_float) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + self.ctx.field_columns = vec![float_column]; + } + if self.all_field_columns_are_native_histograms(input.schema()) { let promql_annotations = self.promql_annotations.clone(); let input = self.projection_for_each_field_column(input, |col| { @@ -4123,9 +4227,11 @@ impl PromPlanner { param: &Option>, input_plan: &LogicalPlan, ) -> Result<(Vec, Vec)> { - let mut non_col_args = Vec::new(); + let mixed_sample_columns = + Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns) + .map(|(float, histogram)| (float.to_string(), histogram.to_string())); let is_group_agg = op.id() == token::T_GROUP; - if is_group_agg { + if is_group_agg && mixed_sample_columns.is_none() { ensure!( self.ctx.field_columns.len() == 1, MultiFieldsNotSupportedSnafu { @@ -4133,46 +4239,28 @@ impl PromPlanner { } ); } - let aggr = match op.id() { - token::T_SUM => sum_udaf(), - token::T_QUANTILE => { - let q = - Self::get_param_as_literal_expr(param, Some(op), Some(ArrowDataType::Float64))?; - non_col_args.push(q); - quantile_udaf() - } - token::T_AVG => avg_udaf(), - token::T_COUNT_VALUES | token::T_COUNT => count_udaf(), - token::T_MIN => min_udaf(), - token::T_MAX => max_udaf(), - // PromQL's `group()` aggregator produces 1 for each group. - // Use `max(1.0)` (per-group) to match semantics and output type (Float64). - token::T_GROUP => max_udaf(), - token::T_STDDEV => stddev_pop_udaf(), - token::T_STDVAR => var_pop_udaf(), - token::T_TOPK | token::T_BOTTOMK => UnsupportedExprSnafu { - name: format!("{op:?}"), - } - .fail()?, - _ => UnexpectedTokenSnafu { token: op }.fail()?, - }; + + if let Some((float, histogram)) = mixed_sample_columns { + return self.create_mixed_aggregate_exprs(op, param, &float, &histogram); + } + + if self.all_field_columns_are_native_histograms(input_plan.schema()) { + return self.create_native_histogram_aggregate_exprs(op, input_plan); + } // perform aggregate operation to each value column - let exprs: Vec = self + let exprs = self .ctx .field_columns .iter() .map(|col| { - if is_group_agg { - aggr.call(vec![lit(1_f64)]) - } else { - non_col_args.push(DfExpr::Column(Column::from_name(col))); - let expr = aggr.call(non_col_args.clone()); - non_col_args.pop(); - expr - } + Self::create_numeric_aggregate_expr( + op, + param, + DfExpr::Column(Column::from_name(col)), + ) }) - .collect::>(); + .collect::>>()?; // if the aggregator is `count_values`, it must be grouped by current fields. let prev_field_exprs = if op.id() == token::T_COUNT_VALUES { @@ -4208,6 +4296,273 @@ impl PromPlanner { Ok((exprs, prev_field_exprs)) } + fn create_numeric_aggregate_expr( + op: TokenType, + param: &Option>, + input: DfExpr, + ) -> Result { + let expr = match op.id() { + token::T_SUM => sum_udaf().call(vec![input]), + token::T_QUANTILE => { + let q = + Self::get_param_as_literal_expr(param, Some(op), Some(ArrowDataType::Float64))?; + quantile_udaf().call(vec![q, input]) + } + token::T_AVG => avg_udaf().call(vec![input]), + token::T_COUNT_VALUES | token::T_COUNT => count_udaf().call(vec![input]), + token::T_MIN => min_udaf().call(vec![input]), + token::T_MAX => max_udaf().call(vec![input]), + // PromQL's `group()` aggregator produces 1 for each group. + // Use `max(1.0)` (per-group) to match semantics and output type (Float64). + token::T_GROUP => max_udaf().call(vec![lit(1_f64)]), + token::T_STDDEV => stddev_pop_udaf().call(vec![input]), + token::T_STDVAR => var_pop_udaf().call(vec![input]), + token::T_TOPK | token::T_BOTTOMK => { + return UnsupportedExprSnafu { + name: format!("{op:?}"), + } + .fail(); + } + _ => return UnexpectedTokenSnafu { token: op }.fail(), + }; + Ok(expr) + } + + fn create_mixed_aggregate_exprs( + &mut self, + op: TokenType, + param: &Option>, + float_column: &str, + histogram_column: &str, + ) -> Result<(Vec, Vec)> { + let float_input = DfExpr::Column(Column::from_name(float_column)); + let histogram_input = DfExpr::Column(Column::from_name(histogram_column)); + let float_count = count_udaf().call(vec![float_input.clone()]); + let histogram_count = count_udaf().call(vec![histogram_input.clone()]); + let mixed_sample_value = || { + DfExpr::ScalarFunction(ScalarFunction { + func: coalesce(), + args: vec![ + DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(PromqlFloatToString::scalar_udf()), + args: vec![float_input.clone()], + }), + DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramToString::scalar_udf()), + args: vec![histogram_input.clone()], + }), + ], + }) + }; + + let (exprs, prev_field_exprs, field_columns) = match op.id() { + token::T_SUM | token::T_AVG => ( + vec![ + Self::create_numeric_aggregate_expr(op, param, float_input)? + .alias(float_column), + self.create_native_histogram_aggregate_expr(op, histogram_column)?, + float_count.alias(Self::mixed_sample_count_name(float_column)), + histogram_count.alias(Self::mixed_sample_count_name(histogram_column)), + ], + vec![], + vec![float_column.to_string(), histogram_column.to_string()], + ), + token::T_COUNT => { + let present = when( + float_input + .clone() + .is_not_null() + .or(histogram_input.clone().is_not_null()), + lit(1_i64), + ) + .otherwise(lit(ScalarValue::Int64(None))) + .context(DataFusionPlanningSnafu)?; + ( + vec![count_udaf().call(vec![present]).alias(float_column)], + vec![], + vec![float_column.to_string()], + ) + } + token::T_GROUP => ( + vec![max_udaf().call(vec![lit(1_f64)]).alias(float_column)], + vec![], + vec![float_column.to_string()], + ), + token::T_COUNT_VALUES => { + let value = mixed_sample_value(); + ( + vec![count_udaf().call(vec![value.clone()]).alias(float_column)], + vec![value], + vec![float_column.to_string()], + ) + } + token::T_MIN | token::T_MAX | token::T_STDDEV | token::T_STDVAR | token::T_QUANTILE => { + ( + vec![ + Self::create_numeric_aggregate_expr(op, param, float_input)? + .alias(float_column), + histogram_count.alias(Self::mixed_sample_count_name(histogram_column)), + ], + vec![], + vec![float_column.to_string()], + ) + } + token::T_TOPK | token::T_BOTTOMK => { + return UnsupportedExprSnafu { + name: format!("{op:?}"), + } + .fail(); + } + _ => return UnexpectedTokenSnafu { token: op }.fail(), + }; + + self.ctx.field_columns = field_columns; + Ok((exprs, prev_field_exprs)) + } + + fn mixed_sample_count_column(column: &str) -> DfExpr { + DfExpr::Column(Column::from_name(Self::mixed_sample_count_name(column))) + } + + fn mixed_sample_count_name(column: &str) -> String { + format!("__promql_sample_count({column})") + } + + fn mixed_aggregate_filter_expr( + &self, + op: TokenType, + float_column: &str, + histogram_column: &str, + ) -> Result { + let float_count = Self::mixed_sample_count_column(float_column); + let histogram_count = Self::mixed_sample_count_column(histogram_column); + let mixed = float_count + .clone() + .gt(lit(0_i64)) + .and(histogram_count.clone().gt(lit(0_i64))); + let drop_mixed = DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramDrop::warning_bool_false_udf( + format!( + "{op}: dropped aggregation result containing both float and native histogram samples" + ), + self.promql_annotations.clone(), + )), + args: vec![float_count, histogram_count], + }); + + when(mixed, drop_mixed) + .otherwise(lit(true)) + .context(DataFusionPlanningSnafu) + } + + fn mixed_ignored_histogram_filter_expr( + &self, + op: TokenType, + histogram_column: &str, + ) -> Result { + let histogram_count = Self::mixed_sample_count_column(histogram_column); + let has_histograms = histogram_count.clone().gt(lit(0_i64)); + let record_info = DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramDrop::bool_true_udf( + format!( + "{op}: dropped native histogram samples because this aggregation is not supported for native histograms" + ), + self.promql_annotations.clone(), + )), + args: vec![histogram_count], + }); + + when(has_histograms, record_info) + .otherwise(lit(true)) + .context(DataFusionPlanningSnafu) + } + + fn create_native_histogram_aggregate_expr( + &self, + op: TokenType, + column: &str, + ) -> Result { + let input = DfExpr::Column(Column::from_name(column)); + let expr = match op.id() { + token::T_SUM => Arc::new(NativeHistogramAggSum::aggregate_udf_with_collector( + self.promql_annotations.clone(), + )) + .call(vec![input]) + .alias(column), + token::T_AVG => Arc::new(NativeHistogramAggAvg::aggregate_udf_with_collector( + self.promql_annotations.clone(), + )) + .call(vec![input]) + .alias(column), + token::T_COUNT_VALUES | token::T_COUNT => { + count_udaf().call(vec![input]).alias(column) + } + token::T_GROUP => max_udaf().call(vec![lit(1_f64)]).alias(column), + token::T_MIN + | token::T_MAX + | token::T_STDDEV + | token::T_STDVAR + | token::T_QUANTILE + | token::T_TOPK + | token::T_BOTTOMK => sum_udaf() + .call(vec![DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramDrop::float_null_udf( + format!( + "{op}: dropped native histogram samples because this aggregation is not supported for native histograms" + ), + self.promql_annotations.clone(), + )), + args: vec![input], + })]) + .alias(column), + _ => return UnexpectedTokenSnafu { token: op }.fail(), + }; + Ok(expr) + } + + fn create_native_histogram_aggregate_exprs( + &mut self, + op: TokenType, + input_plan: &LogicalPlan, + ) -> Result<(Vec, Vec)> { + let prev_field_exprs = if op.id() == token::T_COUNT_VALUES { + ensure!( + self.ctx.field_columns.len() == 1, + UnsupportedExprSnafu { + name: "count_values on multi-value input" + } + ); + self.ctx + .field_columns + .iter() + .map(|col| { + DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramToString::scalar_udf()), + args: vec![DfExpr::Column(Column::from_name(col))], + }) + }) + .collect::>() + } else { + vec![] + }; + + let exprs = self + .ctx + .field_columns + .iter() + .map(|col| self.create_native_histogram_aggregate_expr(op, col)) + .collect::>>()?; + + let normalized_exprs = + normalize_cols(exprs.iter().cloned(), input_plan).context(DataFusionPlanningSnafu)?; + self.ctx.field_columns = normalized_exprs + .into_iter() + .map(|expr| expr.schema_name().to_string()) + .collect(); + + Ok((exprs, prev_field_exprs)) + } + fn get_param_value_as_str(op: TokenType, param: &Option>) -> Result<&str> { let param = param .as_deref() @@ -6346,7 +6701,7 @@ mod test { use common_base::Plugins; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_query::native_histogram::{ - CounterResetHint, NativeHistogram, build_histogram_array, + CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, build_histogram_array, }; use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; @@ -6880,6 +7235,65 @@ mod test { (planner, plan) } + async fn mixed_aggregate_input(histograms: Vec) -> (PromPlanner, LogicalPlan) { + let float_field = format!("{OR_FLOAT_FIELD_PREFIX}0"); + let histogram_field = format!("{OR_HISTOGRAM_FIELD_PREFIX}0"); + let row_count = histograms.len() + 1; + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new( + "ts", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + ), + Field::new("k", ArrowDataType::Utf8, false), + Field::new(&float_field, ArrowDataType::Float64, true), + Field::new( + &histogram_field, + native_histogram_value_type().as_arrow_type(), + true, + ), + ])); + let mut histogram_values = Vec::with_capacity(row_count); + histogram_values.push(None); + histogram_values.extend(histograms.into_iter().map(Some)); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1; row_count])), + Arc::new(StringArray::from_iter_values( + (0..row_count).map(|row| format!("kind_{row}")), + )), + Arc::new(Float64Array::from_iter( + (0..row_count).map(|row| (row == 0).then_some(1.25)), + )), + build_histogram_array(&histogram_values), + ], + ) + .unwrap(); + let table = Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()); + let plan = LogicalPlanBuilder::scan("mixed", provider_as_source(table), None) + .unwrap() + .build() + .unwrap(); + let table_provider = build_test_table_provider_with_fields( + &[(DEFAULT_SCHEMA_NAME.to_string(), "dummy".to_string())], + &[], + ) + .await; + let planner = PromPlanner { + table_provider, + ctx: PromPlannerContext { + table_name: Some("mixed".to_string()), + time_index_column: Some("ts".to_string()), + field_columns: vec![float_field, histogram_field], + tag_columns: vec!["k".to_string()], + ..Default::default() + }, + promql_annotations: None, + }; + (planner, plan) + } + fn assert_no_internal_or_keys(schema: &DFSchema) { assert!( schema @@ -6906,6 +7320,26 @@ mod test { .collect() } + fn numeric_values(batches: &[RecordBatch], column: &str) -> Vec { + batches + .iter() + .flat_map(|batch| { + let values = datafusion::arrow::compute::cast( + batch.column_by_name(column).unwrap(), + &ArrowDataType::Float64, + ) + .unwrap(); + values + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + .collect::>() + }) + .collect() + } + fn histograms(batches: &[RecordBatch], column: &str) -> Vec { batches .iter() @@ -7325,18 +7759,41 @@ mod test { ]; let schema = Arc::new(Schema::new(columns)); let table_meta = TableMetaBuilder::empty() - .schema(schema) + .schema(schema.clone()) .primary_key_indices(vec![0]) .value_indices(vec![2, 3]) .next_column_id(1024) .build() .unwrap(); - let table_info = TableInfoBuilder::default() - .name(table_name) - .meta(table_meta) - .build() - .unwrap(); - let table = EmptyTable::from_table_info(&table_info); + let table_info = Arc::new( + TableInfoBuilder::default() + .name(table_name) + .meta(table_meta) + .build() + .unwrap(), + ); + let batch = RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![ + Arc::new(StringArray::from(vec!["float", "histogram"])), + Arc::new(TimestampMillisecondArray::from(vec![1_000, 1_000])), + build_histogram_array(&[None, Some(direct_or_histogram())]), + Arc::new(Float64Array::from(vec![Some(2.0), None])), + ], + ) + .unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + table_name, + GreptimeRecordBatch::from_df_record_batch(schema, batch), + 1024, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let table = Arc::new(Table::new( + table_info, + FilterPushDownType::Unsupported, + backing.data_source(), + )); assert!( catalog_list @@ -9812,6 +10269,50 @@ mod test { ); } + #[tokio::test] + async fn mixed_or_topk_bottomk_ignore_native_histograms() { + for op in ["topk", "bottomk"] { + let collector = PromqlAnnotationCollector::default(); + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + operator_table_provider(), + &operator_eval_stmt(&format!("{op}(1, lf or on(tag) lh)")), + &state, + Some(collector.clone()), + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + assert!( + plan.schema() + .fields() + .iter() + .all(|field| field.data_type() != &PromPlanner::native_histogram_arrow_type()), + "{plan:?}" + ); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(values(&batches, &float_field), vec![2.0], "{op}"); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + assert_eq!( + infos, + vec![format!( + "{op}: dropped native histogram samples because this aggregation is not supported for native histograms" + )] + ); + } + } + #[tokio::test] async fn native_histogram_scalar_is_ignored_before_scalar_calculate() { let plan = native_histogram_plan("scalar(some_metric)").await; @@ -10142,9 +10643,9 @@ mod test { #[tokio::test] async fn native_histogram_mixed_field_table_behaves() { - // Metric Engine exposes float and histogram samples as alternative nullable fields. - // Histogram functions must select the histogram field without adding a NULL float - // field that would make the empty-value filter reject every row. + // Exercise function planning after float and histogram samples have already been + // represented as alternative nullable fields. Histogram functions must select the + // histogram field without adding a NULL float field that would reject every row. let table_provider = build_test_mixed_native_histogram_table_provider("some_metric").await; let plan = PromPlanner::stmt_to_plan( table_provider, @@ -11126,15 +11627,14 @@ Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo: PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state()) .await .unwrap(); - let expected = "Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.greptime_value ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N, greptime_value:Float64;N]\ - \n Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series, prometheus_tsdb_head_series.greptime_value [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N, greptime_value:Float64;N]\ - \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64]\ - \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ - \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; + let expected = "Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]\ + \n Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N]\ + \n Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64]\ + \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n PromSeriesDivide: tags=[\"ip\"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n Filter: prometheus_tsdb_head_series.ip ~ Utf8(\"^(?:(10.0.160.237:8080|10.0.160.237:9090))$\") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]\ + \n TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"; assert_eq!(plan.display_indent_schema().to_string(), expected); } @@ -11176,15 +11676,14 @@ Filter: up.field_0 IS NOT NULL [timestamp:Timestamp(ms), field_0:Float64;N, foo: .unwrap(); let expected = r#" Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp [my_series:Int64, ip:Utf8, greptime_timestamp:Timestamp(ms)] - Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N] - Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, prometheus_tsdb_head_series.greptime_value ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N, greptime_value:Float64;N] - Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series, prometheus_tsdb_head_series.greptime_value [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N, greptime_value:Float64;N] - Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64] - PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - Filter: prometheus_tsdb_head_series.ip ~ Utf8("^(?:(10.0.160.237:8080|10.0.160.237:9090))$") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] - TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"#; + Sort: prometheus_tsdb_head_series.ip ASC NULLS LAST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS LAST, series ASC NULLS LAST [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N] + Projection: count(prometheus_tsdb_head_series.greptime_value), prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value AS series [count(prometheus_tsdb_head_series.greptime_value):Int64, ip:Utf8, greptime_timestamp:Timestamp(ms), series:Float64;N] + Aggregate: groupBy=[[prometheus_tsdb_head_series.ip, prometheus_tsdb_head_series.greptime_timestamp, prometheus_tsdb_head_series.greptime_value]], aggr=[[count(prometheus_tsdb_head_series.greptime_value)]] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N, count(prometheus_tsdb_head_series.greptime_value):Int64] + PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[greptime_timestamp] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + PromSeriesDivide: tags=["ip"] [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + Sort: prometheus_tsdb_head_series.ip ASC NULLS FIRST, prometheus_tsdb_head_series.greptime_timestamp ASC NULLS FIRST [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + Filter: prometheus_tsdb_head_series.ip ~ Utf8("^(?:(10.0.160.237:8080|10.0.160.237:9090))$") AND prometheus_tsdb_head_series.greptime_timestamp >= TimestampMillisecond(-999, None) AND prometheus_tsdb_head_series.greptime_timestamp <= TimestampMillisecond(100000000, None) [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N] + TableScan: prometheus_tsdb_head_series [ip:Utf8, greptime_timestamp:Timestamp(ms), greptime_value:Float64;N]"#; assert_eq!(format!("\n{}", plan.display_indent_schema()), expected); } @@ -11562,6 +12061,7 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom .contains("OR value fields have incompatible types") ); } + #[tokio::test] async fn test_or_with_histogram_quantile_missing_le_column() { let case = r#"histogram_quantile(0.99, non_existent_histogram_bucket) or normal_metric"#; @@ -11898,6 +12398,55 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom } } + #[tokio::test] + async fn test_histogram_only_min_drops_empty_aggregate_group() { + // `min` over native-histogram-only input drops every sample in the group, so the + // NULL-valued aggregate row must be filtered out. Otherwise an outer expression + // like `group()` resurrects the group Prometheus considers unseen. + let state = build_query_engine_state(); + for query in ["min(lh)", "group(min(lh))"] { + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 0, + "{query}" + ); + } + } + + #[tokio::test] + async fn test_mixed_min_drops_histogram_only_group() { + // With alternative float/histogram fields, `min by (tag)` keeps float-only groups + // (tag=a from `lf`) and drops histogram-only groups (tag=b from `lh`) instead of + // emitting a NULL-valued row for them. + let state = build_query_engine_state(); + let plan = PromPlanner::stmt_to_plan( + operator_table_provider(), + &operator_eval_stmt("min by (tag) (lf or on(tag) lh)"), + &state, + ) + .await + .unwrap(); + let float_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!(values(&batches, &float_field), vec![2.0]); + } + #[tokio::test] async fn test_mixed_or_can_feed_another_or() { let state = build_query_engine_state(); @@ -12171,6 +12720,306 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom } } + #[tokio::test] + async fn test_native_histogram_sum_and_avg_execute_real_batches() { + for op_name in ["sum", "avg"] { + for incompatible in [false, true] { + let mut second = direct_or_histogram(); + if incompatible { + second.schema = CUSTOM_BUCKETS_SCHEMA; + second.custom_values = vec![1.0]; + } + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = + mixed_aggregate_input(vec![direct_or_histogram(), second]).await; + planner.promql_annotations = Some(collector.clone()); + let histogram_column = planner.ctx.field_columns[1].clone(); + planner.ctx.field_columns = vec![histogram_column.clone()]; + let input = LogicalPlanBuilder::from(input) + .project([col("ts"), col(&histogram_column)]) + .unwrap() + .build() + .unwrap(); + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = + parser::parse(&format!("{op_name}(mixed)")).unwrap() + else { + unreachable!() + }; + let (aggregate_exprs, _) = + planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter(planner.create_empty_values_filter_expr(false).unwrap()) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(infos.is_empty()); + if incompatible { + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + assert!(warnings.iter().any(|warning| { + warning + == &format!( + "prom_native_histogram_agg_{op_name}: dropped native histogram aggregate with incompatible schemas" + ) + })); + } else { + let histograms = histograms(&batches, &histogram_column); + assert_eq!(histograms.len(), 1); + let expected = if op_name == "sum" { 2.0 } else { 1.0 }; + assert_eq!(histograms[0].count, expected); + assert_eq!(histograms[0].sum, expected); + assert!(warnings.is_empty()); + } + } + } + } + + #[tokio::test] + async fn test_canonical_mixed_count_group_and_count_values_execute() { + let state = build_query_engine_state(); + for (query, expected) in [ + ("count(some_metric)", vec![2.0]), + ("group(some_metric)", vec![1.0]), + (r#"count_values("sample", some_metric)"#, vec![1.0, 1.0]), + ] { + let plan = PromPlanner::stmt_to_plan( + build_test_mixed_native_histogram_table_provider("some_metric").await, + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + assert!( + plan.schema() + .fields() + .iter() + .all(|field| !field.name().starts_with("__promql_sample_count")), + "{query}: {plan:?}" + ); + let value_fields = plan + .schema() + .fields() + .iter() + .filter(|field| { + matches!( + field.data_type(), + ArrowDataType::Float64 | ArrowDataType::Int64 | ArrowDataType::UInt64 + ) || field.data_type() == &native_histogram_value_type().as_arrow_type() + }) + .collect::>(); + assert_eq!(value_fields.len(), 1, "{query}: {plan:?}"); + assert_ne!( + value_fields[0].data_type(), + &native_histogram_value_type().as_arrow_type(), + "{query}: {plan:?}" + ); + let value_column = value_fields[0].name().clone(); + + let (_, batches) = execute(plan, &state).await; + let mut actual = numeric_values(&batches, &value_column); + actual.sort_by(f64::total_cmp); + assert_eq!(actual, expected, "{query}"); + + if query.starts_with("count_values") { + let mut sample_labels = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("sample") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten() + .map(str::to_string) + }) + .collect::>(); + sample_labels.sort(); + let mut expected_labels = + vec!["2".to_string(), direct_or_histogram().promql_string()]; + expected_labels.sort(); + assert_eq!(sample_labels, expected_labels); + } + } + } + + #[tokio::test] + async fn test_mixed_or_sum_aggregates_each_sample_type() { + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = + parser::parse("sum(lhs)").unwrap() + else { + unreachable!() + }; + + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_direct_or(false).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts"), col("k")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([ + col(&float_column), + col(&histogram_column), + col("ts"), + col("k"), + ]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, &float_column), vec![1.25]); + let histogram = batches + .iter() + .find_map(|batch| { + let values = batch + .column_by_name(&histogram_column)? + .as_any() + .downcast_ref::()?; + (0..values.len()).find_map(|row| { + common_query::native_histogram::read_histogram(values, row).unwrap() + }) + }) + .unwrap(); + assert_eq!(histogram.count, 1.0); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_direct_or(false).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([col(&float_column), col(&histogram_column), col("ts")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert_eq!( + warnings, + vec![ + "sum: dropped aggregation result containing both float and native histogram samples" + ] + ); + } + + #[tokio::test] + async fn test_mixed_or_sum_drops_incompatible_mixed_group() { + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = + parser::parse("sum(lhs)").unwrap() + else { + unreachable!() + }; + let mut custom = direct_or_histogram(); + custom.schema = CUSTOM_BUCKETS_SCHEMA; + custom.custom_values = vec![1.0]; + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_aggregate_input(vec![direct_or_histogram(), custom]).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_aggregate_filter_expr(op, &float_column, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([col(&float_column), col(&histogram_column), col("ts")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.iter().any(|warning| { + warning + == "sum: dropped aggregation result containing both float and native histogram samples" + })); + } + + #[tokio::test] + async fn test_mixed_or_min_records_only_present_histograms() { + let PromExpr::Aggregate(AggregateExpr { op, param, .. }) = + parser::parse("min(lhs)").unwrap() + else { + unreachable!() + }; + let expected_info = "min: dropped native histogram samples because this aggregation is not supported for native histograms"; + + for (histograms, expected_infos) in [ + (vec![], vec![]), + (vec![direct_or_histogram()], vec![expected_info]), + ] { + let collector = PromqlAnnotationCollector::default(); + let (mut planner, input) = mixed_aggregate_input(histograms).await; + planner.promql_annotations = Some(collector.clone()); + let float_column = planner.ctx.field_columns[0].clone(); + let histogram_column = planner.ctx.field_columns[1].clone(); + let (aggregate_exprs, _) = planner.create_aggregate_exprs(op, ¶m, &input).unwrap(); + let plan = LogicalPlanBuilder::from(input) + .aggregate(vec![col("ts")], aggregate_exprs) + .unwrap() + .filter( + planner + .mixed_ignored_histogram_filter_expr(op, &histogram_column) + .unwrap(), + ) + .unwrap() + .project([col(&float_column), col("ts")]) + .unwrap() + .build() + .unwrap(); + + let (_, batches) = execute(plan, &build_query_engine_state()).await; + assert_eq!(values(&batches, &float_column), vec![1.25]); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + assert_eq!(infos, expected_infos); + } + } + #[tokio::test] async fn test_mixed_or_value_aliases_do_not_replace_labels() { let left = source( diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index 030ed27b4e3..e009d80954f 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -40,8 +40,8 @@ use promql::functions::{ NativeHistogramNotEq, NativeHistogramPresentOverTime, NativeHistogramQuantile, NativeHistogramRate, NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar, NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, - NativeHistogramToString, PredictLinear, PresentOverTime, QuantileOverTime, Rate, Resets, Round, - StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf, + NativeHistogramToString, PredictLinear, PresentOverTime, PromqlFloatToString, QuantileOverTime, + Rate, Resets, Round, StddevOverTime, StdvarOverTime, SumOverTime, quantile_udaf, }; use prost::Message; use session::context::QueryContextRef; @@ -188,6 +188,7 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { NativeHistogramDelta::scalar_udf(), NativeHistogramDivScalar::scalar_udf(), NativeHistogramDrop::bool_false_udf(String::new(), None), + NativeHistogramDrop::bool_true_udf(String::new(), None), NativeHistogramDrop::float_null_udf(String::new(), None), NativeHistogramEq::scalar_udf(), NativeHistogramFraction::scalar_udf(), @@ -211,6 +212,7 @@ impl SubstraitPlanDecoder for DefaultPlanDecoder { NativeHistogramSum::scalar_udf(), NativeHistogramSumOverTime::scalar_udf(), NativeHistogramToString::scalar_udf(), + PromqlFloatToString::scalar_udf(), ] { let _ = session_state.register_udf(Arc::new(udf)); } @@ -335,24 +337,47 @@ mod tests { None, ) .unwrap() + .aggregate( + Vec::::new(), + vec![ + Arc::new(NativeHistogramAggSum::aggregate_udf()) + .call(vec![col("histogram")]) + .alias("sum"), + Arc::new(NativeHistogramAggAvg::aggregate_udf()) + .call(vec![col("histogram")]) + .alias("avg"), + ], + ) + .unwrap() .project(vec![ Expr::ScalarFunction(ScalarFunction { func: Arc::new(NativeHistogramCount::scalar_udf()), - args: vec![col("histogram")], + args: vec![col("sum")], }), Expr::ScalarFunction(ScalarFunction { func: Arc::new(NativeHistogramDrop::bool_false_udf( "ignored annotation".to_string(), None, )), - args: vec![col("histogram")], + args: vec![col("sum")], + }), + Expr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramDrop::bool_true_udf( + "ignored annotation".to_string(), + None, + )), + args: vec![col("sum")], }), Expr::ScalarFunction(ScalarFunction { func: Arc::new(NativeHistogramDrop::float_null_udf( "ignored annotation".to_string(), None, )), - args: vec![col("histogram")], + args: vec![col("sum")], + }), + Expr::ScalarFunction(ScalarFunction { + func: Arc::new(PromqlFloatToString::scalar_udf()), + args: vec![lit(2.0)], }), ]) .unwrap() @@ -379,7 +404,11 @@ mod tests { let decoded = decoded.to_string(); assert!(decoded.contains("prom_native_histogram_count")); assert!(decoded.contains("prom_native_histogram_drop_bool")); + assert!(decoded.contains("prom_native_histogram_keep_bool")); assert!(decoded.contains("prom_native_histogram_drop_float")); + assert!(decoded.contains(NativeHistogramAggSum::name())); + assert!(decoded.contains(NativeHistogramAggAvg::name())); + assert!(decoded.contains(PromqlFloatToString::name())); let schema = Arc::new(Schema::new(vec![ Field::new( From 6249623cfb181b0222dfc03224945c625f326f1a Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Wed, 12 Aug 2026 08:15:12 +0000 Subject: [PATCH 34/92] chore: remove iceberg read (#8858) Signed-off-by: shuiyisong --- src/auth/src/lib.rs | 4 ++-- src/auth/src/permission.rs | 2 -- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/src/auth/src/lib.rs b/src/auth/src/lib.rs index d107068940e..d2a1c05b1cd 100644 --- a/src/auth/src/lib.rs +++ b/src/auth/src/lib.rs @@ -31,8 +31,8 @@ pub use common::{ }; pub use permission::{ ALL_ACTIONS, AccessMode, DASHBOARD_DELETE, DASHBOARD_QUERY, DASHBOARD_SAVE, - DefaultPermissionChecker, ICEBERG_READ, INFLUXDB_WRITE, JAEGER_QUERY, LOG_QUERY, LOG_WRITE, - OPENTSDB_WRITE, OTLP_WRITE, PIPELINE_DELETE, PIPELINE_INSERT, PIPELINE_QUERY, PROM_STORE_READ, + DefaultPermissionChecker, INFLUXDB_WRITE, JAEGER_QUERY, LOG_QUERY, LOG_WRITE, OPENTSDB_WRITE, + OTLP_WRITE, PIPELINE_DELETE, PIPELINE_INSERT, PIPELINE_QUERY, PROM_STORE_READ, PROM_STORE_WRITE, PROMQL_QUERY, PermissionAction, PermissionChecker, PermissionReq, PermissionResp, PermissionTableTarget, PermissionTableTargets, SEMANTIC_GRAPH_QUERY, }; diff --git a/src/auth/src/permission.rs b/src/auth/src/permission.rs index 984fe8ea4cf..901b935d6a2 100644 --- a/src/auth/src/permission.rs +++ b/src/auth/src/permission.rs @@ -139,7 +139,6 @@ pub const PROMQL_QUERY: PermissionAction = PermissionAction::read("promql.query" pub const LOG_QUERY: PermissionAction = PermissionAction::read("log.query"); pub const OPENTSDB_WRITE: PermissionAction = PermissionAction::write("opentsdb.write"); pub const INFLUXDB_WRITE: PermissionAction = PermissionAction::write("influxdb.write"); -pub const ICEBERG_READ: PermissionAction = PermissionAction::read("iceberg.read"); pub const PROM_STORE_WRITE: PermissionAction = PermissionAction::write("prom_store.write"); pub const PROM_STORE_READ: PermissionAction = PermissionAction::read("prom_store.read"); pub const OTLP_WRITE: PermissionAction = PermissionAction::write("otlp.write"); @@ -165,7 +164,6 @@ pub const ALL_ACTIONS: &[PermissionAction] = &[ LOG_QUERY, OPENTSDB_WRITE, INFLUXDB_WRITE, - ICEBERG_READ, PROM_STORE_WRITE, PROM_STORE_READ, OTLP_WRITE, From f2ffaef9dfee134fa46d54ba65e6f40b37982cef Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Wed, 12 Aug 2026 14:32:59 +0000 Subject: [PATCH 35/92] feat: update to pgwire 0.40.7 (#8860) --- Cargo.lock | 44 +++++++++++++++++++++++++------------------- 1 file changed, 25 insertions(+), 19 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4d0c271882b..5d8a2d1c6c0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1216,6 +1216,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" + [[package]] name = "base64-simd" version = "0.8.0" @@ -1614,7 +1620,7 @@ dependencies = [ "maybe-owned", "rustix 1.0.7", "rustix-linux-procfs", - "windows-sys 0.61.2", + "windows-sys 0.60.2", "winx", ] @@ -5120,7 +5126,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -5682,7 +5688,7 @@ checksum = "94e7099f6313ecacbe1256e8ff9d617b75d1bcb16a6fddef94866d225a01a14a" dependencies = [ "io-lifetimes 2.0.4", "rustix 1.0.7", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -6650,7 +6656,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core 0.57.0", + "windows-core 0.61.2", ] [[package]] @@ -7078,7 +7084,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20fd6de4ccfcc187e38bc21cfa543cb5a302cb86a8b114eb7f0bf0dc9f8ac00f" dependencies = [ "io-lifetimes 3.0.1", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -7116,7 +7122,7 @@ checksum = "e04d7f318608d35d4b61ddd75cbdaee86b023ebe2bd5a66ee0915f0bf93095a9" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -9087,7 +9093,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -9281,7 +9287,7 @@ version = "0.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77e878c846a8abae00dd069496dbe8751b16ac1c3d6bd2a7283a938e8228f90d" dependencies = [ - "proc-macro-crate 3.3.0", + "proc-macro-crate 1.3.1", "proc-macro2", "quote", "syn 2.0.117", @@ -10439,13 +10445,13 @@ dependencies = [ [[package]] name = "pgwire" -version = "0.40.1" +version = "0.40.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f690fba3b77370e278b606b75c6857d3f33f3e570781a75ef7516e76e44c0b19" +checksum = "8265901ede50d0879fe401c6fe282e7e4ff83ae10a48c1f8e4f89b92f7d5f604" dependencies = [ "async-trait", "aws-lc-rs", - "base64 0.22.1", + "base64 0.23.1", "bytes", "chrono", "derive-new", @@ -11762,7 +11768,7 @@ dependencies = [ "once_cell", "socket2 0.5.10", "tracing", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -12711,7 +12717,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.4.15", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -12724,7 +12730,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.9.4", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -12815,7 +12821,7 @@ dependencies = [ "security-framework 3.7.0", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -13728,7 +13734,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -14114,7 +14120,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -14724,7 +14730,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix 1.0.7", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -16916,7 +16922,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f3fd376f71958b862e7afb20cfe5a22830e1963462f3a17f49d82a6c1d1f42d" dependencies = [ "bitflags 2.12.1", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] From 546625c45af629b75f0ba28239a3a183e1183e19 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Thu, 13 Aug 2026 02:31:56 +0000 Subject: [PATCH 36/92] feat: embedded convention pack for the entity graph (prom/k8s, gen_ai naming) (#8854) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: embed the derivation conventions as data and adopt gen_ai entity naming Move the co-declared edge vocabulary, the agent-edge vocabulary and the virtual-destination candidates from Rust consts into an embedded conventions.yaml (include_str!), parsed once behind a LazyLock and validated against the entity-type grammar and the closed rel_type set; a broken file propagates as a plan error instead of panicking. The agent vocabulary entity types follow the GenAI semantic-convention namespace as written: gen_ai.agent / gen_ai.model / gen_ai.tool. Signed-off-by: Dennis Zhuang * feat: drop the tag requirement for entity identity columns Entity declarations no longer require id columns to be tag/primary-key columns; only column existence is validated. Trace pipelines flatten the identifying attributes (span_attributes.gen_ai.agent.id, ...) into field columns, so the tag rule locked real trace tables out of declaring entities while buying no correctness — the read-time derivation works on any column. Signed-off-by: Dennis Zhuang * feat: implicit declarations for well-known prometheus info metrics Tables stamped signal_type=metric + source=prometheus whose name matches the conventions.yaml whitelist (kube_pod_info, kube_node_info, kube_pod_owner, target_info) get implicit entity declarations: k8s.pod / k8s.node / k8s.workload with name-based identity and target_info's service / service.instance with the remaining tags as the descriptive snapshot. The existing co-declared vocabulary then derives runs_on and part_of from the same rows, so no new edge branch is needed. Explicit declarations of a type always suppress the implicit one, and the metric engine's physical table is excluded (it aggregates every logical table's columns). Signed-off-by: Dennis Zhuang * test: cover the prometheus conventions in sqlness and compact the graph cases Add the whitelisted-info-metric scenario (kube_pod_info, kube_pod_owner, target_info deriving runs_on / part_of, a non-whitelisted metric contributing nothing), fold the single-table calls, cross-table pairing and virtual-node cases into one trace scenario (they exercise the same union-before-join path), merge the two declaring-metric-table cases, and reuse one rename probe for both reserved names. Signed-off-by: Dennis Zhuang * fix: reject entity id columns without a stable string form Review follow-ups: the DDL check now validates against the schema and rejects binary-backed and nested types for identity columns (the derivation renders ids via CAST to Utf8, so the failure used to surface only when the graph was scanned); the agent sqlness case keeps its identity columns as fields to cover the relaxed tag rule end to end; stale tag-rule comments and a dangling const reference are cleaned up. Signed-off-by: Dennis Zhuang * fix: type-check every entity column role, not only ids The registry renders scope and descriptive values through the same CAST-to-string path as ids, so a binary-backed column in any role fails at scan time; the DDL check is now role-independent (and simpler). Signed-off-by: Dennis Zhuang * refactor: name the code-anchored vocabulary constants Entity types and edge attributes the derivation code itself anchors on (service, gen_ai.agent, calls, trace/attribute provenance) become constants in the conventions module; the rest of the vocabulary stays YAML-only data. ImplicitEntity is renamed PromImplicitEntity, and the implicit-declaration path logs each skip of a whitelisted info metric (wrong stamps, suppressed by an explicit declaration, missing id column) so a missing graph entity is diagnosable. Signed-off-by: Dennis Zhuang * refactor: single-source the graph constants The graph tables' column names move to common-catalog (the schemas catalog exposes and the plans operator builds must match column by column), and the conventions module now carries the complete built-in vocabulary — entity types, rel_types, provenances and connection types — with the embedded YAML validated by membership against it, so an edit drifting outside the vocabulary fails the conventions test instead of deriving nothing. Signed-off-by: Dennis Zhuang * fix: treat empty identity components as absent kube-state-metrics emits empty-string labels an entity id must not be built from: an unscheduled pod's node and an owner-less pod's owner_kind / owner_name. Standard Prometheus drops empty labels (they arrive as NULL and the existing predicate handles them), but other remote-write agents may keep them, which produced ghost entities with empty ids and false runs_on / part_of edges. Every identity predicate (registry, co-declared edges, span endpoints) now requires non-NULL and non-empty components through one shared helper. Signed-off-by: Dennis Zhuang * refactor: tighten the conventions DSL semantics Rename the co-declaration rule lists to what they are (co_declared_edges / trace_co_declared_edges — derivation rules, not a relation vocabulary), stop overstating the GenAI entity types (Greptime types derived from GenAI attributes; OTel defines no model/tool entities), move target_info's descriptive snapshot to service.instance (the remaining labels are the target's resource attributes, and instances would write conflicting snapshots onto the logical service), and extend the descriptor whitelist with the stable KSM sources: container info metrics (closing the k8s.pod contains k8s.container rule), kube_service_info (new k8s.service entity type) and the fuller descriptive label sets. Signed-off-by: Dennis Zhuang * fix: guard entity column types on ALTER as well ALTER MODIFY COLUMN could change a declared entity column to a type without a stable string form, deferring the failure to graph scan time; verify_alter now checks the post-alter schema. Dropping a declared column stays allowed — the read-time derivation skips the stale declaration, and semantic options cannot be altered off yet. Signed-off-by: Dennis Zhuang * feat: bridge traces and kube-state-metrics on the pod UID Trace-v1 tables now get implicit declarations from their flattened resource attributes (otlp_trace_entities in conventions.yaml): the service identity — replacing the hardcoded fallback — plus service.instance and k8s.pod, each applied only when its columns exist. A new co-declared rule derives service.instance runs_on k8s.pod, and the whitelisted kube-state-metrics pod identity switches from namespace+pod names to the UID, so the trace-side pod and every KSM descriptor land on one entity while names stay descriptive. This also removes pod identity from the multi-cluster same-name collision. The conventions rejection tests were passing for the wrong reason (a half-renamed fixture key failed deserialization before reaching any validation rule); they now assert the specific error each case targets. Sqlness covers the UID merge across descriptor tables, pod-contains- container, the k8s.service node, and the empty-uid/empty-node rows deriving nothing. Signed-off-by: Dennis Zhuang * test: cover the OTLP-to-graph chain end to end One real OTLP export must come out of semantic_relationships as the zero-configuration chain: service calls service, instance part_of service, instance runs_on pod (bridged by k8s.pod.uid). Resources without service.instance.id or k8s.pod.uid derive nothing extra. Signed-off-by: Dennis Zhuang * fix: identify k8s.service by UID Same reasoning as pods: a recreated same-name service must not merge into the old entity and same-named services across clusters must not collide; kube_service_info carries a stable uid and nothing joins on the service's name. Also drop a stale tag-rule mention from the option validation docs. Signed-off-by: Dennis Zhuang * chore: cut duplicated test coverage and redundant comments The trace service-fallback test collapsed into the resource-entities test (same synthesis path since the fallback moved to YAML; only the invalid-explicit-no-fallback case was distinct), role-duplicate and subsumed DDL cases are gone, the embedded-conventions test is just the parse (its assertions were decorative), and the YAML section comments no longer restate the struct docs. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- Cargo.lock | 2 + Cargo.toml | 1 + .../src/system_schema/semantic_graph.rs | 77 ++- src/common/catalog/src/consts.rs | 29 + src/frontend/src/instance/entity_graph.rs | 365 +++++++++-- src/operator/Cargo.toml | 2 + src/operator/src/error.rs | 8 - src/operator/src/statement/ddl.rs | 180 ++++-- src/operator/src/statement/semantic_graph.rs | 171 +++-- .../statement/semantic_graph/conventions.rs | 291 +++++++++ .../statement/semantic_graph/conventions.yaml | 115 ++++ .../statement/semantic_graph/relationships.rs | 323 +++++----- src/table/src/requests/semantic.rs | 6 +- tests-integration/tests/http.rs | 81 +++ .../common/system/semantic_graph.result | 586 ++++++++++-------- .../common/system/semantic_graph.sql | 346 ++++++----- 16 files changed, 1803 insertions(+), 780 deletions(-) create mode 100644 src/operator/src/statement/semantic_graph/conventions.rs create mode 100644 src/operator/src/statement/semantic_graph/conventions.yaml diff --git a/Cargo.lock b/Cargo.lock index 5d8a2d1c6c0..dcd23451407 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9936,7 +9936,9 @@ dependencies = [ "prost-types 0.14.1", "query", "regex", + "serde", "serde_json", + "serde_yaml_ng", "session", "snafu 0.8.6", "sql", diff --git a/Cargo.toml b/Cargo.toml index a138a880ec6..0e6da7888d3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -232,6 +232,7 @@ sea-query = "0.32" serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["float_roundtrip"] } serde_with = "3" +serde_yaml_ng = "0.10" simd-json = "0.15" similar-asserts = "1.6.0" smallvec = { version = "1", features = ["serde"] } diff --git a/src/catalog/src/system_schema/semantic_graph.rs b/src/catalog/src/system_schema/semantic_graph.rs index 028ce79b906..b25dbb03289 100644 --- a/src/catalog/src/system_schema/semantic_graph.rs +++ b/src/catalog/src/system_schema/semantic_graph.rs @@ -35,9 +35,14 @@ use std::sync::{Arc, LazyLock, Weak}; use common_catalog::consts::{ - DEFAULT_PRIVATE_SCHEMA_NAME, SEMANTIC_ENTITIES_TABLE_ID, + CONFIDENCE_COLUMN, DEFAULT_PRIVATE_SCHEMA_NAME, DST_ID_COLUMN, DST_TYPE_COLUMN, + DURATION_COUNT_COLUMN, DURATION_SUM_COLUMN, EDGE_ATTRIBUTES_COLUMN, ENTITY_DESCRIPTIVE_COLUMN, + ENTITY_ID_ATTRS_COLUMN, ENTITY_ID_COLUMN, ENTITY_SCOPE_COLUMN, ENTITY_TYPE_COLUMN, + ERROR_COUNT_COLUMN, FRESH_UNTIL_COLUMN, OBSERVED_AT_COLUMN, PROVENANCE_COLUMN, REL_TYPE_COLUMN, + REQUEST_COUNT_COLUMN, SEMANTIC_ENTITIES_TABLE_ID, SEMANTIC_ENTITIES_TABLE_NAME as SEMANTIC_ENTITIES, SEMANTIC_RELATIONSHIPS_TABLE_ID, - SEMANTIC_RELATIONSHIPS_TABLE_NAME as SEMANTIC_RELATIONSHIPS, + SEMANTIC_RELATIONSHIPS_TABLE_NAME as SEMANTIC_RELATIONSHIPS, SOURCE_TABLES_COLUMN, + SRC_ID_COLUMN, SRC_TYPE_COLUMN, WINDOW_END_COLUMN, WINDOW_START_COLUMN, }; use common_error::ext::BoxedError; use common_recordbatch::adapter::AsyncRecordBatchStreamAdapter; @@ -190,16 +195,16 @@ fn json() -> ConcreteDataType { /// - `source_tables` — JSON array of the telemetry tables that contributed this entity. static ENTITIES_SCHEMA: LazyLock = LazyLock::new(|| { Arc::new(Schema::new(vec![ - ColumnSchema::new("observed_at", ts(), false).with_time_index(true), - ColumnSchema::new("window_start", ts(), true), - ColumnSchema::new("window_end", ts(), true), - ColumnSchema::new("fresh_until", ts(), true), - ColumnSchema::new("entity_type", string(), false), - ColumnSchema::new("entity_id", string(), false), - ColumnSchema::new("entity_id_attrs", json(), true), - ColumnSchema::new("scope", string(), true), - ColumnSchema::new("descriptive", json(), true), - ColumnSchema::new("source_tables", json(), true), + ColumnSchema::new(OBSERVED_AT_COLUMN, ts(), false).with_time_index(true), + ColumnSchema::new(WINDOW_START_COLUMN, ts(), true), + ColumnSchema::new(WINDOW_END_COLUMN, ts(), true), + ColumnSchema::new(FRESH_UNTIL_COLUMN, ts(), true), + ColumnSchema::new(ENTITY_TYPE_COLUMN, string(), false), + ColumnSchema::new(ENTITY_ID_COLUMN, string(), false), + ColumnSchema::new(ENTITY_ID_ATTRS_COLUMN, json(), true), + ColumnSchema::new(ENTITY_SCOPE_COLUMN, string(), true), + ColumnSchema::new(ENTITY_DESCRIPTIVE_COLUMN, json(), true), + ColumnSchema::new(SOURCE_TABLES_COLUMN, json(), true), ])) }); @@ -232,22 +237,38 @@ static ENTITIES_SCHEMA: LazyLock = LazyLock::new(|| { /// `db.system`, `peer.service`. static RELATIONSHIPS_SCHEMA: LazyLock = LazyLock::new(|| { Arc::new(Schema::new(vec![ - ColumnSchema::new("observed_at", ts(), false).with_time_index(true), - ColumnSchema::new("window_start", ts(), true), - ColumnSchema::new("window_end", ts(), true), - ColumnSchema::new("fresh_until", ts(), true), - ColumnSchema::new("src_type", string(), false), - ColumnSchema::new("src_id", string(), false), - ColumnSchema::new("dst_type", string(), false), - ColumnSchema::new("dst_id", string(), false), - ColumnSchema::new("rel_type", string(), false), - ColumnSchema::new("provenance", string(), false), - ColumnSchema::new("confidence", ConcreteDataType::float64_datatype(), true), - ColumnSchema::new("request_count", ConcreteDataType::int64_datatype(), true), - ColumnSchema::new("error_count", ConcreteDataType::int64_datatype(), true), - ColumnSchema::new("duration_sum", ConcreteDataType::float64_datatype(), true), - ColumnSchema::new("duration_count", ConcreteDataType::int64_datatype(), true), - ColumnSchema::new("attributes", json(), true), + ColumnSchema::new(OBSERVED_AT_COLUMN, ts(), false).with_time_index(true), + ColumnSchema::new(WINDOW_START_COLUMN, ts(), true), + ColumnSchema::new(WINDOW_END_COLUMN, ts(), true), + ColumnSchema::new(FRESH_UNTIL_COLUMN, ts(), true), + ColumnSchema::new(SRC_TYPE_COLUMN, string(), false), + ColumnSchema::new(SRC_ID_COLUMN, string(), false), + ColumnSchema::new(DST_TYPE_COLUMN, string(), false), + ColumnSchema::new(DST_ID_COLUMN, string(), false), + ColumnSchema::new(REL_TYPE_COLUMN, string(), false), + ColumnSchema::new(PROVENANCE_COLUMN, string(), false), + ColumnSchema::new( + CONFIDENCE_COLUMN, + ConcreteDataType::float64_datatype(), + true, + ), + ColumnSchema::new( + REQUEST_COUNT_COLUMN, + ConcreteDataType::int64_datatype(), + true, + ), + ColumnSchema::new(ERROR_COUNT_COLUMN, ConcreteDataType::int64_datatype(), true), + ColumnSchema::new( + DURATION_SUM_COLUMN, + ConcreteDataType::float64_datatype(), + true, + ), + ColumnSchema::new( + DURATION_COUNT_COLUMN, + ConcreteDataType::int64_datatype(), + true, + ), + ColumnSchema::new(EDGE_ATTRIBUTES_COLUMN, json(), true), ])) }); diff --git a/src/common/catalog/src/consts.rs b/src/common/catalog/src/consts.rs index 6053d5a67c8..b976cbc256e 100644 --- a/src/common/catalog/src/consts.rs +++ b/src/common/catalog/src/consts.rs @@ -219,4 +219,33 @@ pub const SEMANTIC_RELATIONSHIPS_TABLE_NAME: &str = "semantic_relationships"; /// Physical table of hand-declared edges, unioned into the computed /// `semantic_relationships`; see [`is_ddl_reserved_table`] for its lifecycle. pub const SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME: &str = "semantic_relationships_declared"; + +// Column names of the graph tables: `catalog` exposes these schemas and the +// read-time plans in `operator` must project exactly them. +pub const OBSERVED_AT_COLUMN: &str = "observed_at"; +pub const WINDOW_START_COLUMN: &str = "window_start"; +pub const WINDOW_END_COLUMN: &str = "window_end"; +pub const FRESH_UNTIL_COLUMN: &str = "fresh_until"; +pub const ENTITY_TYPE_COLUMN: &str = "entity_type"; +pub const ENTITY_ID_COLUMN: &str = "entity_id"; +pub const ENTITY_ID_ATTRS_COLUMN: &str = "entity_id_attrs"; +pub const ENTITY_SCOPE_COLUMN: &str = "scope"; +pub const ENTITY_DESCRIPTIVE_COLUMN: &str = "descriptive"; +pub const SOURCE_TABLES_COLUMN: &str = "source_tables"; +pub const SRC_TYPE_COLUMN: &str = "src_type"; +pub const SRC_ID_COLUMN: &str = "src_id"; +pub const DST_TYPE_COLUMN: &str = "dst_type"; +pub const DST_ID_COLUMN: &str = "dst_id"; +pub const REL_TYPE_COLUMN: &str = "rel_type"; +pub const PROVENANCE_COLUMN: &str = "provenance"; +pub const CONFIDENCE_COLUMN: &str = "confidence"; +pub const REQUEST_COUNT_COLUMN: &str = "request_count"; +pub const ERROR_COUNT_COLUMN: &str = "error_count"; +pub const DURATION_SUM_COLUMN: &str = "duration_sum"; +pub const DURATION_COUNT_COLUMN: &str = "duration_count"; +pub const EDGE_ATTRIBUTES_COLUMN: &str = "attributes"; +// Declared-edge table only. +pub const VALID_FROM_COLUMN: &str = "valid_from"; +pub const VALID_UNTIL_COLUMN: &str = "valid_until"; +pub const GENERATION_ID_COLUMN: &str = "generation_id"; // ---- End of entity relationship graph tables ---- diff --git a/src/frontend/src/instance/entity_graph.rs b/src/frontend/src/instance/entity_graph.rs index 70c57846fd0..14f18ed1b79 100644 --- a/src/frontend/src/instance/entity_graph.rs +++ b/src/frontend/src/instance/entity_graph.rs @@ -34,8 +34,8 @@ use auth::{ use catalog::CatalogManager; use catalog::system_schema::semantic_graph::EntityGraphProvider; use common_catalog::consts::{ - DEFAULT_PRIVATE_SCHEMA_NAME, DEFAULT_SCHEMA_NAME, INFORMATION_SCHEMA_NAME, PG_CATALOG_NAME, - SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SERVICE_NAME_COLUMN, + DEFAULT_PRIVATE_SCHEMA_NAME, DEFAULT_SCHEMA_NAME, INFORMATION_SCHEMA_NAME, OBSERVED_AT_COLUMN, + PG_CATALOG_NAME, SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, }; use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; @@ -47,9 +47,10 @@ use datafusion::dataframe::DataFrame; use datafusion_expr::LogicalPlan; use futures::TryStreamExt; use operator::statement::semantic_graph::{ - CallsSource, CoDeclaredSource, DeclaredSource, EntityDeclaration, GraphQueryWindow, - OBSERVED_AT_COLUMN, RegistrySource, RelationshipSources, build_registry_plan, - build_relationships_plan, declared_relationships_schema_matches, + CallsSource, CoDeclaredSource, Conventions, DeclaredSource, ENTITY_TYPE_GEN_AI_AGENT, + ENTITY_TYPE_SERVICE, EntityDeclaration, GraphQueryWindow, ImplicitEntity, RegistrySource, + RelationshipSources, build_registry_plan, build_relationships_plan, conventions, + declared_relationships_schema_matches, }; use query::QueryEngineRef; use session::context::{QueryContext, QueryContextBuilder, QueryContextRef}; @@ -59,7 +60,8 @@ use table::TableRef; use table::metadata::TableInfo; use table::predicate::{TimeRangeExtraction, extract_time_range_strict}; use table::requests::{ - EntityRole, is_trace_v1_table, parse_entity_columns, parse_entity_option_key, + EntityRole, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_METRIC, SOURCE_PROMETHEUS, + is_trace_v1_table, parse_entity_columns, parse_entity_option_key, }; use crate::error; @@ -183,42 +185,127 @@ impl EntityGraphProviderImpl { .collect() } - /// All entity declarations of one table. Trace-v1 tables created before the - /// ingest-side auto-stamp carry no `entity.service.id` option; synthesize it - /// (their schema is fixed), so their `calls` edges have endpoint entities. - /// A table with an explicit service declaration never gets the synthesized - /// one — not even when the explicit declaration is invalid and skipped: - /// silently falling back would change entity identity behind the user's back. - fn declarations_for(table_info: &TableInfo) -> Vec { + /// All entity declarations of one table: the explicit options plus the + /// zero-configuration conventions (`otlp_trace_entities` for trace-v1 + /// tables — including the `service` identity of tables created before the + /// ingest-side auto-stamp — and the Prometheus descriptor whitelist). An + /// explicit declaration of a type always suppresses the implicit one, even + /// when the explicit declaration is invalid and skipped: silently falling + /// back would change entity identity behind the user's back. + fn declarations_for( + table_info: &TableInfo, + conventions: &Conventions, + ) -> Vec { let mut declarations = Self::parse_declarations(table_info); - let has_explicit_service = table_info.meta.options.extra_options.keys().any(|key| { + if is_trace_v1_table(table_info) { + Self::extend_with_implicit_entities( + table_info, + &conventions.otlp_trace_entities, + &mut declarations, + ); + } + Self::extend_with_prometheus_conventions(table_info, conventions, &mut declarations); + declarations + } + + /// Whether the table carries an explicit `entity..id` option. + fn explicitly_declares(table_info: &TableInfo, entity_type: &str) -> bool { + table_info.meta.options.extra_options.keys().any(|key| { parse_entity_option_key(key) - .is_some_and(|(ty, role)| ty == "service" && role == EntityRole::Id) - }); - if is_trace_v1_table(table_info) - && !has_explicit_service - && table_info - .meta - .schema - .column_schema_by_name(SERVICE_NAME_COLUMN) - .is_some() - && let Some(time_index) = table_info - .meta - .schema - .timestamp_column() - .map(|c| c.name.clone()) + .is_some_and(|(ty, role)| ty == entity_type && role == EntityRole::Id) + }) + } + + /// Implicit declarations of the well-known Prometheus entity-descriptor + /// metrics (the `prometheus_info_metrics` whitelist of `conventions.yaml`), + /// gated on the ingest-stamped `signal_type=metric` + `source=prometheus` + /// options and keyed by table name. The metric engine's physical table + /// aggregates every logical table's columns and must not contribute a + /// duplicate source. + fn extend_with_prometheus_conventions( + table_info: &TableInfo, + conventions: &Conventions, + declarations: &mut Vec, + ) { + let Some(implicit_entities) = conventions.prometheus_info_metrics.get(&table_info.name) + else { + return; + }; + let options = &table_info.meta.options.extra_options; + if options.get(SEMANTIC_SIGNAL_TYPE).map(String::as_str) != Some(SIGNAL_TYPE_METRIC) + || options.get(SEMANTIC_SOURCE).map(String::as_str) != Some(SOURCE_PROMETHEUS) + || table_info.is_physical_table() { + debug!( + "Table `{}` matches the info-metric whitelist but is not a prometheus logical \ + metric table; skipping its implicit declarations", + table_info.name + ); + return; + } + Self::extend_with_implicit_entities(table_info, implicit_entities, declarations); + } + + /// Synthesizes the applicable subset of `entities` on `table_info`: + /// explicit declarations win, every id column must exist (no guessing), + /// descriptive columns are filtered to those present. + fn extend_with_implicit_entities( + table_info: &TableInfo, + entities: &[ImplicitEntity], + declarations: &mut Vec, + ) { + let schema = &table_info.meta.schema; + let Some(time_index) = schema.timestamp_column().map(|c| c.name.clone()) else { + debug!( + "Table `{}` has no time index; skipping its implicit declarations", + table_info.name + ); + return; + }; + for implicit in entities { + if Self::explicitly_declares(table_info, &implicit.entity) { + debug!( + "Table `{}` explicitly declares `{}`; the implicit declaration is suppressed", + table_info.name, implicit.entity + ); + continue; + } + if let Some(missing) = implicit + .id + .iter() + .find(|c| schema.column_schema_by_name(c).is_none()) + { + debug!( + "Table `{}` lacks the id column `{}`; skipping the implicit `{}` declaration", + table_info.name, missing, implicit.entity + ); + continue; + } + let descriptive_columns = if implicit.descriptive_rest { + table_info + .meta + .row_key_column_names() + .filter(|c| !implicit.id.contains(c)) + .cloned() + .collect() + } else { + implicit + .descriptive + .iter() + .filter(|c| schema.column_schema_by_name(c).is_some()) + .cloned() + .collect() + }; declarations.push(EntityDeclaration { schema: table_info.schema_name.clone(), table: table_info.name.clone(), - time_index, - entity_type: "service".to_string(), - id_columns: vec![SERVICE_NAME_COLUMN.to_string()], - descriptive_columns: vec![], + time_index: time_index.clone(), + entity_type: implicit.entity.clone(), + id_columns: implicit.id.clone(), + descriptive_columns, scope_columns: vec![], }); } - declarations } /// Enumerates entity declarations and trace tables across a catalog. Trace @@ -233,6 +320,10 @@ impl EntityGraphProviderImpl { let Some(catalog_manager) = self.catalog_manager.upgrade() else { return Ok((vec![], vec![])); }; + let conventions = conventions() + .map_err(datafusion::error::DataFusionError::Internal) + .context(error::DataFusionSnafu) + .map_err(BoxedError::new)?; // A target-blind checker (e.g. the default mode-based one) answers the // same for every table: ask once up front instead of per table. @@ -273,7 +364,7 @@ impl EntityGraphProviderImpl { let mut tables = catalog_manager.tables(catalog, &schema, query_ctx); while let Some(table) = tables.try_next().await.map_err(BoxedError::new)? { let table_info = table.table_info(); - let table_declarations = Self::declarations_for(&table_info); + let table_declarations = Self::declarations_for(&table_info, conventions); let is_trace = is_trace_v1_table(&table_info); // Authorize only tables that would contribute rows. if per_table_auth @@ -301,8 +392,8 @@ impl EntityGraphProviderImpl { .find(|d| d.entity_type == entity_type) .cloned() }; - let service = find("service"); - let agent = find("agent"); + let service = find(ENTITY_TYPE_SERVICE); + let agent = find(ENTITY_TYPE_GEN_AI_AGENT); if service.is_none() { // No usable service identity: the table cannot // contribute service-calls edges (see declarations_for). @@ -541,6 +632,7 @@ mod tests { use common_catalog::consts::{DEFAULT_CATALOG_NAME, MITO_ENGINE}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, SchemaBuilder}; + use store_api::metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY; use table::metadata::{TableInfoBuilder, TableMeta, TableType}; use table::requests::{TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1, TableOptions}; @@ -564,6 +656,15 @@ mod tests { } fn assemble_table_info(column_schemas: Vec, extra: &[(&str, &str)]) -> TableInfo { + named_table_info("t1", column_schemas, vec![], extra) + } + + fn named_table_info( + name: &str, + column_schemas: Vec, + primary_key_indices: Vec, + extra: &[(&str, &str)], + ) -> TableInfo { let schema = Arc::new( SchemaBuilder::try_from_columns(column_schemas) .unwrap() @@ -579,7 +680,7 @@ mod tests { }; let meta = TableMeta { schema, - primary_key_indices: vec![], + primary_key_indices, value_indices: vec![], engine: MITO_ENGINE.to_string(), next_column_id: 1, @@ -591,7 +692,7 @@ mod tests { }; TableInfoBuilder::default() .table_id(1) - .name("t1") + .name(name) .catalog_name(DEFAULT_CATALOG_NAME) .schema_name(DEFAULT_SCHEMA_NAME) .table_version(0) @@ -601,6 +702,24 @@ mod tests { .unwrap() } + /// A metric-engine-logical-table shape: every label column is a tag. + fn prom_table_info(name: &str, tags: &[&str], extra: &[(&str, &str)]) -> TableInfo { + let mut column_schemas = vec![ + ColumnSchema::new( + "greptime_timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ]; + column_schemas.extend( + tags.iter() + .map(|c| ColumnSchema::new(*c, ConcreteDataType::string_datatype(), true)), + ); + let primary_key_indices = (1..=tags.len()).collect(); + named_table_info(name, column_schemas, primary_key_indices, extra) + } + #[test] fn future_lower_only_window_matches_nothing() { let future_ms = common_time::util::current_time_millis() + 86_400_000; @@ -634,7 +753,7 @@ mod tests { ("greptime.semantic.entity.host.id", "gone"), ], ); - let declarations = EntityGraphProviderImpl::declarations_for(&info); + let declarations = EntityGraphProviderImpl::declarations_for(&info, conventions().unwrap()); assert_eq!(declarations.len(), 1); assert_eq!(declarations[0].entity_type, "service"); @@ -645,45 +764,161 @@ mod tests { ("greptime.semantic.entity.service.descriptive", "gone"), ], ); - assert!(EntityGraphProviderImpl::declarations_for(&info).is_empty()); + assert!( + EntityGraphProviderImpl::declarations_for(&info, conventions().unwrap()).is_empty() + ); + } + + const PROM_STAMPS: &[(&str, &str)] = &[ + (SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC), + (SEMANTIC_SOURCE, SOURCE_PROMETHEUS), + ]; + + fn sorted_declarations(info: &TableInfo) -> Vec { + let mut declarations = + EntityGraphProviderImpl::declarations_for(info, conventions().unwrap()); + declarations.sort_by(|a, b| a.entity_type.cmp(&b.entity_type)); + declarations } #[test] - fn trace_v1_table_gets_implicit_service_declaration() { - let info = table_info( - &["service_name"], + fn prometheus_info_metric_gets_implicit_declarations() { + let info = prom_table_info( + "kube_pod_info", + &["namespace", "pod", "uid", "node", "job", "instance"], + PROM_STAMPS, + ); + let declarations = sorted_declarations(&info); + assert_eq!(declarations.len(), 2); + assert_eq!(declarations[0].entity_type, "k8s.node"); + assert_eq!(declarations[0].id_columns, vec!["node"]); + assert_eq!(declarations[1].entity_type, "k8s.pod"); + assert_eq!(declarations[1].id_columns, vec!["uid"]); + // host_ip/pod_ip/created_by_* are absent from this table; descriptive + // shrinks to the present columns. + assert_eq!( + declarations[1].descriptive_columns, + vec!["namespace", "pod", "node"] + ); + assert_eq!(declarations[1].time_index, "greptime_timestamp"); + } + + #[test] + fn trace_table_gets_implicit_resource_entities() { + let full = table_info( + &[ + "service_name", + "resource_attributes.service.instance.id", + "resource_attributes.k8s.pod.uid", + "resource_attributes.k8s.pod.name", + ], &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)], ); - let declarations = EntityGraphProviderImpl::declarations_for(&info); - assert_eq!(declarations.len(), 1); - let decl = &declarations[0]; - assert_eq!(decl.entity_type, "service"); - assert_eq!(decl.id_columns, vec![SERVICE_NAME_COLUMN.to_string()]); - assert_eq!(decl.time_index, "ts"); - - // An explicit service declaration wins; no duplicate is synthesized. - let info = table_info( - &["service_name"], - &[ - (TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1), - ("greptime.semantic.entity.service.id", "service_name"), - ], + let declarations = sorted_declarations(&full); + let types: Vec<&str> = declarations + .iter() + .map(|d| d.entity_type.as_str()) + .collect(); + assert_eq!(types, vec!["k8s.pod", "service", "service.instance"]); + assert_eq!( + declarations[0].id_columns, + vec!["resource_attributes.k8s.pod.uid"] + ); + assert_eq!( + declarations[0].descriptive_columns, + vec!["resource_attributes.k8s.pod.name"] + ); + assert_eq!( + declarations[2].id_columns, + vec!["service_name", "resource_attributes.service.instance.id"] ); - assert_eq!(EntityGraphProviderImpl::declarations_for(&info).len(), 1); - // A trace-model table without the fixed column synthesizes nothing. - let info = table_info(&[], &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)]); - assert!(EntityGraphProviderImpl::declarations_for(&info).is_empty()); + // Missing uid column: no pod entity synthesized, no name-based guess. + let no_uid = table_info( + &["service_name", "resource_attributes.k8s.pod.name"], + &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)], + ); + let types: Vec = sorted_declarations(&no_uid) + .into_iter() + .map(|d| d.entity_type) + .collect(); + assert_eq!(types, vec!["service"]); - // An explicit but invalid declaration must not silently fall back to - // the synthesized one: identity must not change behind the user's back. - let info = table_info( + // An explicit declaration suppresses the implicit one even when it is + // invalid and skipped: identity must not change behind the user's back. + let invalid_explicit = table_info( &["service_name"], &[ (TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1), ("greptime.semantic.entity.service.id", "gone"), ], ); - assert!(EntityGraphProviderImpl::declarations_for(&info).is_empty()); + assert!(sorted_declarations(&invalid_explicit).is_empty()); + } + + #[test] + fn prometheus_implicit_declarations_are_gated() { + let labels: &[&str] = &["namespace", "pod", "node"]; + // A non-Prometheus source. + assert!( + sorted_declarations(&prom_table_info( + "kube_pod_info", + labels, + &[ + (SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC), + (SEMANTIC_SOURCE, "opentelemetry"), + ], + )) + .is_empty() + ); + // Not a whitelisted metric name. + assert!( + sorted_declarations(&prom_table_info("http_requests_total", labels, PROM_STAMPS)) + .is_empty() + ); + let mut stamps = PROM_STAMPS.to_vec(); + stamps.push((PHYSICAL_TABLE_METADATA_KEY, "true")); + assert!(sorted_declarations(&prom_table_info("kube_pod_info", labels, &stamps)).is_empty()); + + // A missing id column (uid) drops that entity, not the whole table. + let info = prom_table_info("kube_pod_info", &["namespace", "pod", "node"], PROM_STAMPS); + let declarations = sorted_declarations(&info); + assert_eq!(declarations.len(), 1); + assert_eq!(declarations[0].entity_type, "k8s.node"); + } + + #[test] + fn explicit_declaration_suppresses_the_implicit_one() { + let mut stamps = PROM_STAMPS.to_vec(); + stamps.push(("greptime.semantic.entity.k8s.pod.id", "pod")); + let info = prom_table_info("kube_pod_info", &["namespace", "pod", "node"], &stamps); + let declarations = sorted_declarations(&info); + assert_eq!(declarations.len(), 2); + assert_eq!(declarations[0].entity_type, "k8s.node"); + assert_eq!(declarations[1].entity_type, "k8s.pod"); + // The explicit identity wins over the conventional [namespace, pod]. + assert_eq!(declarations[1].id_columns, vec!["pod"]); + } + + #[test] + fn target_info_descriptive_rest_covers_remaining_tags() { + let info = prom_table_info( + "target_info", + &["job", "instance", "k8s_cluster_name", "service_version"], + PROM_STAMPS, + ); + let declarations = sorted_declarations(&info); + assert_eq!(declarations.len(), 2); + assert_eq!(declarations[0].entity_type, "service"); + assert_eq!(declarations[0].id_columns, vec!["job"]); + assert!(declarations[0].descriptive_columns.is_empty()); + // The remaining labels are the target's resource attributes: they + // describe the instance, not the logical service. + assert_eq!(declarations[1].entity_type, "service.instance"); + assert_eq!(declarations[1].id_columns, vec!["job", "instance"]); + assert_eq!( + declarations[1].descriptive_columns, + vec!["k8s_cluster_name", "service_version"] + ); } } diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index 087ea461f10..c9258df2c0f 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -69,7 +69,9 @@ prost.workspace = true prost-types = { workspace = true, optional = true } query.workspace = true regex.workspace = true +serde.workspace = true serde_json.workspace = true +serde_yaml_ng.workspace = true session.workspace = true snafu.workspace = true sql.workspace = true diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs index f1e94b2c616..0ed3c74bbed 100644 --- a/src/operator/src/error.rs +++ b/src/operator/src/error.rs @@ -235,13 +235,6 @@ pub enum Error { location: Location, }, - #[snafu(display("Invalid entity semantic option, reason: {}", reason))] - InvalidEntitySemanticOption { - reason: String, - #[snafu(implicit)] - location: Location, - }, - #[snafu(display("Table not found: {}", table_name))] TableNotFound { table_name: String }, @@ -957,7 +950,6 @@ impl ErrorExt for Error { | Error::InvalidConfigValue { .. } | Error::InvalidInsertRequest { .. } | Error::InvalidDeleteRequest { .. } - | Error::InvalidEntitySemanticOption { .. } | Error::IllegalPrimaryKeysDef { .. } | Error::SchemaNotFound { .. } | Error::SchemaExists { .. } diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 3f637ebebcf..256b849123e 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -99,8 +99,8 @@ use table::TableRef; use table::dist_table::DistTable; use table::metadata::{self, TableId, TableInfo, TableMeta, TableType}; use table::requests::{ - AlterKind, AlterTableRequest, COMMENT_KEY, DDL_TIMEOUT, DDL_WAIT, EntityRole, - REPARTITION_COLUMN_HINT_KEY, TableOptions, parse_entity_columns, parse_entity_option_key, + AlterKind, AlterTableRequest, COMMENT_KEY, DDL_TIMEOUT, DDL_WAIT, REPARTITION_COLUMN_HINT_KEY, + TableOptions, parse_entity_columns, parse_entity_option_key, }; use table::table_name::TableName; use table::table_reference::TableReference; @@ -109,12 +109,12 @@ use crate::error::{ self, AlterExprToRequestSnafu, BuildDfLogicalPlanSnafu, CatalogSnafu, ColumnDataTypeSnafu, ColumnNotFoundSnafu, ConvertSchemaSnafu, CreateLogicalTablesSnafu, DeserializePartitionExprSnafu, EmptyDdlExprSnafu, ExternalSnafu, ExtractTableNamesSnafu, - FlowNotFoundSnafu, InvalidEntitySemanticOptionSnafu, InvalidPartitionRuleSnafu, - InvalidPartitionSnafu, InvalidSqlSnafu, InvalidTableNameSnafu, InvalidViewNameSnafu, - InvalidViewStmtSnafu, NotSupportedSnafu, PartitionExprToPbSnafu, Result, SchemaInUseSnafu, - SchemaNotFoundSnafu, SchemaReadOnlySnafu, SerializePartitionExprSnafu, SubstraitCodecSnafu, - TableAlreadyExistsSnafu, TableDdlReservedSnafu, TableMetadataManagerSnafu, TableNotFoundSnafu, - TableReadOnlySnafu, UnrecognizedTableOptionSnafu, ViewAlreadyExistsSnafu, + FlowNotFoundSnafu, InvalidPartitionRuleSnafu, InvalidPartitionSnafu, InvalidSqlSnafu, + InvalidTableNameSnafu, InvalidViewNameSnafu, InvalidViewStmtSnafu, NotSupportedSnafu, + PartitionExprToPbSnafu, Result, SchemaInUseSnafu, SchemaNotFoundSnafu, SchemaReadOnlySnafu, + SerializePartitionExprSnafu, SubstraitCodecSnafu, TableAlreadyExistsSnafu, + TableDdlReservedSnafu, TableMetadataManagerSnafu, TableNotFoundSnafu, TableReadOnlySnafu, + UnrecognizedTableOptionSnafu, ViewAlreadyExistsSnafu, }; use crate::expr_helper::{self, RepartitionRequest, RepartitionSource}; use crate::statement::StatementExecutor; @@ -2371,6 +2371,19 @@ pub fn verify_alter( .context(error::BuildTableMetaSnafu { table_name })?; validate_json2_columns_append_mode(&new_meta.schema, &new_meta.options)?; + // A column referenced by an entity declaration may be dropped (the + // read-time derivation skips the stale declaration), but must not change + // to a type without a stable string form. + for (key, value) in &new_meta.options.extra_options { + if parse_entity_option_key(key).is_none() { + continue; + } + for col in &parse_entity_columns(value) { + if let Some(column) = new_meta.schema.column_schema_by_name(col) { + ensure_entity_column_renders_as_string(key, col, &column.data_type)?; + } + } + } Ok(true) } @@ -2429,11 +2442,7 @@ pub fn create_table_info( &create_table.time_index, )?; - validate_entity_semantic_options( - &table_options, - &column_name_to_index_map, - &primary_key_indices, - )?; + validate_entity_semantic_options(&table_options, &schema)?; let meta = TableMeta { schema, @@ -2581,32 +2590,54 @@ fn ensure_table_definition_writable(schema: &str, table: &str) -> Result<()> { Ok(()) } -/// Validates `greptime.semantic.entity..{id|descriptive|scope}` options -/// against the table schema: every named column must exist, and `id` columns must -/// additionally be tag (primary-key) columns so entity identity stays indexable -/// and joinable. See `docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md`. -fn validate_entity_semantic_options( - table_options: &TableOptions, - column_name_to_index_map: &HashMap, - primary_key_indices: &[usize], +/// Whether `CAST(column AS Utf8)` yields a stable entity-id string — the +/// binary-backed and nested types do not, and without the DDL check the +/// failure would surface only when the graph is scanned. +fn has_stable_string_form(data_type: &datatypes::prelude::ConcreteDataType) -> bool { + use datatypes::prelude::ConcreteDataType; + !matches!( + data_type, + ConcreteDataType::Binary(_) + | ConcreteDataType::Json(_) + | ConcreteDataType::Vector(_) + | ConcreteDataType::List(_) + | ConcreteDataType::Struct(_) + | ConcreteDataType::Dictionary(_) + | ConcreteDataType::Null(_) + ) +} + +fn ensure_entity_column_renders_as_string( + key: &str, + col: &str, + data_type: &datatypes::prelude::ConcreteDataType, ) -> Result<()> { + ensure!( + has_stable_string_form(data_type), + InvalidSqlSnafu { + err_msg: format!( + "entity column `{col}` (option `{key}`) has type `{data_type}`, which cannot \ + render as a string", + ), + } + ); + Ok(()) +} + +/// Validates `greptime.semantic.entity..{id|descriptive|scope}` options +/// against the table schema: every named column must exist (tag or field) and +/// have a stable string form — the derivation renders id, scope and +/// descriptive values as strings. +fn validate_entity_semantic_options(table_options: &TableOptions, schema: &Schema) -> Result<()> { for (key, value) in &table_options.extra_options { - let Some((_, role)) = parse_entity_option_key(key) else { + if parse_entity_option_key(key).is_none() { continue; }; - let is_id_role = role == EntityRole::Id; for col in &parse_entity_columns(value) { - let idx = column_name_to_index_map - .get(col) + let column = schema + .column_schema_by_name(col) .context(ColumnNotFoundSnafu { msg: col })?; - ensure!( - !is_id_role || primary_key_indices.contains(idx), - InvalidEntitySemanticOptionSnafu { - reason: format!( - "entity id column `{col}` (option `{key}`) must be a tag/primary-key column" - ), - } - ); + ensure_entity_column_renders_as_string(key, col, &column.data_type)?; } } Ok(()) @@ -3142,13 +3173,12 @@ mod test { #[test] fn test_validate_entity_semantic_options() { - let col_map = HashMap::from([ - ("service_name".to_string(), 0usize), - ("host_id".to_string(), 1usize), - ("value".to_string(), 2usize), + let schema = Schema::new(vec![ + ColumnSchema::new("service_name", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("host_id", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true), + ColumnSchema::new("payload", ConcreteDataType::binary_datatype(), true), ]); - // service_name (0) and host_id (1) are tags; value (2) is a field. - let pk = vec![0usize, 1usize]; let opts = |pairs: &[(&str, &str)]| TableOptions { extra_options: pairs .iter() @@ -3157,39 +3187,30 @@ mod test { ..Default::default() }; - // id on a tag column, composite id on tags, and descriptive on a field - // column are all accepted. + // Any existing column with a string form may be an id, tag or field. for pairs in [ - &[("greptime.semantic.entity.service.id", "service_name")][..], &[( "greptime.semantic.entity.process.id", "service_name,host_id", )][..], - &[("greptime.semantic.entity.service.descriptive", "value")][..], + &[("greptime.semantic.entity.service.id", "value")][..], ] { - assert!(validate_entity_semantic_options(&opts(pairs), &col_map, &pk).is_ok()); + assert!(validate_entity_semantic_options(&opts(pairs), &schema).is_ok()); } - // id pointing at a field column is rejected. - let err = validate_entity_semantic_options( - &opts(&[("greptime.semantic.entity.service.id", "value")]), - &col_map, - &pk, - ) - .unwrap_err(); - assert!(matches!( - err, - error::Error::InvalidEntitySemanticOption { .. } - )); - - // id pointing at a missing column is rejected. - let err = validate_entity_semantic_options( + let missing = validate_entity_semantic_options( &opts(&[("greptime.semantic.entity.service.id", "nope")]), - &col_map, - &pk, + &schema, ) .unwrap_err(); - assert!(matches!(err, error::Error::ColumnNotFound { .. })); + assert!(matches!(missing, error::Error::ColumnNotFound { .. })); + + let binary_id = validate_entity_semantic_options( + &opts(&[("greptime.semantic.entity.service.id", "payload")]), + &schema, + ) + .unwrap_err(); + assert!(matches!(binary_id, error::Error::InvalidSql { .. })); } #[test] @@ -3555,6 +3576,43 @@ WITH ('repartition.column.hint' = ' host ')", ); } + #[test] + fn test_verify_alter_guards_entity_column_types() { + let expr = create_expr_from_sql( + "CREATE TABLE t (svc STRING, ts TIMESTAMP TIME INDEX) \ + WITH ('greptime.semantic.entity.service.id'='svc');", + ); + let table_info = Arc::new(create_table_info(&expr, vec![]).unwrap()); + let alter = |kind| AlterTableExpr { + catalog_name: "greptime".to_string(), + schema_name: "public".to_string(), + table_name: "t".to_string(), + kind: Some(kind), + }; + + let modify = alter(Kind::ModifyColumnTypes(api::v1::ModifyColumnTypes { + modify_column_types: vec![api::v1::ModifyColumnType { + column_name: "svc".to_string(), + target_type: api::v1::ColumnDataType::Binary as i32, + target_type_extension: None, + }], + })); + let err = verify_alter(1, table_info.clone(), modify).unwrap_err(); + assert!( + err.to_string().contains("cannot render as a string"), + "{err}" + ); + + // Dropping a declared column stays allowed: the derivation skips the + // stale declaration. + let drop = alter(Kind::DropColumns(api::v1::DropColumns { + drop_columns: vec![api::v1::DropColumn { + name: "svc".to_string(), + }], + })); + assert!(verify_alter(1, table_info, drop).unwrap()); + } + #[test] fn test_json2_requires_append_mode() { let cases = [ diff --git a/src/operator/src/statement/semantic_graph.rs b/src/operator/src/statement/semantic_graph.rs index c14d9d0d04e..777c6f677ea 100644 --- a/src/operator/src/statement/semantic_graph.rs +++ b/src/operator/src/statement/semantic_graph.rs @@ -25,6 +25,7 @@ //! including filter pushdown into the source table scans. See //! `docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md`. +mod conventions; mod relationships; use std::sync::{Arc, LazyLock}; @@ -35,10 +36,25 @@ use api::v1::{ SemanticType, }; use common_catalog::consts::{ - DEFAULT_PRIVATE_SCHEMA_NAME, SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, + CONFIDENCE_COLUMN, DEFAULT_CATALOG_NAME, DEFAULT_PRIVATE_SCHEMA_NAME, DST_ID_COLUMN, + DST_TYPE_COLUMN, DURATION_COUNT_COLUMN, DURATION_SUM_COLUMN, EDGE_ATTRIBUTES_COLUMN, + ENTITY_DESCRIPTIVE_COLUMN, ENTITY_ID_ATTRS_COLUMN, ENTITY_ID_COLUMN, ENTITY_SCOPE_COLUMN, + ENTITY_TYPE_COLUMN, ERROR_COUNT_COLUMN, FRESH_UNTIL_COLUMN, GENERATION_ID_COLUMN, + OBSERVED_AT_COLUMN, PROVENANCE_COLUMN, REL_TYPE_COLUMN, REQUEST_COUNT_COLUMN, + SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SOURCE_TABLES_COLUMN, SRC_ID_COLUMN, + SRC_TYPE_COLUMN, VALID_FROM_COLUMN, VALID_UNTIL_COLUMN, WINDOW_END_COLUMN, WINDOW_START_COLUMN, }; use common_function::function::FunctionContext; use common_function::function_registry::FUNCTION_REGISTRY; +pub use conventions::{ + Conventions, ENTITY_TYPE_GEN_AI_AGENT, ENTITY_TYPE_GEN_AI_MODEL, ENTITY_TYPE_GEN_AI_TOOL, + ENTITY_TYPE_HOST, ENTITY_TYPE_K8S_CONTAINER, ENTITY_TYPE_K8S_NODE, ENTITY_TYPE_K8S_POD, + ENTITY_TYPE_K8S_WORKLOAD, ENTITY_TYPE_PROCESS, ENTITY_TYPE_SERVICE, + ENTITY_TYPE_SERVICE_INSTANCE, ImplicitEntity, PROVENANCE_AGENT, PROVENANCE_ATTRIBUTE, + PROVENANCE_DECLARED, PROVENANCE_TRACE, REL_TYPE_CALLS, REL_TYPE_CONTAINS, REL_TYPE_DEPENDS_ON, + REL_TYPE_INVOKES, REL_TYPE_OWNS, REL_TYPE_PART_OF, REL_TYPE_RUNS_ON, REL_TYPE_USES, + conventions, +}; use datafusion::arrow::datatypes::{DataType, TimeUnit}; use datafusion::dataframe::DataFrame; use datafusion::functions::{core as core_fns, datetime as datetime_fns, string as string_fns}; @@ -86,7 +102,7 @@ pub fn declared_relationships_schema_matches(table_info: &table::metadata::Table return false; } - let canonical = build_declared_relationships_expr("greptime"); + let canonical = build_declared_relationships_expr(DEFAULT_CATALOG_NAME); if schema.column_schemas().len() != canonical.column_defs.len() { return false; } @@ -110,9 +126,6 @@ pub fn declared_relationships_schema_matches(table_info: &table::metadata::Table /// service-graph convention. const BIN_NANOS: i64 = 60 * 1_000_000_000; -/// Time index of the graph tables: when an observation was recorded. -pub const OBSERVED_AT_COLUMN: &str = "observed_at"; - /// Default retention for the declared-edge table; expiry slides the topology window. const DEFAULT_DECLARED_RELATIONSHIPS_TTL: &str = "90d"; /// Environment variable overriding the declared-edge table's TTL at creation @@ -133,14 +146,14 @@ fn declared_relationships_ttl() -> String { /// `provenance` and `generation_id` are in the key so a declared edge and a /// (future) derived edge for the same pair coexist without clobbering. pub const DECLARED_PRIMARY_KEY_COLUMNS: [&str; 8] = [ - "src_type", - "src_id", - "rel_type", - "dst_type", - "dst_id", - "provenance", - "scope", - "generation_id", + SRC_TYPE_COLUMN, + SRC_ID_COLUMN, + REL_TYPE_COLUMN, + DST_TYPE_COLUMN, + DST_ID_COLUMN, + PROVENANCE_COLUMN, + ENTITY_SCOPE_COLUMN, + GENERATION_ID_COLUMN, ]; /// The externally visible edge identity: the primary key minus `scope` and @@ -148,12 +161,12 @@ pub const DECLARED_PRIMARY_KEY_COLUMNS: [&str; 8] = [ /// uses this identity, or assertions differing only in those two columns would /// surface as indistinguishable duplicate rows. const DECLARED_EDGE_IDENTITY_COLUMNS: [&str; 6] = [ - "src_type", - "src_id", - "rel_type", - "dst_type", - "dst_id", - "provenance", + SRC_TYPE_COLUMN, + SRC_ID_COLUMN, + REL_TYPE_COLUMN, + DST_TYPE_COLUMN, + DST_ID_COLUMN, + PROVENANCE_COLUMN, ]; fn column( @@ -210,32 +223,32 @@ pub fn build_declared_relationships_expr(catalog: &str) -> CreateTableExpr { SemanticType::Timestamp, false, ), - field("window_start", ColumnDataType::TimestampMillisecond), - field("window_end", ColumnDataType::TimestampMillisecond), - field("fresh_until", ColumnDataType::TimestampMillisecond), + field(WINDOW_START_COLUMN, ColumnDataType::TimestampMillisecond), + field(WINDOW_END_COLUMN, ColumnDataType::TimestampMillisecond), + field(FRESH_UNTIL_COLUMN, ColumnDataType::TimestampMillisecond), // Declared-only business validity. NULL valid_from = valid since the // declaration; NULL valid_until = valid for as long as the row exists // (TTL expiry retires the edge with the row). - field("valid_from", ColumnDataType::TimestampMillisecond), - field("valid_until", ColumnDataType::TimestampMillisecond), + field(VALID_FROM_COLUMN, ColumnDataType::TimestampMillisecond), + field(VALID_UNTIL_COLUMN, ColumnDataType::TimestampMillisecond), // Endpoints + edge identity (all tags, in primary-key order). - tag("src_type"), - tag("src_id"), - tag("rel_type"), - tag("dst_type"), - tag("dst_id"), - tag("provenance"), - tag("scope"), - tag("generation_id"), + tag(SRC_TYPE_COLUMN), + tag(SRC_ID_COLUMN), + tag(REL_TYPE_COLUMN), + tag(DST_TYPE_COLUMN), + tag(DST_ID_COLUMN), + tag(PROVENANCE_COLUMN), + tag(ENTITY_SCOPE_COLUMN), + tag(GENERATION_ID_COLUMN), // Confidence + RED metrics (populated for derived edges; usually NULL here). - field("confidence", ColumnDataType::Float64), - field("request_count", ColumnDataType::Int64), - field("error_count", ColumnDataType::Int64), - field("duration_sum", ColumnDataType::Float64), - field("duration_count", ColumnDataType::Int64), + field(CONFIDENCE_COLUMN, ColumnDataType::Float64), + field(REQUEST_COUNT_COLUMN, ColumnDataType::Int64), + field(ERROR_COUNT_COLUMN, ColumnDataType::Int64), + field(DURATION_SUM_COLUMN, ColumnDataType::Float64), + field(DURATION_COUNT_COLUMN, ColumnDataType::Int64), // JSONB, so the union matches the computed table's json column // without a per-scan parse. - json_field("attributes"), + json_field(EDGE_ATTRIBUTES_COLUMN), ]; let table_options = [(TTL_KEY.to_string(), declared_relationships_ttl())] @@ -408,6 +421,16 @@ fn cast_string_or_empty(column: &str) -> Expr { /// a single column, the sorted `k=v,k=v` rendering for a composite. `col` /// constructs the column reference (unqualified for registry branches, /// join-side-qualified for the calls derivation). +/// A row identifies an entity only when every identity component is present +/// and non-empty: kube-state-metrics descriptors emit empty-string labels (an +/// unscheduled pod's `node`, an owner-less pod's `owner_*`), and an empty +/// string is never a meaningful entity id. +pub(crate) fn identifies(column: &str) -> Expr { + ident(column) + .is_not_null() + .and(cast(ident(column), DataType::Utf8).not_eq(lit(""))) +} + fn entity_id_expr(id_columns: &[String], col: &dyn Fn(&str) -> Expr) -> Expr { if let [id] = id_columns { cast(col(id), DataType::Utf8) @@ -501,16 +524,16 @@ fn sorted_kv_expr_with(sorted_cols: &[String], nullable: bool, col: &dyn Fn(&str } const REGISTRY_COLUMNS: [&str; 10] = [ - "observed_at", - "window_start", - "window_end", - "fresh_until", - "entity_type", - "entity_id", - "entity_id_attrs", - "scope", - "descriptive", - "source_tables", + OBSERVED_AT_COLUMN, + WINDOW_START_COLUMN, + WINDOW_END_COLUMN, + FRESH_UNTIL_COLUMN, + ENTITY_TYPE_COLUMN, + ENTITY_ID_COLUMN, + ENTITY_ID_ATTRS_COLUMN, + ENTITY_SCOPE_COLUMN, + ENTITY_DESCRIPTIVE_COLUMN, + SOURCE_TABLES_COLUMN, ]; const REGISTRY_VALID_COLUMN: &str = "__entity_valid"; @@ -568,10 +591,10 @@ fn registry_source( let mut rows = Vec::with_capacity(1 + rest.len()); for decl in std::iter::once(first).chain(rest) { - // CAST even a single-column id: id columns must be tags but not - // necessarily strings, and the computed table declares entity_id - // STRING. Composite ids additionally carry a JSON object of the id - // columns in entity_id_attrs. + // CAST even a single-column id: id columns need not be strings, and + // the computed table declares entity_id STRING. Composite ids + // additionally carry a JSON object of the id columns in + // entity_id_attrs. let entity_id = entity_id_expr(&decl.id_columns, &|c| ident(c)); let entity_id_attrs = if decl.id_columns.len() == 1 { null_json() @@ -603,12 +626,12 @@ fn registry_source( json_quote(&format!("{}.{}", decl.schema, decl.table)) ))); - // Tag columns may still be nullable; a NULL identity component - // identifies nothing. Keep this predicate per declaration so a NULL - // identity for one entity does not remove other entities on the row. - let valid = decl.id_columns.iter().fold(lit(true), |predicate, id| { - predicate.and(ident(id).is_not_null()) - }); + // Keep this predicate per declaration so an absent identity for one + // entity does not remove other entities on the row. + let valid = decl + .id_columns + .iter() + .fold(lit(true), |predicate, id| predicate.and(identifies(id))); rows.push(vec![ valid, @@ -961,8 +984,31 @@ mod tests { } #[tokio::test] - async fn registry_skips_null_identity_rows() { - let ctx = metric_table_ctx(); + async fn registry_skips_absent_identity_rows() { + // NULL identifies nothing, and so does a kube-state-metrics-style + // empty label (an unscheduled pod's `node` arrives as ""). + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("host", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000, 2_000, 3_000])) as ArrayRef, + Arc::new(StringArray::from(vec![Some("h2"), None, Some("")])), + ], + ) + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_table( + "app_latency", + Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), + ) + .unwrap(); let df = ctx.table("app_latency").await.unwrap(); let plan = build_registry_plan( vec![RegistrySource { @@ -974,9 +1020,8 @@ mod tests { .unwrap() .unwrap(); let batches = collect(&ctx, plan).await; - let mut ids: Vec = batches.iter().flat_map(|b| strings(b, 5)).collect(); - ids.sort(); - assert_eq!(ids, vec!["h2", r#"we"ird\host"#]); + let ids: Vec = batches.iter().flat_map(|b| strings(b, 5)).collect(); + assert_eq!(ids, vec!["h2"]); } #[tokio::test] diff --git a/src/operator/src/statement/semantic_graph/conventions.rs b/src/operator/src/statement/semantic_graph/conventions.rs new file mode 100644 index 00000000000..36d058921b5 --- /dev/null +++ b/src/operator/src/statement/semantic_graph/conventions.rs @@ -0,0 +1,291 @@ +// Copyright 2023 Greptime Team +// +// 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. + +//! The built-in derivation conventions of the entity graph, embedded as +//! [`conventions.yaml`](./conventions.yaml): the co-declared edge vocabulary, +//! the virtual-destination candidates, and the implicit declarations of +//! well-known Prometheus entity-descriptor metrics. The file is data shipped +//! with the binary, not an operator-editable configuration surface; explicit +//! `greptime.semantic.entity.*` declarations always override it. + +use std::collections::{BTreeMap, HashSet}; +use std::sync::LazyLock; + +use serde::Deserialize; + +/// A same-row co-declaration rule: a source-table row carrying both entity +/// identities witnesses `src -rel-> dst`. +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct EdgeRule { + pub src: String, + pub dst: String, + pub rel: String, +} + +/// A span-attribute column that may name an uninstrumented peer, and the +/// `connection_type` its match implies. +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct VirtualDstCandidate { + pub column: String, + pub connection_type: String, +} + +/// One implicit entity declaration: of a whitelisted Prometheus info metric, +/// or of a trace-v1 table's flattened resource attributes. +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ImplicitEntity { + pub entity: String, + /// Identifying label columns; every one must exist on the table for the + /// declaration to apply. + pub id: Vec, + /// Descriptive label columns, filtered to those present (kube-state-metrics + /// label sets vary across versions). + #[serde(default)] + pub descriptive: Vec, + /// Snapshot every tag column except the id columns as descriptive — + /// `target_info`-style enrichment over an open attribute set. + #[serde(default)] + pub descriptive_rest: bool, +} + +/// The parsed, validated conventions file. +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Conventions { + pub co_declared_edges: Vec, + pub trace_co_declared_edges: Vec, + pub virtual_dst_candidates: Vec, + pub otlp_trace_entities: Vec, + pub prometheus_info_metrics: BTreeMap>, +} + +/// The built-in entity-type vocabulary. User-declared types are open-ended; +/// the embedded conventions file must stay inside this set. +pub const ENTITY_TYPE_SERVICE: &str = "service"; +pub const ENTITY_TYPE_SERVICE_INSTANCE: &str = "service.instance"; +pub const ENTITY_TYPE_HOST: &str = "host"; +pub const ENTITY_TYPE_PROCESS: &str = "process"; +pub const ENTITY_TYPE_K8S_POD: &str = "k8s.pod"; +pub const ENTITY_TYPE_K8S_NODE: &str = "k8s.node"; +pub const ENTITY_TYPE_K8S_CONTAINER: &str = "k8s.container"; +pub const ENTITY_TYPE_K8S_WORKLOAD: &str = "k8s.workload"; +pub const ENTITY_TYPE_K8S_SERVICE: &str = "k8s.service"; +pub const ENTITY_TYPE_GEN_AI_AGENT: &str = "gen_ai.agent"; +pub const ENTITY_TYPE_GEN_AI_MODEL: &str = "gen_ai.model"; +pub const ENTITY_TYPE_GEN_AI_TOOL: &str = "gen_ai.tool"; + +const ENTITY_TYPES: [&str; 12] = [ + ENTITY_TYPE_SERVICE, + ENTITY_TYPE_SERVICE_INSTANCE, + ENTITY_TYPE_HOST, + ENTITY_TYPE_PROCESS, + ENTITY_TYPE_K8S_POD, + ENTITY_TYPE_K8S_NODE, + ENTITY_TYPE_K8S_CONTAINER, + ENTITY_TYPE_K8S_WORKLOAD, + ENTITY_TYPE_K8S_SERVICE, + ENTITY_TYPE_GEN_AI_AGENT, + ENTITY_TYPE_GEN_AI_MODEL, + ENTITY_TYPE_GEN_AI_TOOL, +]; + +/// The relationship vocabulary (the RFC's rel_type table) and the edge +/// provenances. +pub const REL_TYPE_CALLS: &str = "calls"; +pub const REL_TYPE_RUNS_ON: &str = "runs_on"; +pub const REL_TYPE_CONTAINS: &str = "contains"; +pub const REL_TYPE_PART_OF: &str = "part_of"; +pub const REL_TYPE_USES: &str = "uses"; +pub const REL_TYPE_INVOKES: &str = "invokes"; +pub const REL_TYPE_DEPENDS_ON: &str = "depends_on"; +pub const REL_TYPE_OWNS: &str = "owns"; +pub const PROVENANCE_TRACE: &str = "trace"; +pub const PROVENANCE_ATTRIBUTE: &str = "attribute"; +pub const PROVENANCE_DECLARED: &str = "declared"; +pub const PROVENANCE_AGENT: &str = "agent"; + +const REL_TYPES: [&str; 8] = [ + REL_TYPE_CALLS, + REL_TYPE_RUNS_ON, + REL_TYPE_CONTAINS, + REL_TYPE_PART_OF, + REL_TYPE_USES, + REL_TYPE_INVOKES, + REL_TYPE_DEPENDS_ON, + REL_TYPE_OWNS, +]; + +/// The `connection_type` values a virtual-destination candidate may imply; +/// the calls derivation branches on them when building edge attributes. +pub const CONNECTION_TYPE_DATABASE: &str = "database"; +pub const CONNECTION_TYPE_VIRTUAL_NODE: &str = "virtual_node"; + +const CONNECTION_TYPES: [&str; 2] = [CONNECTION_TYPE_DATABASE, CONNECTION_TYPE_VIRTUAL_NODE]; + +static CONVENTIONS: LazyLock> = + LazyLock::new(|| parse(include_str!("conventions.yaml"))); + +/// The embedded conventions. `Err` means the embedded file is broken — pinned +/// by unit test, and propagated by the derivation paths rather than panicking. +pub fn conventions() -> Result<&'static Conventions, String> { + CONVENTIONS.as_ref().map_err(Clone::clone) +} + +fn parse(yaml: &str) -> Result { + let conventions: Conventions = + serde_yaml_ng::from_str(yaml).map_err(|e| format!("malformed conventions: {e}"))?; + validate(&conventions)?; + Ok(conventions) +} + +fn validate(conventions: &Conventions) -> Result<(), String> { + let mut seen_edges = HashSet::new(); + for rule in conventions + .co_declared_edges + .iter() + .chain(&conventions.trace_co_declared_edges) + { + for ty in [&rule.src, &rule.dst] { + if !ENTITY_TYPES.contains(&ty.as_str()) { + return Err(format!("unknown entity type `{ty}` in edge vocabulary")); + } + } + if !REL_TYPES.contains(&rule.rel.as_str()) { + return Err(format!( + "unknown rel_type `{}` in edge vocabulary", + rule.rel + )); + } + if !seen_edges.insert((&rule.src, &rule.dst, &rule.rel)) { + return Err(format!( + "duplicate edge rule `{} -{}-> {}`", + rule.src, rule.rel, rule.dst + )); + } + } + + let mut seen_columns = HashSet::new(); + for candidate in &conventions.virtual_dst_candidates { + if candidate.column.is_empty() { + return Err("empty virtual destination column".to_string()); + } + if !CONNECTION_TYPES.contains(&candidate.connection_type.as_str()) { + return Err(format!( + "unknown connection_type `{}` for virtual destination `{}`", + candidate.connection_type, candidate.column + )); + } + if !seen_columns.insert(&candidate.column) { + return Err(format!( + "duplicate virtual destination column `{}`", + candidate.column + )); + } + } + + let per_table = conventions + .prometheus_info_metrics + .iter() + .map(|(table, entities)| (table.as_str(), entities)) + .chain(std::iter::once(( + "otlp traces", + &conventions.otlp_trace_entities, + ))); + for (table, entities) in per_table { + let mut seen_types = HashSet::new(); + for implicit in entities { + if !ENTITY_TYPES.contains(&implicit.entity.as_str()) { + return Err(format!( + "unknown entity type `{}` for info metric `{table}`", + implicit.entity + )); + } + if !seen_types.insert(&implicit.entity) { + return Err(format!( + "duplicate entity type `{}` for info metric `{table}`", + implicit.entity + )); + } + if implicit.id.is_empty() || implicit.id.iter().any(String::is_empty) { + return Err(format!( + "entity `{}` of info metric `{table}` needs non-empty id columns", + implicit.entity + )); + } + if implicit.descriptive_rest && !implicit.descriptive.is_empty() { + return Err(format!( + "entity `{}` of info metric `{table}` sets both descriptive and \ + descriptive_rest", + implicit.entity + )); + } + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn embedded_conventions_parse_and_validate() { + conventions().unwrap(); + } + + /// Each case must fail on exactly the rule it names, so the shared + /// boilerplate is valid and the mutated part is inside the vocabulary. + fn broken(edges: &str, info_metrics: &str) -> String { + format!( + "co_declared_edges: [{edges}]\ntrace_co_declared_edges: []\n\ + virtual_dst_candidates: []\notlp_trace_entities: []\n\ + prometheus_info_metrics: {{{info_metrics}}}" + ) + } + + #[test] + fn validation_rejects_broken_conventions() { + let err = |edges, info| parse(&broken(edges, info)).unwrap_err(); + + assert!(err("{src: host, dst: service, rel: pets}", "").contains("unknown rel_type")); + assert!( + err("{src: k8s.pods, dst: k8s.node, rel: runs_on}", "").contains("unknown entity type") + ); + assert!( + err( + "{src: host, dst: service, rel: uses}, {src: host, dst: service, rel: uses}", + "" + ) + .contains("duplicate edge rule") + ); + assert!( + err( + "", + "t: [{entity: host, id: [x], descriptive: [y], descriptive_rest: true}]" + ) + .contains("descriptive_rest") + ); + // Unknown YAML keys are rejected, catching typos in the embedded file. + assert!( + parse(&broken( + "{src: host, dst: service, rel: uses, direction: down}", + "" + )) + .is_err() + ); + } +} diff --git a/src/operator/src/statement/semantic_graph/conventions.yaml b/src/operator/src/statement/semantic_graph/conventions.yaml new file mode 100644 index 00000000000..9fc37fe2a02 --- /dev/null +++ b/src/operator/src/statement/semantic_graph/conventions.yaml @@ -0,0 +1,115 @@ +# The built-in derivation conventions of the entity graph. Embedded and +# versioned with the binary (include_str!), not an operator-editable surface; +# non-standard deployments override per table with explicit +# greptime.semantic.entity.* declarations, which always win. + +co_declared_edges: + - { src: service.instance, dst: host, rel: runs_on } + - { src: process, dst: host, rel: runs_on } + - { src: k8s.pod, dst: k8s.node, rel: runs_on } + - { src: k8s.pod, dst: k8s.container, rel: contains } + - { src: service.instance, dst: service, rel: part_of } + - { src: k8s.pod, dst: k8s.workload, rel: part_of } + - { src: service.instance, dst: k8s.pod, rel: runs_on } + +# Applied only to trace sources; the endpoints are Greptime entity types +# derived from GenAI semantic-convention attributes. +trace_co_declared_edges: + - { src: gen_ai.agent, dst: gen_ai.model, rel: uses } + - { src: gen_ai.agent, dst: gen_ai.tool, rel: invokes } + +# Span-attribute columns naming an uninstrumented peer, in precedence order: +# current OTel names first (service.peer.name / db.namespace replaced the +# deprecated peer.service / db.name in semconv 1.39/1.26, kept for existing +# telemetry), the bare network endpoint last. +virtual_dst_candidates: + - { column: span_attributes.service.peer.name, connection_type: virtual_node } + - { column: span_attributes.peer.service, connection_type: virtual_node } + - { column: span_attributes.db.namespace, connection_type: database } + - { column: span_attributes.db.name, connection_type: database } + - { column: span_attributes.server.address, connection_type: virtual_node } + +# Implicit declarations for greptime_trace_v1 tables; the pod UID bridges +# trace-side pods onto the kube-state-metrics entities below. +otlp_trace_entities: + - entity: service + id: [service_name] + - entity: service.instance + id: [service_name, resource_attributes.service.instance.id] + - entity: k8s.pod + id: [resource_attributes.k8s.pod.uid] + descriptive: + - resource_attributes.k8s.pod.name + - resource_attributes.k8s.namespace.name + +# Implicit declarations for Prometheus entity-descriptor metrics, keyed by +# table name. Pod and service identities are the UIDs (names stay +# descriptive); k8s.workload is a Greptime abstraction with the owner kind +# encoded in its identity. +prometheus_info_metrics: + kube_pod_info: + - entity: k8s.pod + id: [uid] + descriptive: + - namespace + - pod + - node + - host_ip + - pod_ip + - created_by_kind + - created_by_name + - priority_class + - host_network + - entity: k8s.node + id: [node] + kube_node_info: + - entity: k8s.node + id: [node] + descriptive: + - kernel_version + - os_image + - container_runtime_version + - kubelet_version + - pod_cidr + - provider_id + - system_uuid + - internal_ip + kube_pod_owner: + - entity: k8s.pod + id: [uid] + descriptive: [namespace, pod] + - entity: k8s.workload + id: [namespace, owner_kind, owner_name] + descriptive: [owner_is_controller] + kube_pod_container_info: + - entity: k8s.pod + id: [uid] + descriptive: [namespace, pod] + - entity: k8s.container + id: [uid, container] + descriptive: [image, image_id, image_spec, container_id] + kube_pod_init_container_info: + - entity: k8s.pod + id: [uid] + descriptive: [namespace, pod] + - entity: k8s.container + id: [uid, container] + descriptive: [image, image_id, image_spec, container_id, restart_policy] + kube_service_info: + - entity: k8s.service + id: [uid] + descriptive: + - namespace + - service + - cluster_ip + - external_name + - load_balancer_ip + - external_traffic_policy + # The remaining labels are the target's resource attributes: the snapshot + # belongs to the instance, not the logical service. + target_info: + - entity: service + id: [job] + - entity: service.instance + id: [job, instance] + descriptive_rest: true diff --git a/src/operator/src/statement/semantic_graph/relationships.rs b/src/operator/src/statement/semantic_graph/relationships.rs index 416c8dc10e1..3b5c6d37b78 100644 --- a/src/operator/src/statement/semantic_graph/relationships.rs +++ b/src/operator/src/statement/semantic_graph/relationships.rs @@ -19,9 +19,13 @@ //! expression helpers and the entity registry live in the parent module. use common_catalog::consts::{ - DURATION_NANO_COLUMN, PARENT_SPAN_ID_COLUMN, SPAN_ID_COLUMN, SPAN_KIND_CLIENT, - SPAN_KIND_COLUMN, SPAN_KIND_SERVER, SPAN_STATUS_CODE_COLUMN, SPAN_STATUS_ERROR, - TRACE_ID_COLUMN, TRACE_TIMESTAMP_COLUMN, + CONFIDENCE_COLUMN, DST_ID_COLUMN, DST_TYPE_COLUMN, DURATION_COUNT_COLUMN, DURATION_NANO_COLUMN, + DURATION_SUM_COLUMN, EDGE_ATTRIBUTES_COLUMN, ENTITY_SCOPE_COLUMN, ERROR_COUNT_COLUMN, + FRESH_UNTIL_COLUMN, GENERATION_ID_COLUMN, OBSERVED_AT_COLUMN, PARENT_SPAN_ID_COLUMN, + PROVENANCE_COLUMN, REL_TYPE_COLUMN, REQUEST_COUNT_COLUMN, SPAN_ID_COLUMN, SPAN_KIND_CLIENT, + SPAN_KIND_COLUMN, SPAN_KIND_SERVER, SPAN_STATUS_CODE_COLUMN, SPAN_STATUS_ERROR, SRC_ID_COLUMN, + SRC_TYPE_COLUMN, TRACE_ID_COLUMN, TRACE_TIMESTAMP_COLUMN, VALID_FROM_COLUMN, + VALID_UNTIL_COLUMN, WINDOW_END_COLUMN, WINDOW_START_COLUMN, }; use datafusion::arrow::datatypes::DataType; use datafusion::dataframe::DataFrame; @@ -31,12 +35,21 @@ use datafusion::functions_window::expr_fn::row_number; use datafusion_common::{Result as DfResult, ScalarValue}; use datafusion_expr::{Expr, ExprFunctionExt, JoinType, LogicalPlan, cast, ident, lit, when}; +use crate::statement::semantic_graph::conventions::{ + CONNECTION_TYPE_DATABASE, CONNECTION_TYPE_VIRTUAL_NODE, Conventions, ENTITY_TYPE_GEN_AI_AGENT, + ENTITY_TYPE_SERVICE, PROVENANCE_ATTRIBUTE, PROVENANCE_TRACE, REL_TYPE_CALLS, +}; use crate::statement::semantic_graph::{ - DECLARED_EDGE_IDENTITY_COLUMNS, EntityDeclaration, GraphQueryWindow, OBSERVED_AT_COLUMN, - bin_interval, bin_ms, entity_id_expr, interval, null_json, parse_json_expr, qcol, union_all, + DECLARED_EDGE_IDENTITY_COLUMNS, EntityDeclaration, GraphQueryWindow, bin_interval, bin_ms, + conventions, entity_id_expr, identifies, interval, null_json, parse_json_expr, qcol, union_all, unnest_rows, }; +/// The embedded conventions, with a broken file surfaced as a plan error. +fn builtin() -> DfResult<&'static Conventions> { + conventions().map_err(datafusion_common::DataFusionError::Internal) +} + /// The projected columns of `semantic_relationships`, in order. Every derived /// branch and the declared-edge branch must project exactly these so the /// top-level `UNION ALL` type-aligns; `build_relationships_plan` re-selects @@ -44,22 +57,22 @@ use crate::statement::semantic_graph::{ /// additionally stores `valid_from`/`valid_until`, which feed the validity /// filter and the projected window columns.) const RELATIONSHIP_COLUMNS: [&str; 16] = [ - "observed_at", - "window_start", - "window_end", - "fresh_until", - "src_type", - "src_id", - "dst_type", - "dst_id", - "rel_type", - "provenance", - "confidence", - "request_count", - "error_count", - "duration_sum", - "duration_count", - "attributes", + OBSERVED_AT_COLUMN, + WINDOW_START_COLUMN, + WINDOW_END_COLUMN, + FRESH_UNTIL_COLUMN, + SRC_TYPE_COLUMN, + SRC_ID_COLUMN, + DST_TYPE_COLUMN, + DST_ID_COLUMN, + REL_TYPE_COLUMN, + PROVENANCE_COLUMN, + CONFIDENCE_COLUMN, + REQUEST_COUNT_COLUMN, + ERROR_COUNT_COLUMN, + DURATION_SUM_COLUMN, + DURATION_COUNT_COLUMN, + EDGE_ATTRIBUTES_COLUMN, ]; /// A child server span starts no earlier than 5 minutes before its client span @@ -129,25 +142,6 @@ pub fn build_relationships_plan( .transpose() } -/// The same-row co-declaration vocabulary open to any declaring table: a row -/// carrying both identities witnesses the edge, and the built-in direction is -/// `src -> dst`. Rows are `(src_type, dst_type, rel_type)`. -const ATTRIBUTE_EDGE_VOCABULARY: [(&str, &str, &str); 6] = [ - ("service.instance", "host", "runs_on"), - ("process", "host", "runs_on"), - ("k8s.pod", "k8s.node", "runs_on"), - ("k8s.pod", "k8s.container", "contains"), - ("service.instance", "service", "part_of"), - ("k8s.pod", "k8s.workload", "part_of"), -]; - -/// The agent-edge vocabulary, applied only to trace sources: the RFC derives -/// `agent uses model` / `agent invokes tool` from span structure (an LLM- or -/// tool-call span carries both identities), so a non-trace table co-declaring -/// these types must not fabricate invocations. -const AGENT_EDGE_VOCABULARY: [(&str, &str, &str); 2] = - [("agent", "model", "uses"), ("agent", "tool", "invokes")]; - const CO_DECLARED_VALID_COLUMN: &str = "__edge_valid"; /// The same-row co-declared branch: for each source table, every vocabulary @@ -159,19 +153,22 @@ fn co_declared_branch( sources: Vec, window: &GraphQueryWindow, ) -> DfResult> { + let vocabulary = builtin()?; let mut union_df: Option = None; for source in sources { let find = |ty: &str| source.declarations.iter().find(|d| d.entity_type == ty); let mut pairs = Vec::new(); - for (src_type, dst_type, rel_type) in ATTRIBUTE_EDGE_VOCABULARY { - if let (Some(src), Some(dst)) = (find(src_type), find(dst_type)) { - pairs.push((src, dst, rel_type, "attribute")); + for rule in &vocabulary.co_declared_edges { + if let (Some(src), Some(dst)) = (find(&rule.src), find(&rule.dst)) { + pairs.push((src, dst, rule.rel.as_str(), PROVENANCE_ATTRIBUTE)); } } if source.is_trace { - for (src_type, dst_type, rel_type) in AGENT_EDGE_VOCABULARY { - if let (Some(src), Some(dst)) = (find(src_type), find(dst_type)) { - pairs.push((src, dst, rel_type, "trace")); + // Agent edges are tied to span structure by the RFC: a non-trace + // table co-declaring these types must not fabricate invocations. + for rule in &vocabulary.trace_co_declared_edges { + if let (Some(src), Some(dst)) = (find(&rule.src), find(&rule.dst)) { + pairs.push((src, dst, rule.rel.as_str(), PROVENANCE_TRACE)); } } } @@ -195,9 +192,7 @@ fn co_declared_branch( .id_columns .iter() .chain(&dst.id_columns) - .fold(lit(true), |predicate, id| { - predicate.and(ident(id).is_not_null()) - }); + .fold(lit(true), |predicate, id| predicate.and(identifies(id))); vec![ valid, bin.clone(), @@ -249,9 +244,9 @@ const DECLARED_REVISION_COLUMN: &str = "__declared_revision"; /// above the computed table, and the physical revision time would fail them. fn declared_edges(scan: DataFrame, window: &GraphQueryWindow) -> DfResult { let eff_valid_from = - core_fns::coalesce().call(vec![ident("valid_from"), ident(OBSERVED_AT_COLUMN)]); + core_fns::coalesce().call(vec![ident(VALID_FROM_COLUMN), ident(OBSERVED_AT_COLUMN)]); let eff_valid_until = - core_fns::coalesce().call(vec![ident("valid_until"), window.observed_end()]); + core_fns::coalesce().call(vec![ident(VALID_UNTIL_COLUMN), window.observed_end()]); // As-of the queried window: a revision recorded after its end, or whose // validity starts after it, was not the edge's state inside the window and @@ -271,17 +266,17 @@ fn declared_edges(scan: DataFrame, window: &GraphQueryWindow) -> DfResult DfResult DfResult DfResult DfResult DfResult DfResult { /// `connection_type`. fn virtual_attrs_expr() -> DfResult { when( - ident("virtual_conn").eq(lit("database")), - parse_json_expr(lit(r#"{"connection_type":"database"}"#)), + ident("virtual_conn").eq(lit(CONNECTION_TYPE_DATABASE)), + parse_json_expr(lit(format!( + r#"{{"connection_type":"{CONNECTION_TYPE_DATABASE}"}}"# + ))), ) - .otherwise(parse_json_expr(lit( - r#"{"connection_type":"virtual_node"}"#, - ))) + .otherwise(parse_json_expr(lit(format!( + r#"{{"connection_type":"{CONNECTION_TYPE_VIRTUAL_NODE}"}}"# + )))) } /// The window + span-kind + identity predicate shared by both join sides. @@ -506,9 +493,9 @@ fn span_predicate(service: &EntityDeclaration, window: &GraphQueryWindow, strict .gt_eq(window.source_start() - interval(CHILD_SPAN_EARLY_NANOS)) .and(ts.lt(window.source_end() + interval(CHILD_SPAN_LATE_NANOS))) }; - // A NULL identity component identifies nothing, on either endpoint. + // An absent identity component identifies nothing, on either endpoint. for id in &service.id_columns { - predicate = predicate.and(ident(id).is_not_null()); + predicate = predicate.and(identifies(id)); } predicate } @@ -524,16 +511,22 @@ fn client_spans( scan: &DataFrame, window: &GraphQueryWindow, ) -> DfResult { - // NULLIF('') lets an empty value fall through to the next candidate. + // Attribute columns are dynamic in `greptime_trace_v1` (created on first + // use), so only candidate columns present in the table's schema + // participate. NULLIF('') lets an empty value fall through to the next + // candidate. let present: Vec<(Expr, &str)> = { let schema = scan.schema(); - VIRTUAL_DST_CANDIDATES + builtin()? + .virtual_dst_candidates .iter() - .filter(|(column, _)| schema.has_column_with_unqualified_name(column)) - .map(|(column, conn)| { - let value = - core_fns::nullif().call(vec![cast(ident(*column), DataType::Utf8), lit("")]); - (value, *conn) + .filter(|candidate| schema.has_column_with_unqualified_name(&candidate.column)) + .map(|candidate| { + let value = core_fns::nullif().call(vec![ + cast(ident(&candidate.column), DataType::Utf8), + lit(""), + ]); + (value, candidate.connection_type.as_str()) }) .collect() }; @@ -560,7 +553,7 @@ fn client_spans( ident(SPAN_ID_COLUMN), // The cast inside entity_id_expr also normalizes tag columns, which // come out of the storage engine dictionary-encoded. - entity_id_expr(&service.id_columns, &|c| ident(c)).alias("src_id"), + entity_id_expr(&service.id_columns, &|c| ident(c)).alias(SRC_ID_COLUMN), ident(SPAN_STATUS_CODE_COLUMN), ident(DURATION_NANO_COLUMN), virtual_dst.alias("virtual_dst"), @@ -585,7 +578,7 @@ fn server_spans( ident(TRACE_TIMESTAMP_COLUMN), ident(TRACE_ID_COLUMN), ident(PARENT_SPAN_ID_COLUMN), - entity_id_expr(&service.id_columns, &|c| ident(c)).alias("dst_id"), + entity_id_expr(&service.id_columns, &|c| ident(c)).alias(DST_ID_COLUMN), ident(SPAN_STATUS_CODE_COLUMN), ident(DURATION_NANO_COLUMN), ]) @@ -616,7 +609,7 @@ fn agent_calls_branch( ident(TRACE_TIMESTAMP_COLUMN), ident(TRACE_ID_COLUMN), ident(SPAN_ID_COLUMN), - entity_id_expr(&agent.id_columns, &|c| ident(c)).alias("src_id"), + entity_id_expr(&agent.id_columns, &|c| ident(c)).alias(SRC_ID_COLUMN), ])?; let child = trace .scan @@ -626,7 +619,7 @@ fn agent_calls_branch( ident(TRACE_TIMESTAMP_COLUMN), ident(TRACE_ID_COLUMN), ident(PARENT_SPAN_ID_COLUMN), - entity_id_expr(&agent.id_columns, &|c| ident(c)).alias("dst_id"), + entity_id_expr(&agent.id_columns, &|c| ident(c)).alias(DST_ID_COLUMN), ident(SPAN_STATUS_CODE_COLUMN), ident(DURATION_NANO_COLUMN), ])?; @@ -651,44 +644,48 @@ fn agent_calls_branch( let df = parent .join_on(child, JoinType::Inner, join_conditions)? .select(vec![ - bin_ms(qcol("parent", TRACE_TIMESTAMP_COLUMN)).alias("observed_at"), - qcol("parent", "src_id").alias("src_id"), - qcol("child", "dst_id").alias("dst_id"), + bin_ms(qcol("parent", TRACE_TIMESTAMP_COLUMN)).alias(OBSERVED_AT_COLUMN), + qcol("parent", SRC_ID_COLUMN).alias(SRC_ID_COLUMN), + qcol("child", DST_ID_COLUMN).alias(DST_ID_COLUMN), qcol("child", SPAN_STATUS_CODE_COLUMN).alias("status_code"), - qcol("child", DURATION_NANO_COLUMN).alias("duration_nano"), + qcol("child", DURATION_NANO_COLUMN).alias(DURATION_NANO_COLUMN), ])? // A sub-span of the same agent is internal structure, not a call // between two agents. - .filter(ident("src_id").not_eq(ident("dst_id")))? + .filter(ident(SRC_ID_COLUMN).not_eq(ident(DST_ID_COLUMN)))? .aggregate( - vec![ident("observed_at"), ident("src_id"), ident("dst_id")], vec![ - count(lit(1)).alias("request_count"), + ident(OBSERVED_AT_COLUMN), + ident(SRC_ID_COLUMN), + ident(DST_ID_COLUMN), + ], + vec![ + count(lit(1)).alias(REQUEST_COUNT_COLUMN), count(lit(1)) .filter(ident("status_code").eq(lit(SPAN_STATUS_ERROR))) .build()? - .alias("error_count"), - sum(ident("duration_nano")).alias("duration_nano_sum"), + .alias(ERROR_COUNT_COLUMN), + sum(ident(DURATION_NANO_COLUMN)).alias("duration_nano_sum"), ], )? .select(vec![ - ident("observed_at"), - ident("observed_at").alias("window_start"), - (ident("observed_at") + bin_interval()).alias("window_end"), - (ident("observed_at") + bin_interval()).alias("fresh_until"), - lit("agent").alias("src_type"), - ident("src_id"), - lit("agent").alias("dst_type"), - ident("dst_id"), - lit("calls").alias("rel_type"), - lit("trace").alias("provenance"), - lit(1.0_f64).alias("confidence"), - ident("request_count"), - ident("error_count"), + ident(OBSERVED_AT_COLUMN), + ident(OBSERVED_AT_COLUMN).alias(WINDOW_START_COLUMN), + (ident(OBSERVED_AT_COLUMN) + bin_interval()).alias(WINDOW_END_COLUMN), + (ident(OBSERVED_AT_COLUMN) + bin_interval()).alias(FRESH_UNTIL_COLUMN), + lit(ENTITY_TYPE_GEN_AI_AGENT).alias(SRC_TYPE_COLUMN), + ident(SRC_ID_COLUMN), + lit(ENTITY_TYPE_GEN_AI_AGENT).alias(DST_TYPE_COLUMN), + ident(DST_ID_COLUMN), + lit(REL_TYPE_CALLS).alias(REL_TYPE_COLUMN), + lit(PROVENANCE_TRACE).alias(PROVENANCE_COLUMN), + lit(1.0_f64).alias(CONFIDENCE_COLUMN), + ident(REQUEST_COUNT_COLUMN), + ident(ERROR_COUNT_COLUMN), (cast(ident("duration_nano_sum"), DataType::Float64) / lit(1e9_f64)) - .alias("duration_sum"), - ident("request_count").alias("duration_count"), - null_json().alias("attributes"), + .alias(DURATION_SUM_COLUMN), + ident(REQUEST_COUNT_COLUMN).alias(DURATION_COUNT_COLUMN), + null_json().alias(EDGE_ATTRIBUTES_COLUMN), ])?; Ok(Some(df)) } @@ -1173,7 +1170,7 @@ mod tests { RelationshipSources { traces: vec![CallsSource { service: None, - agent: Some(co_decl("agent", &["agent_id"])), + agent: Some(co_decl("gen_ai.agent", &["agent_id"])), scan, }], co_declared: vec![], @@ -1190,9 +1187,9 @@ mod tests { // no service declaration. assert_eq!(total, 1); let batch = &batches[0]; - assert_eq!(strings(batch, 4), vec!["agent"]); + assert_eq!(strings(batch, 4), vec!["gen_ai.agent"]); assert_eq!(strings(batch, 5), vec!["orchestrator"]); - assert_eq!(strings(batch, 6), vec!["agent"]); + assert_eq!(strings(batch, 6), vec!["gen_ai.agent"]); assert_eq!(strings(batch, 7), vec!["researcher"]); assert_eq!(strings(batch, 8), vec!["calls"]); assert_eq!(strings(batch, 9), vec!["trace"]); @@ -1229,7 +1226,7 @@ mod tests { RelationshipSources { traces: vec![CallsSource { service: None, - agent: Some(co_decl("agent", &["agent_id"])), + agent: Some(co_decl("gen_ai.agent", &["agent_id"])), scan, }], co_declared: vec![], @@ -1444,9 +1441,9 @@ mod tests { let ctx = co_declared_ctx(); let declarations = || { vec![ - co_decl("agent", &["agent_id"]), - co_decl("model", &["model_name"]), - co_decl("tool", &["tool_name"]), + co_decl("gen_ai.agent", &["agent_id"]), + co_decl("gen_ai.model", &["model_name"]), + co_decl("gen_ai.tool", &["tool_name"]), ] }; // A non-trace table co-declaring agent/model/tool must not fabricate @@ -1462,17 +1459,17 @@ mod tests { rows, vec![ ( - "agent".to_string(), + "gen_ai.agent".to_string(), "agent-1".to_string(), - "model".to_string(), + "gen_ai.model".to_string(), "gpt".to_string(), "uses".to_string(), "trace".to_string(), ), ( - "agent".to_string(), + "gen_ai.agent".to_string(), "agent-1".to_string(), - "tool".to_string(), + "gen_ai.tool".to_string(), "search".to_string(), "invokes".to_string(), "trace".to_string(), diff --git a/src/table/src/requests/semantic.rs b/src/table/src/requests/semantic.rs index cb9d6aa7476..5c0e4c3caed 100644 --- a/src/table/src/requests/semantic.rs +++ b/src/table/src/requests/semantic.rs @@ -209,9 +209,9 @@ pub fn is_semantic_option_key(key: &str) -> bool { /// non-empty string. Closed-domain keys accept a fixed set, plus the `unknown` /// sentinel, plus `mixed` for the keys where one long-lived table can /// legitimately see multiple values. Entity keys ([`is_entity_option_key`]) take a -/// comma-separated column-name list (each token non-empty); the "id columns must -/// be tag columns" invariant is enforced later against the table schema at DDL -/// time, not here. Keys that are neither whitelisted nor a well-formed entity key +/// comma-separated column-name list (each token non-empty); column existence and +/// the stable-string-form rule for entity columns are enforced later against +/// the table schema at DDL time, not here. Keys that are neither whitelisted nor a well-formed entity key /// are rejected. pub fn validate_semantic_option(key: &str, value: &str) -> bool { if is_entity_option_key(key) { diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 042be385a6d..b2c8fe394e2 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -161,6 +161,7 @@ macro_rules! http_tests { test_otlp_metric_translation_strategies, test_otlp_traces_v0, test_otlp_traces_v1, + test_otlp_traces_v1_entity_graph, test_otlp_logs, test_loki_pb_logs, test_loki_pb_logs_with_pipeline, @@ -6886,6 +6887,86 @@ pub async fn test_otlp_traces_v0(store_type: StorageType) { guard.remove_all().await; } +/// One real OTLP export must come out of `semantic_relationships` as the +/// zero-configuration chain; resources missing an identity column derive +/// nothing extra. +pub async fn test_otlp_traces_v1_entity_graph(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = + setup_test_http_app_with_frontend(store_type, "test_otlp_traces_entity_graph").await; + let client = TestClient::new(app).await; + + let span = |trace_id: &str, span_id: &str, parent: &str, kind: u32| { + json!({ + "traceId": trace_id, + "spanId": span_id, + "parentSpanId": parent, + "name": "op", + "kind": kind, + "startTimeUnixNano": "1736480942444376000", + "endTimeUnixNano": "1736480942444499000", + "attributes": [], + "status": { "message": "", "code": 0 } + }) + }; + let resource_spans = |attributes: Vec, spans: Vec| { + json!({ + "resource": { "attributes": attributes }, + "scopeSpans": [{ "scope": { "name": "graph-tests" }, "spans": spans }], + "schemaUrl": "https://opentelemetry.io/schemas/1.4.0" + }) + }; + + // frontend: full identity (instance id + pod uid); cart: bare service — + // no instance entity; batch: instance id but no pod uid — no runs_on. + let req: ExportTraceServiceRequest = serde_json::from_value(json!({ + "resourceSpans": [ + resource_spans( + vec![ + make_string_attr("service.name", "frontend"), + make_string_attr("service.instance.id", "inst-1"), + make_string_attr("k8s.pod.uid", "uid-9"), + make_string_attr("k8s.pod.name", "api-1"), + make_string_attr("k8s.namespace.name", "default"), + ], + vec![span("c05d7a4ec8e1f231f02ed6e8da8655b4", "d24f921c75f68e23", "", 3)], + ), + resource_spans( + vec![make_string_attr("service.name", "cart")], + vec![span( + "c05d7a4ec8e1f231f02ed6e8da8655b4", + "9630f2916e2f7909", + "d24f921c75f68e23", + 2, + )], + ), + resource_spans( + vec![ + make_string_attr("service.name", "batch"), + make_string_attr("service.instance.id", "binst-1"), + ], + vec![span("cc9e0991a2e63d274984bd44ee669203", "8f847259b0f6e1ab", "", 1)], + ), + ] + })) + .unwrap(); + let res = send_trace_v1_req(&client, "graph_traces", req, false).await; + assert_eq!(StatusCode::OK, res.status()); + + let expected = r#"[["service","frontend","service","cart","calls"],["service.instance","resource_attributes.service.instance.id=binst-1,service_name=batch","service","batch","part_of"],["service.instance","resource_attributes.service.instance.id=inst-1,service_name=frontend","service","frontend","part_of"],["service.instance","resource_attributes.service.instance.id=inst-1,service_name=frontend","k8s.pod","uid-9","runs_on"]]"#; + validate_data( + "otlp_traces_entity_graph", + &client, + "select src_type, src_id, dst_type, dst_id, rel_type \ + from greptime_private.semantic_relationships \ + where observed_at >= '2025-01-01 00:00:00' order by rel_type, src_id;", + expected, + ) + .await; + + guard.remove_all().await; +} + pub async fn test_otlp_traces_v1(store_type: StorageType) { // init common_telemetry::init_default_ut_logging(); diff --git a/tests/cases/standalone/common/system/semantic_graph.result b/tests/cases/standalone/common/system/semantic_graph.result index 6617d9f839d..ed6a3d3f16f 100644 --- a/tests/cases/standalone/common/system/semantic_graph.result +++ b/tests/cases/standalone/common/system/semantic_graph.result @@ -44,47 +44,8 @@ drop table greptime_private.semantic_relationships; Error: 1001(Unsupported), Table `semantic_relationships` is read-only --- Read-time derivation: declare entity identities on a user table and observe --- them in the computed registry (single-column and composite ids, scope and --- descriptive columns, JSON output columns). -create table graph_app_latency ( - ts timestamp time index, - service_name string, - host string, - env string, - latency double, - primary key (service_name, host, env) -) with ( - 'greptime.semantic.entity.service.id' = 'service_name', - 'greptime.semantic.entity.service.scope' = 'env', - 'greptime.semantic.entity.process.id' = 'service_name,host', - 'greptime.semantic.entity.process.descriptive' = 'env' -); - -Affected Rows: 0 - -insert into graph_app_latency values (now(), 'cart', 'h1', 'us-east', 3.5); - -Affected Rows: 1 - --- SQLNESS PROTOCOL MYSQL -select entity_type, entity_id, entity_id_attrs, scope, descriptive, source_tables -from greptime_private.semantic_entities -order by entity_type, entity_id; - -+-------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ -| entity_type | entity_id | entity_id_attrs | scope | descriptive | source_tables | -+-------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ -| process | host=h1,service_name=cart | {"host":"h1","service_name":"cart"} | | {"env":"us-east"} | ["public.graph_app_latency"] | -| service | cart | | us-east | | ["public.graph_app_latency"] | -+-------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ - -drop table graph_app_latency; - -Affected Rows: 0 - --- Renaming a physical table INTO a reserved computed-table name would let the --- overlay shadow it, orphaning its data. +-- Renaming a physical table INTO a reserved name would let the overlay shadow +-- it (or squat the declared table's canonical schema), orphaning its data. create table greptime_private.graph_rename_probe (ts timestamp time index); Affected Rows: 0 @@ -93,13 +54,96 @@ alter table greptime_private.graph_rename_probe rename semantic_entities; Error: 1001(Unsupported), Table `semantic_entities` is read-only +alter table greptime_private.graph_rename_probe rename semantic_relationships_declared; + +Error: 1001(Unsupported), The definition of table `semantic_relationships_declared` is managed by GreptimeDB; it cannot be created or altered (DROP recreates it on the next write) + drop table greptime_private.graph_rename_probe; Affected Rows: 0 --- Calls derivation: a trace-v1-model table with client/server span pairs. --- Trace tables need no entity declaration: the service entities are implicit. -create table graph_traces ( +-- Read-time derivation from one declaring metric table: entity identities in +-- the computed registry (single-column and composite ids, scope and +-- descriptive columns, JSON output columns), and the same rows witnessing the +-- co-declared vocabulary edges (runs_on / part_of, direction built in). +create table graph_app_metrics ( + ts timestamp time index, + service_name string, + instance string, + host string, + env string, + latency double, + primary key (service_name, instance, host, env) +) with ( + 'greptime.semantic.entity.service.id' = 'service_name', + 'greptime.semantic.entity.service.scope' = 'env', + 'greptime.semantic.entity.service.instance.id' = 'instance', + 'greptime.semantic.entity.host.id' = 'host', + 'greptime.semantic.entity.process.id' = 'service_name,host', + 'greptime.semantic.entity.process.descriptive' = 'env' +); + +Affected Rows: 0 + +insert into graph_app_metrics values (now(), 'cart', 'cart-0', 'h1', 'us-east', 3.5); + +Affected Rows: 1 + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, entity_id_attrs, scope, descriptive, source_tables +from greptime_private.semantic_entities +order by entity_type, entity_id; + ++------------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ +| entity_type | entity_id | entity_id_attrs | scope | descriptive | source_tables | ++------------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ +| host | h1 | | | | ["public.graph_app_metrics"] | +| process | host=h1,service_name=cart | {"host":"h1","service_name":"cart"} | | {"env":"us-east"} | ["public.graph_app_metrics"] | +| service | cart | | us-east | | ["public.graph_app_metrics"] | +| service.instance | cart-0 | | | | ["public.graph_app_metrics"] | ++------------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ + +select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence +from greptime_private.semantic_relationships +order by rel_type, src_id; + ++------------------+---------------------------+----------+--------+----------+------------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | confidence | ++------------------+---------------------------+----------+--------+----------+------------+------------+ +| service.instance | cart-0 | service | cart | part_of | attribute | 1.0 | +| service.instance | cart-0 | host | h1 | runs_on | attribute | 1.0 | +| process | host=h1,service_name=cart | host | h1 | runs_on | attribute | 1.0 | ++------------------+---------------------------+----------+--------+----------+------------+------------+ + +drop table graph_app_metrics; + +Affected Rows: 0 + +-- Calls derivation over trace-v1 tables, all branches in one scan: the client +-- span and its child server span land in different tables and still pair into +-- one edge with RED metrics; the epoch-timestamped pair falls outside the +-- default one-hour window; unmatched clients become virtual-node edges named +-- from span attributes (confidence < 1.0, connection type); a trace-model +-- table that lost the fixed span columns is skipped instead of failing the +-- scan. Trace tables need no entity declaration: service entities are +-- implicit. +create table graph_traces_a ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + "span_attributes.peer.service" string, + "span_attributes.db.name" string, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +Affected Rows: 0 + +create table graph_traces_b ( "timestamp" timestamp(9) time index, trace_id string, span_id string, @@ -113,43 +157,68 @@ create table graph_traces ( Affected Rows: 0 --- The epoch-timestamped pair falls outside the default one-hour derivation --- window and must not appear in the results. -insert into graph_traces values - (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0), - (now(), 't1', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_ERROR', 'cart', 1500000000), - (now(), 't2', 'c2', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0), - (now(), 't2', 's2', 'c2', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'cart', 500000000), - (0, 't0', 'c0', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'stale-src', 0), - (0, 't0', 's0', 'c0', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'stale-dst', 100); +create table graph_traces_malformed ( + ts timestamp time index, + note string +) with ('table_data_model' = 'greptime_trace_v1'); + +Affected Rows: 0 + +insert into graph_traces_a values + (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0, NULL, NULL), + (now(), 't2', 'c2', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0, NULL, NULL), + (0, 't0', 'c0', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'stale-src', 0, NULL, NULL), + (0, 't0', 's0', 'c0', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'stale-dst', 100, NULL, NULL), + (now(), 't3', 'c3', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 250000000, 'redis', NULL), + (now(), 't4', 'c4', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 100000000, NULL, 'orders-db'); Affected Rows: 6 +insert into graph_traces_b values + (now(), 't1', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_ERROR', 'cart', 1500000000), + (now(), 't2', 's2', 'c2', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'cart', 500000000); + +Affected Rows: 2 + +insert into graph_traces_malformed values (now(), 'not a trace'); + +Affected Rows: 1 + -- SQLNESS PROTOCOL MYSQL -select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence, +select src_id, dst_id, rel_type, provenance, confidence, request_count, error_count, duration_sum, duration_count, attributes from greptime_private.semantic_relationships -order by src_id; +order by dst_id; -+----------+----------+----------+--------+----------+------------+------------+---------------+-------------+--------------+----------------+------------+ -| src_type | src_id | dst_type | dst_id | rel_type | provenance | confidence | request_count | error_count | duration_sum | duration_count | attributes | -+----------+----------+----------+--------+----------+------------+------------+---------------+-------------+--------------+----------------+------------+ -| service | frontend | service | cart | calls | trace | 1 | 2 | 1 | 2 | 2 | | -+----------+----------+----------+--------+----------+------------+------------+---------------+-------------+--------------+----------------+------------+ ++----------+-----------+----------+------------+------------+---------------+-------------+--------------+----------------+------------------------------------+ +| src_id | dst_id | rel_type | provenance | confidence | request_count | error_count | duration_sum | duration_count | attributes | ++----------+-----------+----------+------------+------------+---------------+-------------+--------------+----------------+------------------------------------+ +| frontend | cart | calls | trace | 1 | 2 | 1 | 2 | 2 | | +| frontend | orders-db | calls | trace | 0.5 | 1 | 0 | 0.1 | 1 | {"connection_type":"database"} | +| frontend | redis | calls | trace | 0.5 | 1 | 0 | 0.25 | 1 | {"connection_type":"virtual_node"} | ++----------+-----------+----------+------------+------------+---------------+-------------+--------------+----------------+------------------------------------+ -- SQLNESS PROTOCOL MYSQL select entity_type, entity_id, entity_id_attrs, scope, source_tables from greptime_private.semantic_entities order by entity_id; -+-------------+-----------+-----------------+-------+-------------------------+ -| entity_type | entity_id | entity_id_attrs | scope | source_tables | -+-------------+-----------+-----------------+-------+-------------------------+ -| service | cart | | | ["public.graph_traces"] | -| service | frontend | | | ["public.graph_traces"] | -+-------------+-----------+-----------------+-------+-------------------------+ ++-------------+-----------+-----------------+-------+---------------------------+ +| entity_type | entity_id | entity_id_attrs | scope | source_tables | ++-------------+-----------+-----------------+-------+---------------------------+ +| service | cart | | | ["public.graph_traces_b"] | +| service | frontend | | | ["public.graph_traces_a"] | ++-------------+-----------+-----------------+-------+---------------------------+ -drop table graph_traces; +drop table graph_traces_a; + +Affected Rows: 0 + +drop table graph_traces_b; + +Affected Rows: 0 + +drop table graph_traces_malformed; Affected Rows: 0 @@ -222,8 +291,7 @@ where observed_at < '2001-01-02 00:00:00'; Error: 3001(EngineExecuteQuery), Invalid SQL, error: the observed_at filter has no lower bound; the graph cannot derive over unbounded history — add e.g. observed_at >= '2026-01-01 00:00:00' -- The declared-edge table's definition is system-owned: user CREATE/ALTER are --- rejected and a physical table cannot be renamed into the name, while plain --- DML stays allowed. +-- rejected, while plain DML stays allowed. create table greptime_private.semantic_relationships_declared (ts timestamp time index); Error: 1001(Unsupported), The definition of table `semantic_relationships_declared` is managed by GreptimeDB; it cannot be created or altered (DROP recreates it on the next write) @@ -232,18 +300,6 @@ alter table greptime_private.semantic_relationships_declared add column extra st Error: 1001(Unsupported), The definition of table `semantic_relationships_declared` is managed by GreptimeDB; it cannot be created or altered (DROP recreates it on the next write) -create table greptime_private.declared_rename_probe (ts timestamp time index); - -Affected Rows: 0 - -alter table greptime_private.declared_rename_probe rename semantic_relationships_declared; - -Error: 1001(Unsupported), The definition of table `semantic_relationships_declared` is managed by GreptimeDB; it cannot be created or altered (DROP recreates it on the next write) - -drop table greptime_private.declared_rename_probe; - -Affected Rows: 0 - delete from greptime_private.semantic_relationships_declared; Affected Rows: 4 @@ -277,161 +333,9 @@ drop table greptime_private.semantic_relationships_declared; Affected Rows: 0 --- Same-row co-declaration: a table declaring both entity types of a built-in --- vocabulary pair witnesses the edge on every row carrying both identities; --- the vocabulary fixes the direction. -create table graph_pod_metrics ( - ts timestamp time index, - instance string, - host string, - "service" string, - cpu double, - primary key (instance, host, "service") -) with ( - 'greptime.semantic.entity.service.instance.id' = 'instance', - 'greptime.semantic.entity.host.id' = 'host', - 'greptime.semantic.entity.service.id' = 'service' -); - -Affected Rows: 0 - -insert into graph_pod_metrics values (now(), 'cart-0', 'node-1', 'cart', 0.5); - -Affected Rows: 1 - -select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence -from greptime_private.semantic_relationships -order by rel_type, src_id; - -+------------------+--------+----------+--------+----------+------------+------------+ -| src_type | src_id | dst_type | dst_id | rel_type | provenance | confidence | -+------------------+--------+----------+--------+----------+------------+------------+ -| service.instance | cart-0 | service | cart | part_of | attribute | 1.0 | -| service.instance | cart-0 | host | node-1 | runs_on | attribute | 1.0 | -+------------------+--------+----------+--------+----------+------------+------------+ - -drop table graph_pod_metrics; - -Affected Rows: 0 - --- Cross-table calls pairing: the client span and its child server span land in --- different trace tables and must still pair into one edge. A trace-model --- table that lost the fixed span columns is skipped instead of failing the --- scan. -create table graph_traces_a ( - "timestamp" timestamp(9) time index, - trace_id string, - span_id string, - parent_span_id string, - span_kind string, - span_status_code string, - service_name string, - duration_nano bigint unsigned, - primary key (service_name) -) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); - -Affected Rows: 0 - -create table graph_traces_b ( - "timestamp" timestamp(9) time index, - trace_id string, - span_id string, - parent_span_id string, - span_kind string, - span_status_code string, - service_name string, - duration_nano bigint unsigned, - primary key (service_name) -) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); - -Affected Rows: 0 - -create table graph_traces_malformed ( - ts timestamp time index, - note string -) with ('table_data_model' = 'greptime_trace_v1'); - -Affected Rows: 0 - -insert into graph_traces_a values - (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0); - -Affected Rows: 1 - -insert into graph_traces_b values - (now(), 't1', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_ERROR', 'cart', 1500000000); - -Affected Rows: 1 - -insert into graph_traces_malformed values (now(), 'not a trace'); - -Affected Rows: 1 - -select src_id, dst_id, rel_type, provenance, confidence, request_count, error_count -from greptime_private.semantic_relationships -order by src_id; - -+----------+--------+----------+------------+------------+---------------+-------------+ -| src_id | dst_id | rel_type | provenance | confidence | request_count | error_count | -+----------+--------+----------+------------+------------+---------------+-------------+ -| frontend | cart | calls | trace | 1.0 | 1 | 1 | -+----------+--------+----------+------------+------------+---------------+-------------+ - -drop table graph_traces_a; - -Affected Rows: 0 - -drop table graph_traces_b; - -Affected Rows: 0 - -drop table graph_traces_malformed; - -Affected Rows: 0 - --- Virtual nodes: a client span with no matching server span names its peer --- through span attributes; the edge carries confidence < 1.0 and the --- connection type. -create table graph_traces_virtual ( - "timestamp" timestamp(9) time index, - trace_id string, - span_id string, - parent_span_id string, - span_kind string, - span_status_code string, - service_name string, - duration_nano bigint unsigned, - "span_attributes.peer.service" string, - "span_attributes.db.name" string, - primary key (service_name) -) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); - -Affected Rows: 0 - -insert into graph_traces_virtual values - (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 250000000, 'redis', NULL), - (now(), 't2', 'c2', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 100000000, NULL, 'orders-db'); - -Affected Rows: 2 - --- SQLNESS PROTOCOL MYSQL -select src_id, dst_id, rel_type, confidence, request_count, attributes -from greptime_private.semantic_relationships -order by dst_id; - -+----------+-----------+----------+------------+---------------+------------------------------------+ -| src_id | dst_id | rel_type | confidence | request_count | attributes | -+----------+-----------+----------+------------+---------------+------------------------------------+ -| frontend | orders-db | calls | 0.5 | 1 | {"connection_type":"database"} | -| frontend | redis | calls | 0.5 | 1 | {"connection_type":"virtual_node"} | -+----------+-----------+----------+------------+---------------+------------------------------------+ - -drop table graph_traces_virtual; - -Affected Rows: 0 - -- Agent edges: span structure derives parent_agent calls agent, and span rows --- co-declaring agent+model / agent+tool witness uses / invokes. +-- co-declaring gen_ai.agent+gen_ai.model / gen_ai.agent+gen_ai.tool witness +-- uses / invokes. The identity columns are fields, not tags. create table graph_agent_traces ( "timestamp" timestamp(9) time index, trace_id string, @@ -444,13 +348,13 @@ create table graph_agent_traces ( agent_id string, model_name string, tool_name string, - primary key (service_name, agent_id, model_name, tool_name) + primary key (service_name) ) with ( 'table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true', - 'greptime.semantic.entity.agent.id' = 'agent_id', - 'greptime.semantic.entity.model.id' = 'model_name', - 'greptime.semantic.entity.tool.id' = 'tool_name' + 'greptime.semantic.entity.gen_ai.agent.id' = 'agent_id', + 'greptime.semantic.entity.gen_ai.model.id' = 'model_name', + 'greptime.semantic.entity.gen_ai.tool.id' = 'tool_name' ); Affected Rows: 0 @@ -466,15 +370,205 @@ select src_type, src_id, dst_type, dst_id, rel_type, provenance from greptime_private.semantic_relationships order by rel_type, dst_id; -+----------+--------------+----------+------------+----------+------------+ -| src_type | src_id | dst_type | dst_id | rel_type | provenance | -+----------+--------------+----------+------------+----------+------------+ -| agent | orchestrator | agent | researcher | calls | trace | -| agent | researcher | tool | web_search | invokes | trace | -| agent | researcher | model | gpt-5 | uses | trace | -+----------+--------------+----------+------------+----------+------------+ ++--------------+--------------+--------------+------------+----------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | ++--------------+--------------+--------------+------------+----------+------------+ +| gen_ai.agent | orchestrator | gen_ai.agent | researcher | calls | trace | +| gen_ai.agent | researcher | gen_ai.tool | web_search | invokes | trace | +| gen_ai.agent | researcher | gen_ai.model | gpt-5 | uses | trace | ++--------------+--------------+--------------+------------+----------+------------+ drop table graph_agent_traces; Affected Rows: 0 +-- Prometheus conventions: whitelisted entity-descriptor metrics (stamped +-- signal_type=metric + source=prometheus by the remote-write path) get +-- implicit declarations, and the co-declared rules derive runs_on / part_of / +-- contains from their rows. Pod identity is the UID, so every KSM descriptor +-- lands on one entity; the empty labels kube-state-metrics emits (unscheduled +-- pod's node, ownerless pod's owner_*) identify nothing. A non-whitelisted +-- metric contributes nothing. +create table kube_pod_info ( + greptime_timestamp timestamp(3) time index, + uid string, + namespace string, + pod string, + node string, + greptime_value double, + primary key (uid, namespace, pod, node) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +Affected Rows: 0 + +create table kube_pod_owner ( + greptime_timestamp timestamp(3) time index, + uid string, + namespace string, + pod string, + owner_kind string, + owner_name string, + greptime_value double, + primary key (uid, namespace, pod, owner_kind, owner_name) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +Affected Rows: 0 + +create table kube_pod_container_info ( + greptime_timestamp timestamp(3) time index, + uid string, + container string, + image string, + greptime_value double, + primary key (uid, container, image) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +Affected Rows: 0 + +create table kube_service_info ( + greptime_timestamp timestamp(3) time index, + uid string, + namespace string, + "service" string, + cluster_ip string, + greptime_value double, + primary key (uid, namespace, "service", cluster_ip) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +Affected Rows: 0 + +create table target_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + k8s_cluster_name string, + greptime_value double, + primary key (job, instance, k8s_cluster_name) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +Affected Rows: 0 + +create table http_requests_total ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + greptime_value double, + primary key (job, instance) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +Affected Rows: 0 + +insert into kube_pod_info values + (now(), 'uid-1', 'default', 'api-1', 'node-a', 1), + (now(), 'uid-2', 'default', 'api-2', 'node-a', 1), + (now(), 'uid-3', 'default', 'pending-1', '', 1), + (now(), '', 'default', 'ghost', 'node-a', 1); + +Affected Rows: 4 + +insert into kube_pod_owner values + (now(), 'uid-1', 'default', 'api-1', 'ReplicaSet', 'api-rs', 1), + (now(), 'uid-2', 'default', 'api-2', 'ReplicaSet', 'api-rs', 1); + +Affected Rows: 2 + +insert into kube_pod_container_info values + (now(), 'uid-1', 'main', 'nginx:1.25', 1); + +Affected Rows: 1 + +insert into kube_service_info values + (now(), 'svc-uid-1', 'default', 'api-svc', '10.0.0.1', 1); + +Affected Rows: 1 + +insert into target_info values + (now(), 'shop/api', 'inst-1', 'prod', 1); + +Affected Rows: 1 + +insert into http_requests_total values + (now(), 'shop/api', 'inst-1', 42); + +Affected Rows: 1 + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +order by entity_type, entity_id, source_tables; + ++------------------+-----------------------------------------------------------+------------------------------------+ +| entity_type | entity_id | source_tables | ++------------------+-----------------------------------------------------------+------------------------------------+ +| k8s.container | container=main,uid=uid-1 | ["public.kube_pod_container_info"] | +| k8s.node | node-a | ["public.kube_pod_info"] | +| k8s.pod | uid-1 | ["public.kube_pod_info"] | +| k8s.pod | uid-1 | ["public.kube_pod_owner"] | +| k8s.pod | uid-1 | ["public.kube_pod_container_info"] | +| k8s.pod | uid-2 | ["public.kube_pod_info"] | +| k8s.pod | uid-2 | ["public.kube_pod_owner"] | +| k8s.pod | uid-3 | ["public.kube_pod_info"] | +| k8s.service | svc-uid-1 | ["public.kube_service_info"] | +| k8s.workload | namespace=default,owner_kind=ReplicaSet,owner_name=api-rs | ["public.kube_pod_owner"] | +| service | shop/api | ["public.target_info"] | +| service.instance | instance=inst-1,job=shop/api | ["public.target_info"] | ++------------------+-----------------------------------------------------------+------------------------------------+ + +-- SQLNESS PROTOCOL MYSQL +select src_type, src_id, dst_type, dst_id, rel_type, provenance +from greptime_private.semantic_relationships +order by rel_type, src_id, dst_id; + ++------------------+------------------------------+---------------+-----------------------------------------------------------+----------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | ++------------------+------------------------------+---------------+-----------------------------------------------------------+----------+------------+ +| k8s.pod | uid-1 | k8s.container | container=main,uid=uid-1 | contains | attribute | +| service.instance | instance=inst-1,job=shop/api | service | shop/api | part_of | attribute | +| k8s.pod | uid-1 | k8s.workload | namespace=default,owner_kind=ReplicaSet,owner_name=api-rs | part_of | attribute | +| k8s.pod | uid-2 | k8s.workload | namespace=default,owner_kind=ReplicaSet,owner_name=api-rs | part_of | attribute | +| k8s.pod | uid-1 | k8s.node | node-a | runs_on | attribute | +| k8s.pod | uid-2 | k8s.node | node-a | runs_on | attribute | ++------------------+------------------------------+---------------+-----------------------------------------------------------+----------+------------+ + +drop table kube_pod_info; + +Affected Rows: 0 + +drop table kube_pod_owner; + +Affected Rows: 0 + +drop table kube_pod_container_info; + +Affected Rows: 0 + +drop table kube_service_info; + +Affected Rows: 0 + +drop table target_info; + +Affected Rows: 0 + +drop table http_requests_total; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/system/semantic_graph.sql b/tests/cases/standalone/common/system/semantic_graph.sql index 9acb003331a..181a8da66a9 100644 --- a/tests/cases/standalone/common/system/semantic_graph.sql +++ b/tests/cases/standalone/common/system/semantic_graph.sql @@ -22,43 +22,73 @@ drop table greptime_private.semantic_entities; drop table greptime_private.semantic_relationships; --- Read-time derivation: declare entity identities on a user table and observe --- them in the computed registry (single-column and composite ids, scope and --- descriptive columns, JSON output columns). -create table graph_app_latency ( +-- Renaming a physical table INTO a reserved name would let the overlay shadow +-- it (or squat the declared table's canonical schema), orphaning its data. +create table greptime_private.graph_rename_probe (ts timestamp time index); + +alter table greptime_private.graph_rename_probe rename semantic_entities; + +alter table greptime_private.graph_rename_probe rename semantic_relationships_declared; + +drop table greptime_private.graph_rename_probe; + +-- Read-time derivation from one declaring metric table: entity identities in +-- the computed registry (single-column and composite ids, scope and +-- descriptive columns, JSON output columns), and the same rows witnessing the +-- co-declared vocabulary edges (runs_on / part_of, direction built in). +create table graph_app_metrics ( ts timestamp time index, service_name string, + instance string, host string, env string, latency double, - primary key (service_name, host, env) + primary key (service_name, instance, host, env) ) with ( 'greptime.semantic.entity.service.id' = 'service_name', 'greptime.semantic.entity.service.scope' = 'env', + 'greptime.semantic.entity.service.instance.id' = 'instance', + 'greptime.semantic.entity.host.id' = 'host', 'greptime.semantic.entity.process.id' = 'service_name,host', 'greptime.semantic.entity.process.descriptive' = 'env' ); -insert into graph_app_latency values (now(), 'cart', 'h1', 'us-east', 3.5); +insert into graph_app_metrics values (now(), 'cart', 'cart-0', 'h1', 'us-east', 3.5); -- SQLNESS PROTOCOL MYSQL select entity_type, entity_id, entity_id_attrs, scope, descriptive, source_tables from greptime_private.semantic_entities order by entity_type, entity_id; -drop table graph_app_latency; +select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence +from greptime_private.semantic_relationships +order by rel_type, src_id; --- Renaming a physical table INTO a reserved computed-table name would let the --- overlay shadow it, orphaning its data. -create table greptime_private.graph_rename_probe (ts timestamp time index); +drop table graph_app_metrics; -alter table greptime_private.graph_rename_probe rename semantic_entities; +-- Calls derivation over trace-v1 tables, all branches in one scan: the client +-- span and its child server span land in different tables and still pair into +-- one edge with RED metrics; the epoch-timestamped pair falls outside the +-- default one-hour window; unmatched clients become virtual-node edges named +-- from span attributes (confidence < 1.0, connection type); a trace-model +-- table that lost the fixed span columns is skipped instead of failing the +-- scan. Trace tables need no entity declaration: service entities are +-- implicit. +create table graph_traces_a ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + "span_attributes.peer.service" string, + "span_attributes.db.name" string, + primary key (service_name) +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); -drop table greptime_private.graph_rename_probe; - --- Calls derivation: a trace-v1-model table with client/server span pairs. --- Trace tables need no entity declaration: the service entities are implicit. -create table graph_traces ( +create table graph_traces_b ( "timestamp" timestamp(9) time index, trace_id string, span_id string, @@ -70,28 +100,41 @@ create table graph_traces ( primary key (service_name) ) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); --- The epoch-timestamped pair falls outside the default one-hour derivation --- window and must not appear in the results. -insert into graph_traces values - (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0), +create table graph_traces_malformed ( + ts timestamp time index, + note string +) with ('table_data_model' = 'greptime_trace_v1'); + +insert into graph_traces_a values + (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0, NULL, NULL), + (now(), 't2', 'c2', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0, NULL, NULL), + (0, 't0', 'c0', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'stale-src', 0, NULL, NULL), + (0, 't0', 's0', 'c0', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'stale-dst', 100, NULL, NULL), + (now(), 't3', 'c3', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 250000000, 'redis', NULL), + (now(), 't4', 'c4', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 100000000, NULL, 'orders-db'); + +insert into graph_traces_b values (now(), 't1', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_ERROR', 'cart', 1500000000), - (now(), 't2', 'c2', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0), - (now(), 't2', 's2', 'c2', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'cart', 500000000), - (0, 't0', 'c0', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'stale-src', 0), - (0, 't0', 's0', 'c0', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'stale-dst', 100); + (now(), 't2', 's2', 'c2', 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'cart', 500000000); + +insert into graph_traces_malformed values (now(), 'not a trace'); -- SQLNESS PROTOCOL MYSQL -select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence, +select src_id, dst_id, rel_type, provenance, confidence, request_count, error_count, duration_sum, duration_count, attributes from greptime_private.semantic_relationships -order by src_id; +order by dst_id; -- SQLNESS PROTOCOL MYSQL select entity_type, entity_id, entity_id_attrs, scope, source_tables from greptime_private.semantic_entities order by entity_id; -drop table graph_traces; +drop table graph_traces_a; + +drop table graph_traces_b; + +drop table graph_traces_malformed; -- Declared edges: the physical table is created by the system on the first -- INSERT with its canonical schema. Re-asserting an edge stores a new revision; @@ -134,18 +177,11 @@ select src_id from greptime_private.semantic_relationships where observed_at < '2001-01-02 00:00:00'; -- The declared-edge table's definition is system-owned: user CREATE/ALTER are --- rejected and a physical table cannot be renamed into the name, while plain --- DML stays allowed. +-- rejected, while plain DML stays allowed. create table greptime_private.semantic_relationships_declared (ts timestamp time index); alter table greptime_private.semantic_relationships_declared add column extra string; -create table greptime_private.declared_rename_probe (ts timestamp time index); - -alter table greptime_private.declared_rename_probe rename semantic_relationships_declared; - -drop table greptime_private.declared_rename_probe; - delete from greptime_private.semantic_relationships_declared; select src_id from greptime_private.semantic_relationships order by src_id; @@ -162,111 +198,9 @@ select src_id from greptime_private.semantic_relationships order by src_id; drop table greptime_private.semantic_relationships_declared; --- Same-row co-declaration: a table declaring both entity types of a built-in --- vocabulary pair witnesses the edge on every row carrying both identities; --- the vocabulary fixes the direction. -create table graph_pod_metrics ( - ts timestamp time index, - instance string, - host string, - "service" string, - cpu double, - primary key (instance, host, "service") -) with ( - 'greptime.semantic.entity.service.instance.id' = 'instance', - 'greptime.semantic.entity.host.id' = 'host', - 'greptime.semantic.entity.service.id' = 'service' -); - -insert into graph_pod_metrics values (now(), 'cart-0', 'node-1', 'cart', 0.5); - -select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence -from greptime_private.semantic_relationships -order by rel_type, src_id; - -drop table graph_pod_metrics; - --- Cross-table calls pairing: the client span and its child server span land in --- different trace tables and must still pair into one edge. A trace-model --- table that lost the fixed span columns is skipped instead of failing the --- scan. -create table graph_traces_a ( - "timestamp" timestamp(9) time index, - trace_id string, - span_id string, - parent_span_id string, - span_kind string, - span_status_code string, - service_name string, - duration_nano bigint unsigned, - primary key (service_name) -) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); - -create table graph_traces_b ( - "timestamp" timestamp(9) time index, - trace_id string, - span_id string, - parent_span_id string, - span_kind string, - span_status_code string, - service_name string, - duration_nano bigint unsigned, - primary key (service_name) -) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); - -create table graph_traces_malformed ( - ts timestamp time index, - note string -) with ('table_data_model' = 'greptime_trace_v1'); - -insert into graph_traces_a values - (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 0); - -insert into graph_traces_b values - (now(), 't1', 's1', 'c1', 'SPAN_KIND_SERVER', 'STATUS_CODE_ERROR', 'cart', 1500000000); - -insert into graph_traces_malformed values (now(), 'not a trace'); - -select src_id, dst_id, rel_type, provenance, confidence, request_count, error_count -from greptime_private.semantic_relationships -order by src_id; - -drop table graph_traces_a; - -drop table graph_traces_b; - -drop table graph_traces_malformed; - --- Virtual nodes: a client span with no matching server span names its peer --- through span attributes; the edge carries confidence < 1.0 and the --- connection type. -create table graph_traces_virtual ( - "timestamp" timestamp(9) time index, - trace_id string, - span_id string, - parent_span_id string, - span_kind string, - span_status_code string, - service_name string, - duration_nano bigint unsigned, - "span_attributes.peer.service" string, - "span_attributes.db.name" string, - primary key (service_name) -) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); - -insert into graph_traces_virtual values - (now(), 't1', 'c1', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 250000000, 'redis', NULL), - (now(), 't2', 'c2', NULL, 'SPAN_KIND_CLIENT', 'STATUS_CODE_UNSET', 'frontend', 100000000, NULL, 'orders-db'); - --- SQLNESS PROTOCOL MYSQL -select src_id, dst_id, rel_type, confidence, request_count, attributes -from greptime_private.semantic_relationships -order by dst_id; - -drop table graph_traces_virtual; - -- Agent edges: span structure derives parent_agent calls agent, and span rows --- co-declaring agent+model / agent+tool witness uses / invokes. +-- co-declaring gen_ai.agent+gen_ai.model / gen_ai.agent+gen_ai.tool witness +-- uses / invokes. The identity columns are fields, not tags. create table graph_agent_traces ( "timestamp" timestamp(9) time index, trace_id string, @@ -279,13 +213,13 @@ create table graph_agent_traces ( agent_id string, model_name string, tool_name string, - primary key (service_name, agent_id, model_name, tool_name) + primary key (service_name) ) with ( 'table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true', - 'greptime.semantic.entity.agent.id' = 'agent_id', - 'greptime.semantic.entity.model.id' = 'model_name', - 'greptime.semantic.entity.tool.id' = 'tool_name' + 'greptime.semantic.entity.gen_ai.agent.id' = 'agent_id', + 'greptime.semantic.entity.gen_ai.model.id' = 'model_name', + 'greptime.semantic.entity.gen_ai.tool.id' = 'tool_name' ); insert into graph_agent_traces values @@ -298,3 +232,129 @@ from greptime_private.semantic_relationships order by rel_type, dst_id; drop table graph_agent_traces; + +-- Prometheus conventions: whitelisted entity-descriptor metrics (stamped +-- signal_type=metric + source=prometheus by the remote-write path) get +-- implicit declarations, and the co-declared rules derive runs_on / part_of / +-- contains from their rows. Pod identity is the UID, so every KSM descriptor +-- lands on one entity; the empty labels kube-state-metrics emits (unscheduled +-- pod's node, ownerless pod's owner_*) identify nothing. A non-whitelisted +-- metric contributes nothing. +create table kube_pod_info ( + greptime_timestamp timestamp(3) time index, + uid string, + namespace string, + pod string, + node string, + greptime_value double, + primary key (uid, namespace, pod, node) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +create table kube_pod_owner ( + greptime_timestamp timestamp(3) time index, + uid string, + namespace string, + pod string, + owner_kind string, + owner_name string, + greptime_value double, + primary key (uid, namespace, pod, owner_kind, owner_name) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +create table kube_pod_container_info ( + greptime_timestamp timestamp(3) time index, + uid string, + container string, + image string, + greptime_value double, + primary key (uid, container, image) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +create table kube_service_info ( + greptime_timestamp timestamp(3) time index, + uid string, + namespace string, + "service" string, + cluster_ip string, + greptime_value double, + primary key (uid, namespace, "service", cluster_ip) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +create table target_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + k8s_cluster_name string, + greptime_value double, + primary key (job, instance, k8s_cluster_name) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +create table http_requests_total ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + greptime_value double, + primary key (job, instance) +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +insert into kube_pod_info values + (now(), 'uid-1', 'default', 'api-1', 'node-a', 1), + (now(), 'uid-2', 'default', 'api-2', 'node-a', 1), + (now(), 'uid-3', 'default', 'pending-1', '', 1), + (now(), '', 'default', 'ghost', 'node-a', 1); + +insert into kube_pod_owner values + (now(), 'uid-1', 'default', 'api-1', 'ReplicaSet', 'api-rs', 1), + (now(), 'uid-2', 'default', 'api-2', 'ReplicaSet', 'api-rs', 1); + +insert into kube_pod_container_info values + (now(), 'uid-1', 'main', 'nginx:1.25', 1); + +insert into kube_service_info values + (now(), 'svc-uid-1', 'default', 'api-svc', '10.0.0.1', 1); + +insert into target_info values + (now(), 'shop/api', 'inst-1', 'prod', 1); + +insert into http_requests_total values + (now(), 'shop/api', 'inst-1', 42); + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +order by entity_type, entity_id, source_tables; + +-- SQLNESS PROTOCOL MYSQL +select src_type, src_id, dst_type, dst_id, rel_type, provenance +from greptime_private.semantic_relationships +order by rel_type, src_id, dst_id; + +drop table kube_pod_info; + +drop table kube_pod_owner; + +drop table kube_pod_container_info; + +drop table kube_service_info; + +drop table target_info; + +drop table http_requests_total; From d0fecdd6b0876b91f2730b70cdb291a6899d5d34 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 13 Aug 2026 06:57:46 +0000 Subject: [PATCH 37/92] fix(mito2): publish committed sequence only after rows are installed (#8862) The committed-sequence watermark must never cover rows that are not yet physically visible. Previously write_memtable() published next_sequence - 1 before bulk parts were installed, so a scan opening a snapshot could bind H to invisible sequences and permanently miss rows after checkpoint advance. Publish once, after both ordinary and bulk memtable writes complete, in the single-region fast path, the multi-region spawned tasks, and WAL replay; skip publication for contexts whose WAL entry could not be built. Add a deterministic worker-level race test using a cfg(test) bulk-install barrier proving the committed sequence stays put until installation. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/mito2/src/region/opener.rs | 3 + src/mito2/src/region_write_ctx.rs | 166 ++++++++++++++++++++++++++- src/mito2/src/worker/handle_write.rs | 164 +++++++++++++++++++++++++- 3 files changed, 326 insertions(+), 7 deletions(-) diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 5d2d36df107..23695ca424c 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -908,6 +908,9 @@ where region_write_ctx.set_next_entry_id(last_entry_id + 1); region_write_ctx.write_memtable().await; region_write_ctx.write_bulk().await; + // Publish the replayed sequences only after all rows (including bulk + // parts) are installed, matching the write path ordering. + region_write_ctx.publish_sequence_and_entry_id(); } // TODO(weny): We need to update `flushed_entry_id` in the region manifest diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index b71e3e5ab17..f64f9810f03 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -199,6 +199,12 @@ impl RegionWriteCtx { &self.version } + /// Returns the version control of the region. + #[cfg(test)] + pub(crate) fn version_control(&self) -> &VersionControlRef { + &self.version_control + } + /// Returns whether writes in this context should skip WAL. pub(crate) fn skip_wal(&self) -> bool { self.provider == Provider::Noop || self.version.options.skip_wal @@ -286,10 +292,6 @@ impl RegionWriteCtx { let bytes = new_memory_usage.saturating_sub(prev_memory_usage.unwrap_or_default()); written_bytes.fetch_add(bytes as u64, Ordering::Relaxed); } - // Updates region sequence and entry id. Since we stores last sequence and entry id in region, we need - // to decrease `next_sequence` and `next_entry_id` by 1. - self.version_control - .set_sequence_and_entry_id(self.next_sequence - 1, self.next_entry_id - 1); } pub(crate) fn push_bulk( @@ -324,6 +326,8 @@ impl RegionWriteCtx { if self.failed || self.bulk_parts.is_empty() { return; } + #[cfg(test)] + test_hooks::pause_before_bulk_install(self.region_id).await; let _timer = metrics::REGION_WORKER_HANDLE_WRITE_ELAPSED .with_label_values(&["write_bulk"]) .start_timer(); @@ -369,11 +373,165 @@ impl RegionWriteCtx { let bytes = new_memory_usage.saturating_sub(prev_memory_usage.unwrap_or_default()); written_bytes.fetch_add(bytes as u64, Ordering::Relaxed); } + } + + /// Publishes the sequences and entry id assigned by this context to the + /// region's committed watermark. + /// + /// Must be called after both [`write_memtable`](Self::write_memtable) and + /// [`write_bulk`](Self::write_bulk) have completed, so the committed + /// sequence never covers rows that are not yet physically installed in the + /// memtable (bulk part sequences are assigned in [`push_bulk`](Self::push_bulk) + /// before installation). Since we store the last sequence and entry id in + /// the region, we decrease `next_sequence` and `next_entry_id` by 1. + /// + /// If the write operation failed (e.g. the WAL entry could not be built), + /// no rows were installed and nothing must be published: advancing the + /// committed watermark here would make it cover rows that were never + /// written. + pub(crate) fn publish_sequence_and_entry_id(&self) { + if self.failed { + return; + } self.version_control .set_sequence_and_entry_id(self.next_sequence - 1, self.next_entry_id - 1); } } +/// Test-only hooks to make write ordering races deterministic. +#[cfg(test)] +pub(crate) mod test_hooks { + use std::sync::Mutex; + use std::sync::atomic::{AtomicU64, Ordering}; + + use store_api::storage::RegionId; + use tokio::sync::watch; + + /// Channels of an armed bulk-install barrier. + /// + /// `reached` signals that a bulk write paused between ordinary-memtable + /// handling and bulk installation; `release` unblocks it. Dropping the + /// senders (by disarming the barrier) also unblocks paused writes because + /// their `wait_for` fails on a closed channel. + struct ActiveBarrier { + id: u64, + /// Only bulk writes for this region pause at the barrier; writes for + /// other regions pass through untouched, so concurrently running tests + /// can never satisfy (or be blocked by) each other's barrier. + target_region_id: RegionId, + reached: watch::Sender, + release: watch::Sender, + } + + /// The currently armed barrier, if any. Wrapped in a `Mutex` so a test can + /// arm a fresh barrier after a previous one was released (or dropped), and + /// so a barrier can never leak past the test that owns it. + static ACTIVE_BARRIER: Mutex> = Mutex::new(None); + static NEXT_BARRIER_ID: AtomicU64 = AtomicU64::new(1); + + fn lock_active_barrier() -> std::sync::MutexGuard<'static, Option> { + // Never let a poisoned mutex (e.g. a panic in another test while + // holding the lock) hang or break unrelated tests. + ACTIVE_BARRIER + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } + + /// RAII guard for the bulk-install barrier. + /// + /// The guard owns the armed barrier. Releasing it — or dropping the guard, + /// e.g. when a test panics — unblocks any write paused at the barrier and + /// disarms it, so a paused write can never hang and later tests can arm a + /// fresh barrier. + pub(crate) struct BulkInstallBarrier { + id: u64, + reached_rx: watch::Receiver, + release_tx: watch::Sender, + released: bool, + } + + impl BulkInstallBarrier { + /// Waits until a bulk write paused at the barrier. + pub(crate) async fn wait_until_reached(&mut self) { + if !*self.reached_rx.borrow() { + let _ = self.reached_rx.wait_for(|reached| *reached).await; + } + } + + /// Unblocks the paused write and disarms the barrier. + pub(crate) fn release(&mut self) { + if self.released { + return; + } + self.released = true; + // Flip the release value so writes paused on this barrier proceed. + let _ = self.release_tx.send(true); + disarm_barrier(self.id); + } + } + + impl Drop for BulkInstallBarrier { + fn drop(&mut self) { + self.release(); + } + } + + /// Arms the bulk-install barrier for `target_region_id` and returns a + /// guard that owns it. + /// + /// Any previously armed barrier is replaced: its senders are dropped, which + /// unblocks writes that were paused on it instead of letting them hang. + pub(crate) fn arm_bulk_install_barrier(target_region_id: RegionId) -> BulkInstallBarrier { + let (reached_tx, reached_rx) = watch::channel(false); + let (release_tx, _release_rx) = watch::channel(false); + let id = NEXT_BARRIER_ID.fetch_add(1, Ordering::Relaxed); + let mut active = lock_active_barrier(); + *active = Some(ActiveBarrier { + id, + target_region_id, + reached: reached_tx, + release: release_tx.clone(), + }); + BulkInstallBarrier { + id, + reached_rx, + release_tx, + released: false, + } + } + + /// Removes the barrier with `id` from the statics if it is still active. + fn disarm_barrier(id: u64) { + let mut active = lock_active_barrier(); + if active.as_ref().is_some_and(|barrier| barrier.id == id) { + *active = None; + } + } + + /// Pauses a bulk write for `region_id` before installing its parts until + /// the test releases the barrier (or it is disarmed). Bulk writes for + /// other regions return immediately. + pub(crate) async fn pause_before_bulk_install(region_id: RegionId) { + let (reached_tx, release_rx) = { + let active = lock_active_barrier(); + match active.as_ref() { + Some(barrier) if barrier.target_region_id == region_id => { + (barrier.reached.clone(), barrier.release.subscribe()) + } + _ => return, + } + }; + // Signal that a bulk write reached the pause point. + let _ = reached_tx.send(true); + let mut release_rx = release_rx; + if !*release_rx.borrow() { + // The sender is dropped when the barrier is disarmed, which makes + // `wait_for` return an error instead of hanging forever. + let _ = release_rx.wait_for(|released| *released).await; + } + } +} + #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 3be40553145..0c2eecdea20 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -128,6 +128,10 @@ impl RegionWorkerLoop { let mut region_ctx = region_ctxs.into_values().next().unwrap(); region_ctx.write_memtable().await; region_ctx.write_bulk().await; + // Publish only after all rows (including bulk parts) are + // physically installed, so a scan opening on the committed + // sequence can never bind H to invisible rows. + region_ctx.publish_sequence_and_entry_id(); put_rows += region_ctx.put_num; delete_rows += region_ctx.delete_num; } else { @@ -138,6 +142,9 @@ impl RegionWorkerLoop { common_runtime::spawn_global(async move { region_ctx.write_memtable().await; region_ctx.write_bulk().await; + // The spawned task owns the moved ctx, so publish + // inside the task after the memtable writes. + region_ctx.publish_sequence_and_entry_id(); (region_ctx.put_num, region_ctx.delete_num) }) }) @@ -734,7 +741,12 @@ fn check_partition_expr_version( #[cfg(test)] mod tests { - use api::v1::{Row, Rows}; + use api::v1::helper::{tag_column_schema, time_index_column_schema}; + use api::v1::value::ValueData; + use api::v1::{ColumnDataType, Row, Rows}; + use common_recordbatch::DfRecordBatch; + use datatypes::arrow::array::{ArrayRef, StringArray, TimestampMillisecondArray}; + use datatypes::arrow::datatypes::{DataType, Field, Schema}; use futures::stream; use log_store::error::{ Error as LogStoreError, IllegalStateSnafu, InvalidProviderSnafu, Result as LogStoreResult, @@ -746,9 +758,42 @@ mod tests { use tokio::sync::oneshot; use super::*; + use crate::memtable::bulk::part::BulkPart; use crate::request::OptionOutputTx; + use crate::test_util::ts_ms_value; use crate::test_util::version_util::VersionControlBuilder; + /// Creates a bulk part with `num_rows` rows. The schema carries the + /// builder metadata's columns (`tag_0` primary key + `ts` time index) so + /// the bulk-install conversion succeeds; `timestamp_index` points at the + /// `ts` column. + fn new_bulk_part(num_rows: i64) -> BulkPart { + let schema = Arc::new(Schema::new(vec![ + Field::new("tag_0", DataType::Utf8, true), + Field::new( + "ts", + DataType::Timestamp(datatypes::arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + ])); + let tag = Arc::new(StringArray::from_iter_values( + (0..num_rows).map(|value| value.to_string()), + )) as ArrayRef; + let ts = Arc::new(TimestampMillisecondArray::from( + (0..num_rows).collect::>(), + )) as ArrayRef; + let batch = DfRecordBatch::try_new(schema, vec![tag, ts]).unwrap(); + + BulkPart { + batch, + max_timestamp: num_rows - 1, + min_timestamp: 0, + sequence: 0, + timestamp_index: 1, + raw_data: None, + } + } + /// A log store that fails to build entries for `failing_region` and fails the /// whole batch when `fail_append` is true. #[derive(Debug, Default)] @@ -857,8 +902,18 @@ mod tests { ctx.push_mutation( OpType::Put as i32, Some(Rows { - schema: vec![], - rows: vec![Row { values: vec![] }], + schema: vec![ + time_index_column_schema("ts", ColumnDataType::TimestampMillisecond), + tag_column_schema("tag_0", ColumnDataType::String), + ], + rows: vec![Row { + values: vec![ + ts_ms_value(0), + api::v1::Value { + value_data: Some(ValueData::StringValue("a".to_string())), + }, + ], + }], }), None, OptionOutputTx::from(tx), @@ -878,8 +933,10 @@ mod tests { let mut region_ctxs = HashMap::new(); let (ctx, failing_rx) = new_region_ctx(failing_region); + let failing_committed_sequence = ctx.version_control().committed_sequence(); region_ctxs.insert(failing_region, ctx); let (ctx, ok_rx) = new_region_ctx(ok_region); + let ok_committed_sequence = ctx.version_control().committed_sequence(); region_ctxs.insert(ok_region, ctx); let entry_id = region_ctxs[&ok_region].next_entry_id(); @@ -890,12 +947,113 @@ mod tests { assert!(!region_ctxs[&ok_region].is_failed()); assert_eq!(entry_id + 1, region_ctxs[&ok_region].next_entry_id()); + // Run the memtable and publication phase the worker runs after `write_wal`. + // The failed region installed no rows, so its committed sequence must not + // advance; the successful sibling region advances normally. + for region_ctx in region_ctxs.values_mut() { + region_ctx.write_memtable().await; + region_ctx.write_bulk().await; + region_ctx.publish_sequence_and_entry_id(); + } + + assert_eq!( + failing_committed_sequence, + region_ctxs[&failing_region] + .version_control() + .committed_sequence() + ); + assert_eq!( + ok_committed_sequence + 1, + region_ctxs[&ok_region] + .version_control() + .committed_sequence() + ); + // Waiters of the failed region get the error while others get the result. drop(region_ctxs); assert!(failing_rx.await.unwrap().is_err()); assert_eq!(1, ok_rx.await.unwrap().unwrap()); } + // The committed sequence must not be published before the bulk part's rows + // are physically installed in the memtable: publication happens strictly + // after `write_bulk` returns. Armed with the bulk-install test barrier, + // this deterministically pauses the worker's memtable phase between the + // ordinary-memtable handling and the bulk installation and asserts the + // region's committed sequence is still its initial value (0) — read via + // the `version_control()` accessor, never through a scanner API. After the + // barrier is released, the committed sequence must advance to cover the + // bulk rows. + #[tokio::test] + async fn test_bulk_write_sequence_not_committed_before_install_worker_level() { + let region_id = RegionId::new(1, 1); + let version_control = Arc::new(VersionControlBuilder::new().build()); + assert_eq!( + 0, + version_control.committed_sequence(), + "the builder must start at sequence 0" + ); + + let mut region_ctxs = HashMap::new(); + let mut ctx = RegionWriteCtx::new( + region_id, + &version_control, + Provider::raft_engine_provider(region_id.as_u64()), + None, + ); + let (tx, rx) = oneshot::channel(); + // Push a 3-row bulk part without an explicit sequence, exactly like the + // worker's `process_bulk_requests`. + assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(3), None)); + region_ctxs.insert(region_id, ctx); + + // Run the WAL phase the worker runs before the memtable phase. + let wal = Wal::new(Arc::new(MockLogStore::default())); + assert!(write_wal(&wal, &mut region_ctxs).await); + assert!(!region_ctxs[®ion_id].is_failed()); + + // Arm the bulk-install barrier: `write_bulk` pauses right before the + // parts are physically installed into the memtable. + let mut barrier = crate::region_write_ctx::test_hooks::arm_bulk_install_barrier(region_id); + + // Run the worker's memtable and publication phases in the background. + let write_handle = tokio::spawn(async move { + let mut region_ctx = region_ctxs.remove(®ion_id).unwrap(); + region_ctx.write_memtable().await; + region_ctx.write_bulk().await; + region_ctx.publish_sequence_and_entry_id(); + }); + + // Wait until the write paused at the barrier: deterministic, no sleeps. + tokio::time::timeout( + std::time::Duration::from_secs(10), + barrier.wait_until_reached(), + ) + .await + .expect("bulk write never reached the install barrier"); + + // The committed sequence must not have advanced yet: publication + // happens strictly after the bulk part is installed. + assert_eq!( + 0, + version_control.committed_sequence(), + "committed sequence leaked before the bulk part was installed" + ); + + // Release the barrier: the bulk part installs and the committed + // sequence advances to cover all 3 bulk rows. + barrier.release(); + write_handle.await.expect("bulk write should complete"); + assert_eq!( + 3, + version_control.committed_sequence(), + "committed sequence must cover the installed bulk rows" + ); + + // The bulk waiter is notified with the number of installed rows. + assert_eq!(3, rx.await.unwrap().unwrap()); + } + #[tokio::test] async fn test_write_wal_all_regions_failed_to_build_entries() { let failing_region = RegionId::new(1, 1); From 1af4c335242e7e5e06ee91bc428a71bb631d0798 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 13 Aug 2026 07:29:00 +0000 Subject: [PATCH 38/92] refactor(event): separate procedure submission context (#8856) * refactor(event): separate procedure submission context Signed-off-by: WenyXu * fix(event): map extensions and forward GC context Signed-off-by: WenyXu * fix(gc): initialize integration test context Signed-off-by: WenyXu * fix(test): pass procedure context to DDL helpers Signed-off-by: WenyXu * refactor: simplify procedure submission contexts Signed-off-by: WenyXu * refactor(event): separate procedure and query contexts Signed-off-by: WenyXu * fix(event): clarify procedure context propagation Signed-off-by: WenyXu * fix(event): preserve procedure submission context Signed-off-by: WenyXu * refactor(event): move DDL context by value Signed-off-by: WenyXu * fix(test): retain manual GC event context Signed-off-by: WenyXu * refactor(event): tighten procedure context API Signed-off-by: WenyXu * chore: update greptime-proto Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- Cargo.lock | 3 +- Cargo.toml | 2 +- src/common/event-recorder/src/context.rs | 119 ++++++ src/common/event-recorder/src/lib.rs | 2 +- src/common/function/src/admin/gc.rs | 52 ++- .../function/src/admin/migrate_region.rs | 17 +- src/common/function/src/admin/purge_table.rs | 14 +- src/common/function/src/handlers.rs | 20 +- src/common/function/src/state.rs | 15 +- src/common/meta/src/ddl/tests/event/flow.rs | 12 +- src/common/meta/src/ddl/tests/event/view.rs | 12 +- src/common/meta/src/ddl_manager.rs | 388 +++++++++++++----- src/common/meta/src/procedure_executor.rs | 16 +- src/common/meta/src/rpc/ddl.rs | 146 +------ src/common/procedure/src/local.rs | 83 +++- src/common/procedure/src/procedure.rs | 12 +- src/common/procedure/src/store.rs | 8 + src/frontend/src/instance.rs | 4 +- src/meta-client/Cargo.toml | 1 + src/meta-client/src/client.rs | 80 ++-- src/meta-client/src/client/procedure.rs | 227 ++++++++-- src/meta-client/src/error.rs | 7 + src/meta-srv/src/gc/ctx.rs | 13 +- src/meta-srv/src/gc/handler.rs | 18 +- src/meta-srv/src/gc/mock.rs | 10 +- src/meta-srv/src/gc/mock/basic.rs | 21 +- src/meta-srv/src/gc/mock/concurrent.rs | 12 +- src/meta-srv/src/gc/mock/full_list.rs | 12 +- src/meta-srv/src/gc/mock/misc.rs | 6 +- src/meta-srv/src/gc/scheduler.rs | 35 +- .../src/procedure/region_migration.rs | 4 +- .../src/procedure/region_migration/manager.rs | 103 ++++- src/meta-srv/src/service/procedure.rs | 195 +++++++-- src/operator/src/procedure.rs | 73 ++-- src/operator/src/statement/comment.rs | 19 +- src/operator/src/statement/ddl.rs | 227 +++++----- src/operator/src/utils.rs | 63 ++- tests-integration/tests/gc_event.rs | 10 +- tests-integration/tests/region_migration.rs | 174 +++++--- tests-integration/tests/repartition.rs | 6 +- tests-integration/tests/table_ddl_event.rs | 16 +- 41 files changed, 1487 insertions(+), 770 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index dcd23451407..97b2fc76a79 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6039,7 +6039,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=e7be20ff855b7522efe235383d82238d56e676c3#e7be20ff855b7522efe235383d82238d56e676c3" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=32f467fa2ba2b3588a58381a24af83de09fbb00a#32f467fa2ba2b3588a58381a24af83de09fbb00a" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", @@ -8312,6 +8312,7 @@ dependencies = [ "async-trait", "common-base", "common-error", + "common-event-recorder", "common-grpc", "common-macro", "common-meta", diff --git a/Cargo.toml b/Cargo.toml index 0e6da7888d3..63dbc76f383 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -159,7 +159,7 @@ fs2 = "0.4" fst = "0.4.7" futures = "0.3" futures-util = "0.3" -greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "e7be20ff855b7522efe235383d82238d56e676c3" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "32f467fa2ba2b3588a58381a24af83de09fbb00a" } hex = "0.4" http = "1" humantime = "2.1" diff --git a/src/common/event-recorder/src/context.rs b/src/common/event-recorder/src/context.rs index 601276967e3..1220c25419f 100644 --- a/src/common/event-recorder/src/context.rs +++ b/src/common/event-recorder/src/context.rs @@ -12,9 +12,51 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; + +use api::v1::meta::ProcedureEventContext as PbProcedureEventContext; use serde::{Deserialize, Serialize}; use strum::{AsRefStr, EnumString}; +/// Event metadata supplied when a procedure is submitted. +/// +/// Protocol is deliberately absent. The trusted submission adapter derives it +/// from the typed query channel before serializing the protobuf request. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct ProcedureEventInput { + pub reason: TriggerReason, + pub extensions: HashMap, +} + +impl ProcedureEventInput { + /// Creates procedure event input with a trigger reason and no extensions. + pub fn new(reason: TriggerReason) -> Self { + Self { + reason, + extensions: Default::default(), + } + } +} + +impl From<&ProcedureEventInput> for PbProcedureEventContext { + fn from(input: &ProcedureEventInput) -> Self { + Self { + reason: input.reason.as_ref().to_string(), + protocol: String::new(), + extensions: input.extensions.clone(), + } + } +} + +impl From for ProcedureEventInput { + fn from(context: PbProcedureEventContext) -> Self { + Self { + reason: TriggerReason::from_extension(&context.reason), + extensions: context.extensions, + } + } +} + /// Stable context recorded for a procedure event. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct PersistentEventContext { @@ -42,12 +84,40 @@ impl PersistentEventContext { } } +impl From<(ProcedureEventInput, Option)> for PersistentEventContext { + fn from((input, protocol): (ProcedureEventInput, Option)) -> Self { + Self { + reason: input.reason, + protocol, + extensions: input + .extensions + .into_iter() + .map(|(key, value)| (key, serde_json::Value::String(value))) + .collect(), + } + } +} + impl Default for PersistentEventContext { fn default() -> Self { Self::new(TriggerReason::default()) } } +impl From for PersistentEventContext { + fn from(context: PbProcedureEventContext) -> Self { + Self { + reason: TriggerReason::from_extension(&context.reason), + protocol: (!context.protocol.is_empty()).then_some(context.protocol), + extensions: context + .extensions + .into_iter() + .map(|(key, value)| (key, serde_json::Value::String(value))) + .collect(), + } + } +} + /// Stable classification of a procedure trigger. #[derive( Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, AsRefStr, EnumString, @@ -96,6 +166,55 @@ mod tests { ); } + #[test] + fn test_event_context_from_protobuf() { + let protobuf = PbProcedureEventContext { + reason: "auto_create".to_string(), + protocol: "postgres".to_string(), + extensions: HashMap::from([ + ("source".to_string(), "sql".to_string()), + ("tenant".to_string(), "a".to_string()), + ]), + }; + assert_eq!( + PersistentEventContext::from(protobuf), + PersistentEventContext { + reason: TriggerReason::AutoCreate, + protocol: Some("postgres".to_string()), + extensions: serde_json::Map::from_iter([ + ("source".to_string(), json!("sql")), + ("tenant".to_string(), json!("a")), + ]), + } + ); + } + + #[test] + fn test_event_input_protobuf_has_no_protocol() { + let input = ProcedureEventInput { + reason: TriggerReason::AutoCreate, + extensions: HashMap::from([("source".to_string(), "sql".to_string())]), + }; + + assert_eq!( + PbProcedureEventContext::from(&input), + PbProcedureEventContext { + reason: "auto_create".to_string(), + protocol: String::new(), + extensions: input.extensions.clone(), + } + ); + + assert_eq!( + ProcedureEventInput::from(PbProcedureEventContext { + reason: "auto_create".to_string(), + protocol: "untrusted".to_string(), + extensions: input.extensions.clone(), + }), + input + ); + } + #[test] fn test_trigger_reason_deserializes_unknown_value() { let reason: TriggerReason = serde_json::from_str("\"future_reason\"").unwrap(); diff --git a/src/common/event-recorder/src/lib.rs b/src/common/event-recorder/src/lib.rs index 5b9cb4a7b97..7ee71722229 100644 --- a/src/common/event-recorder/src/lib.rs +++ b/src/common/event-recorder/src/lib.rs @@ -22,5 +22,5 @@ pub mod recorder; #[cfg(any(test, feature = "testing"))] pub mod testing; -pub use context::{PersistentEventContext, TriggerReason}; +pub use context::{PersistentEventContext, ProcedureEventInput, TriggerReason}; pub use recorder::*; diff --git a/src/common/function/src/admin/gc.rs b/src/common/function/src/admin/gc.rs index e9ba33a050b..0090140a460 100644 --- a/src/common/function/src/admin/gc.rs +++ b/src/common/function/src/admin/gc.rs @@ -38,17 +38,20 @@ const DEFAULT_FULL_FILE_LISTING: bool = false; )] pub(crate) async fn gc_regions( procedure_service_handler: &ProcedureServiceHandlerRef, - _ctx: &QueryContextRef, + query_ctx: &QueryContextRef, params: &[ValueRef<'_>], ) -> Result { let (region_ids, full_file_listing) = parse_gc_regions_params(params)?; let resp = procedure_service_handler - .gc_regions(GcRegionsRequest { - region_ids, - full_file_listing, - timeout: None, - }) + .gc_regions( + query_ctx.clone(), + GcRegionsRequest { + region_ids, + full_file_listing, + timeout: None, + }, + ) .await?; Ok(Value::from(resp.processed_regions)) @@ -69,13 +72,16 @@ pub(crate) async fn gc_table( parse_gc_table_params(params, query_ctx)?; let resp = procedure_service_handler - .gc_table(GcTableRequest { - catalog_name, - schema_name, - table_name, - full_file_listing, - timeout: None, - }) + .gc_table( + query_ctx.clone(), + GcTableRequest { + catalog_name, + schema_name, + table_name, + full_file_listing, + timeout: None, + }, + ) .await?; Ok(Value::from(resp.processed_regions)) @@ -268,13 +274,17 @@ mod tests { impl ProcedureServiceHandler for MockProcedureServiceHandler { async fn purge_table( &self, - _table_name: table::table_name::TableName, _query_ctx: QueryContextRef, + _table_name: table::table_name::TableName, ) -> Result<()> { unreachable!() } - async fn migrate_region(&self, _request: MigrateRegionRequest) -> Result> { + async fn migrate_region( + &self, + _query_ctx: QueryContextRef, + _request: MigrateRegionRequest, + ) -> Result> { unreachable!() } @@ -297,12 +307,20 @@ mod tests { unreachable!() } - async fn gc_regions(&self, request: GcRegionsRequest) -> Result { + async fn gc_regions( + &self, + _query_ctx: QueryContextRef, + request: GcRegionsRequest, + ) -> Result { *self.gc_regions_request.lock().unwrap() = Some(request); Ok(GcResponse::default()) } - async fn gc_table(&self, request: GcTableRequest) -> Result { + async fn gc_table( + &self, + _query_ctx: QueryContextRef, + request: GcTableRequest, + ) -> Result { *self.gc_table_request.lock().unwrap() = Some(request); Ok(GcResponse::default()) } diff --git a/src/common/function/src/admin/migrate_region.rs b/src/common/function/src/admin/migrate_region.rs index d958b3cce95..5d258e14946 100644 --- a/src/common/function/src/admin/migrate_region.rs +++ b/src/common/function/src/admin/migrate_region.rs @@ -47,7 +47,7 @@ const DEFAULT_TIMEOUT_SECS: u64 = 300; )] pub(crate) async fn migrate_region( procedure_service_handler: &ProcedureServiceHandlerRef, - _ctx: &QueryContextRef, + query_ctx: &QueryContextRef, params: &[ValueRef<'_>], ) -> Result { let (region_id, from_peer, to_peer, timeout) = match params.len() { @@ -82,12 +82,15 @@ pub(crate) async fn migrate_region( match (region_id, from_peer, to_peer, timeout) { (Some(region_id), Some(from_peer), Some(to_peer), Some(timeout)) => { let pid = procedure_service_handler - .migrate_region(MigrateRegionRequest { - region_id, - from_peer, - to_peer, - timeout: Duration::from_secs(timeout), - }) + .migrate_region( + query_ctx.clone(), + MigrateRegionRequest { + region_id, + from_peer, + to_peer, + timeout: Duration::from_secs(timeout), + }, + ) .await?; match pid { diff --git a/src/common/function/src/admin/purge_table.rs b/src/common/function/src/admin/purge_table.rs index 59613e155f0..7e5fcd84371 100644 --- a/src/common/function/src/admin/purge_table.rs +++ b/src/common/function/src/admin/purge_table.rs @@ -64,8 +64,8 @@ pub(crate) async fn purge_table( procedure_service_handler .purge_table( - TableName::new(catalog_name, schema_name, table_name), query_ctx.clone(), + TableName::new(catalog_name, schema_name, table_name), ) .await?; Ok(Value::from(0_u64)) @@ -131,8 +131,8 @@ mod tests { impl ProcedureServiceHandler for RecordingHandler { async fn purge_table( &self, - table_name: TableName, query_ctx: QueryContextRef, + table_name: TableName, ) -> Result<()> { if self.fail { return InvalidFuncArgsSnafu { @@ -144,7 +144,11 @@ mod tests { Ok(()) } - async fn migrate_region(&self, _: MigrateRegionRequest) -> Result> { + async fn migrate_region( + &self, + _: QueryContextRef, + _: MigrateRegionRequest, + ) -> Result> { unreachable!() } async fn reconcile(&self, _: ReconcileRequest) -> Result> { @@ -159,10 +163,10 @@ mod tests { fn catalog_manager(&self) -> &CatalogManagerRef { unreachable!() } - async fn gc_regions(&self, _: GcRegionsRequest) -> Result { + async fn gc_regions(&self, _: QueryContextRef, _: GcRegionsRequest) -> Result { unreachable!() } - async fn gc_table(&self, _: GcTableRequest) -> Result { + async fn gc_table(&self, _: QueryContextRef, _: GcTableRequest) -> Result { unreachable!() } } diff --git a/src/common/function/src/handlers.rs b/src/common/function/src/handlers.rs index 24ce115a0d6..6ee8064d6c7 100644 --- a/src/common/function/src/handlers.rs +++ b/src/common/function/src/handlers.rs @@ -89,10 +89,14 @@ pub trait TableMutationHandler: Send + Sync { #[async_trait] pub trait ProcedureServiceHandler: Send + Sync { /// Permanently purge a dropped table. - async fn purge_table(&self, table_name: TableName, query_ctx: QueryContextRef) -> Result<()>; + async fn purge_table(&self, query_ctx: QueryContextRef, table_name: TableName) -> Result<()>; /// Migrate a region from source peer to target peer, returns the procedure id if success. - async fn migrate_region(&self, request: MigrateRegionRequest) -> Result>; + async fn migrate_region( + &self, + query_ctx: QueryContextRef, + request: MigrateRegionRequest, + ) -> Result>; /// Reconcile a table, database or catalog, returns the procedure id if success. async fn reconcile(&self, request: ReconcileRequest) -> Result>; @@ -107,10 +111,18 @@ pub trait ProcedureServiceHandler: Send + Sync { fn catalog_manager(&self) -> &CatalogManagerRef; /// Manually trigger GC for specific regions. - async fn gc_regions(&self, request: MetaGcRegionsRequest) -> Result; + async fn gc_regions( + &self, + query_ctx: QueryContextRef, + request: MetaGcRegionsRequest, + ) -> Result; /// Manually trigger GC for a table. - async fn gc_table(&self, request: MetaGcTableRequest) -> Result; + async fn gc_table( + &self, + query_ctx: QueryContextRef, + request: MetaGcTableRequest, + ) -> Result; } /// This flow service handler is only use for flush flow for now. diff --git a/src/common/function/src/state.rs b/src/common/function/src/state.rs index b1802b62578..2b9f3a517a5 100644 --- a/src/common/function/src/state.rs +++ b/src/common/function/src/state.rs @@ -60,14 +60,15 @@ impl FunctionState { impl ProcedureServiceHandler for MockProcedureServiceHandler { async fn purge_table( &self, - _table_name: table::table_name::TableName, _query_ctx: QueryContextRef, + _table_name: table::table_name::TableName, ) -> Result<()> { Ok(()) } async fn migrate_region( &self, + _ctx: QueryContextRef, _request: MigrateRegionRequest, ) -> Result> { Ok(Some("test_pid".to_string())) @@ -92,7 +93,11 @@ impl FunctionState { Ok(()) } - async fn gc_regions(&self, _request: GcRegionsRequest) -> Result { + async fn gc_regions( + &self, + _context: QueryContextRef, + _request: GcRegionsRequest, + ) -> Result { Ok(GcResponse { processed_regions: 1, need_retry_regions: vec![], @@ -101,7 +106,11 @@ impl FunctionState { }) } - async fn gc_table(&self, _request: GcTableRequest) -> Result { + async fn gc_table( + &self, + _context: QueryContextRef, + _request: GcTableRequest, + ) -> Result { Ok(GcResponse { processed_regions: 1, need_retry_regions: vec![], diff --git a/src/common/meta/src/ddl/tests/event/flow.rs b/src/common/meta/src/ddl/tests/event/flow.rs index aae3613a087..cec0162adf8 100644 --- a/src/common/meta/src/ddl/tests/event/flow.rs +++ b/src/common/meta/src/ddl/tests/event/flow.rs @@ -19,9 +19,8 @@ use api::v1::{ColumnSchema, Row, Value}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_event_recorder::Event; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, - PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, - jsonb_value, + CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, jsonb_value, + procedure_event_column_schemas, }; use common_event_recorder::testing::assert_event_contract; use common_procedure::{EventTrigger, ProcedureEvent, ProcedureId, ProcedureState}; @@ -240,12 +239,7 @@ fn assert_procedure_event_contract( trigger: &str, locator: FlowEventLocator<'_>, ) { - let mut schema = vec![ - PROCEDURE_ID_COLUMN.column_schema(), - PROCEDURE_STATE_COLUMN.column_schema(), - PROCEDURE_ERROR_COLUMN.column_schema(), - PROCEDURE_TRIGGER_COLUMN.column_schema(), - ]; + let mut schema = procedure_event_column_schemas(); schema.extend(flow_schema()); schema.push(EVENT_CONTEXT_COLUMN.column_schema()); assert_event_contract( diff --git a/src/common/meta/src/ddl/tests/event/view.rs b/src/common/meta/src/ddl/tests/event/view.rs index 5cb5a6bb7d9..6ebb62cda82 100644 --- a/src/common/meta/src/ddl/tests/event/view.rs +++ b/src/common/meta/src/ddl/tests/event/view.rs @@ -17,9 +17,8 @@ use std::sync::Arc; use api::v1::value::ValueData; use api::v1::{ColumnSchema, Row, Value}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, - PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, - VIEW_NAME_COLUMN, jsonb_value, + CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, + VIEW_NAME_COLUMN, jsonb_value, procedure_event_column_schemas, }; use common_event_recorder::testing::assert_event_contract; use common_event_recorder::{Event, EventTypeFilter}; @@ -334,12 +333,7 @@ fn assert_procedure_event_contract( trigger: &str, locator: ViewEventLocator<'_>, ) { - let mut schema = vec![ - PROCEDURE_ID_COLUMN.column_schema(), - PROCEDURE_STATE_COLUMN.column_schema(), - PROCEDURE_ERROR_COLUMN.column_schema(), - PROCEDURE_TRIGGER_COLUMN.column_schema(), - ]; + let mut schema = procedure_event_column_schemas(); schema.extend(view_schema()); schema.push(EVENT_CONTEXT_COLUMN.column_schema()); diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index 2eed4628ac3..7ba408ca2d9 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -19,9 +19,12 @@ use api::v1::Repartition; use api::v1::alter_table_expr::Kind; use api::v1::repartition::Source as PbRepartitionSource; use common_error::ext::BoxedError; +use common_event_recorder::PersistentEventContext; +#[cfg(feature = "enterprise")] +use common_event_recorder::TriggerReason; use common_procedure::{ - BoxedProcedure, BoxedProcedureLoader, Output, ProcedureId, ProcedureManagerRef, - ProcedureWithId, watcher, + BoxedProcedure, BoxedProcedureLoader, Output, ProcedureContext, ProcedureId, + ProcedureManagerRef, ProcedureWithId, watcher, }; use common_telemetry::tracing_context::{FutureExt, TracingContext}; use common_telemetry::{debug, info, tracing}; @@ -54,7 +57,7 @@ use crate::error::{ PersistRepartitionGcRequirementSnafu, ProcedureOutputSnafu, RegisterProcedureLoaderSnafu, RegisterRepartitionProcedureLoaderSnafu, Result, SubmitProcedureSnafu, TableInfoNotFoundSnafu, TableNotFoundSnafu, TableRouteNotFoundSnafu, UnexpectedLogicalRouteTableSnafu, - WaitProcedureSnafu, + UnsupportedSnafu, WaitProcedureSnafu, }; use crate::key::table_info::TableInfoValue; use crate::key::table_name::TableNameKey; @@ -70,16 +73,14 @@ use crate::rpc::ddl::DdlTask::{ DropTable, DropView, PurgeDroppedTable, TruncateTable, UndropTable, }; #[cfg(feature = "enterprise")] -use crate::rpc::ddl::TriggerReason; -#[cfg(feature = "enterprise")] use crate::rpc::ddl::trigger::CreateTriggerTask; #[cfg(feature = "enterprise")] use crate::rpc::ddl::trigger::DropTriggerTask; use crate::rpc::ddl::{ AlterDatabaseTask, AlterTableTask, CommentOnTask, CreateDatabaseTask, CreateFlowTask, CreateTableTask, CreateViewTask, DropDatabaseTask, DropFlowTask, DropTableTask, DropViewTask, - PersistentEventContext, PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest, - SubmitDdlTaskResponse, TruncateTableTask, UndropTableTask, + PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest, SubmitDdlTaskResponse, + TruncateTableTask, UndropTableTask, }; const MAX_REGION_ROUTE_CHANGE_RETRIES: usize = 3; @@ -122,6 +123,7 @@ pub trait TriggerDdlManager: Send + Sync { procedure_manager: ProcedureManagerRef, ddl_context: DdlContext, query_context: QueryContext, + procedure_context: ProcedureContext, ) -> Result; async fn drop_trigger( @@ -130,6 +132,7 @@ pub trait TriggerDdlManager: Send + Sync { procedure_manager: ProcedureManagerRef, ddl_context: DdlContext, query_context: QueryContext, + procedure_context: ProcedureContext, ) -> Result; fn as_any(&self) -> &dyn std::any::Any; @@ -341,7 +344,7 @@ impl DdlManager { repartition: Repartition, wait: bool, timeout: Duration, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); @@ -384,7 +387,7 @@ impl DdlManager { .await .context(PersistRepartitionGcRequirementSnafu)?; let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); if wait { self.execute_procedure_and_wait(procedure_with_id).await } else { @@ -400,7 +403,7 @@ impl DdlManager { &self, table_id: TableId, alter_table_task: AlterTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ddl_options: DdlOptions, ) -> Result<(ProcedureId, Option)> { // make alter_table_task mutable so we can call .take() on its field @@ -421,7 +424,7 @@ impl DdlManager { repartition, ddl_options.wait, ddl_options.timeout, - event_context, + procedure_context, ) .await; } @@ -460,7 +463,7 @@ impl DdlManager { )?; let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) - .with_event_context(event_context.clone()); + .with_context(procedure_context.clone()); let result = self.execute_procedure_and_wait(procedure_with_id).await?; if result .1 @@ -489,7 +492,7 @@ impl DdlManager { &self, create_table_task: CreateTableTask, query_context: QueryContext, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); @@ -500,7 +503,7 @@ impl DdlManager { )?; let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -510,14 +513,14 @@ impl DdlManager { pub async fn submit_create_view_task( &self, create_view_task: CreateViewTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = CreateViewProcedure::new(create_view_task, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -528,7 +531,7 @@ impl DdlManager { &self, create_table_tasks: Vec, physical_table_id: TableId, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); @@ -536,7 +539,7 @@ impl DdlManager { CreateLogicalTablesProcedure::new(create_table_tasks, physical_table_id, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -547,7 +550,7 @@ impl DdlManager { &self, alter_table_tasks: Vec, physical_table_id: TableId, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); @@ -555,7 +558,7 @@ impl DdlManager { AlterLogicalTablesProcedure::new(alter_table_tasks, physical_table_id, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -565,14 +568,14 @@ impl DdlManager { pub async fn submit_drop_table_task( &self, drop_table_task: DropTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = DropTableProcedure::new(drop_table_task, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -583,7 +586,7 @@ impl DdlManager { pub async fn submit_undrop_table_task( &self, undrop_table_task: UndropTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { #[cfg(not(feature = "enterprise"))] { @@ -611,7 +614,7 @@ impl DdlManager { Some(original_table_name), ); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) - .with_event_context(event_context); + .with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -623,7 +626,7 @@ impl DdlManager { pub async fn submit_purge_dropped_table_task( &self, purge_dropped_table_task: PurgeDroppedTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { #[cfg(not(feature = "enterprise"))] { @@ -639,7 +642,7 @@ impl DdlManager { let context = self.create_context(); let procedure = PurgeDroppedTableProcedure::new(purge_dropped_table_task, context); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) - .with_event_context(event_context); + .with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -665,11 +668,13 @@ impl DdlManager { #[cfg(feature = "enterprise")] { let context = self.create_context(); - let event_context = PersistentEventContext::new(TriggerReason::ScheduledGc); + let procedure_context = ProcedureContext::from_event_context( + PersistentEventContext::new(TriggerReason::ScheduledGc), + ); let procedure = PurgeDroppedTableProcedure::new_if_expired(purge_dropped_table_task, context); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)) - .with_event_context(event_context); + .with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -686,7 +691,7 @@ impl DdlManager { options, creator, }: CreateDatabaseTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = CreateDatabaseProcedure::new( @@ -698,7 +703,7 @@ impl DdlManager { context, ); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -712,12 +717,12 @@ impl DdlManager { schema, drop_if_exists, }: DropDatabaseTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = DropDatabaseProcedure::new(catalog, schema, drop_if_exists, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -725,12 +730,12 @@ impl DdlManager { pub async fn submit_alter_database( &self, alter_database_task: AlterDatabaseTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = AlterDatabaseProcedure::new(alter_database_task, context)?; let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -741,12 +746,12 @@ impl DdlManager { &self, create_flow: CreateFlowTask, query_context: QueryContext, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = CreateFlowProcedure::new(create_flow, query_context, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -756,12 +761,12 @@ impl DdlManager { pub async fn submit_drop_flow_task( &self, drop_flow: DropFlowTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = DropFlowProcedure::new(drop_flow, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -771,12 +776,12 @@ impl DdlManager { pub async fn submit_drop_view_task( &self, drop_view: DropViewTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = DropViewProcedure::new(drop_view, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -787,13 +792,13 @@ impl DdlManager { &self, truncate_table_task: TruncateTableTask, table_info_value: DeserializedValueWithBytes, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); let procedure = TruncateTableProcedure::new(truncate_table_task, table_info_value, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -803,7 +808,7 @@ impl DdlManager { pub async fn submit_comment_on_task( &self, mut comment_on_task: CommentOnTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); comment_on_task @@ -814,7 +819,7 @@ impl DdlManager { .await?; let procedure = CommentOnProcedure::new(comment_on_task, context); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); self.execute_procedure_and_wait(procedure_with_id).await } @@ -853,89 +858,125 @@ impl DdlManager { pub async fn submit_ddl_task( &self, - ctx: &ExecutorContext, + context: ExecutorContext, request: SubmitDdlTaskRequest, ) -> Result { - let span = ctx - .tracing_context + let ExecutorContext { + tracing_context, + query_context, + actor, + event_input, + } = context; + let query_context = query_context.context(UnsupportedSnafu { + operation: "submit_ddl_task without query context", + })?; + let procedure_context = ProcedureContext { + actor, + event_context: event_input + .map(|input| PersistentEventContext::from((input, query_context.protocol()))), + }; + let span = tracing_context .as_ref() .map(TracingContext::from_w3c) .unwrap_or_else(TracingContext::from_current_span) .attach(tracing::info_span!("DdlManager::submit_ddl_task")); let SubmitDdlTaskRequest { - query_context, wait, timeout, task, } = request; - let event_context = ctx - .event_context - .clone() - .unwrap_or_else(|| crate::rpc::ddl::event_context_from_query_context(&query_context)); let ddl_options = DdlOptions { wait, timeout }; async move { debug!("Submitting Ddl task: {:?}", task); match task { CreateTable(create_table_task) => { - handle_create_table_task(self, create_table_task, query_context, event_context) - .await + handle_create_table_task( + self, + create_table_task, + query_context, + procedure_context, + ) + .await } DropTable(drop_table_task) => { - handle_drop_table_task(self, drop_table_task, event_context).await + handle_drop_table_task(self, drop_table_task, procedure_context).await } UndropTable(undrop_table_task) => { - handle_undrop_table_task(self, undrop_table_task, event_context).await + handle_undrop_table_task(self, undrop_table_task, procedure_context).await } PurgeDroppedTable(purge_dropped_table_task) => { - handle_purge_dropped_table_task(self, purge_dropped_table_task, event_context) - .await + handle_purge_dropped_table_task( + self, + purge_dropped_table_task, + procedure_context, + ) + .await } AlterTable(alter_table_task) => { - handle_alter_table_task(self, alter_table_task, ddl_options, event_context) + handle_alter_table_task(self, alter_table_task, ddl_options, procedure_context) .await } TruncateTable(truncate_table_task) => { - handle_truncate_table_task(self, truncate_table_task, event_context).await + handle_truncate_table_task(self, truncate_table_task, procedure_context).await } CreateLogicalTables(create_table_tasks) => { - handle_create_logical_table_tasks(self, create_table_tasks, event_context).await + handle_create_logical_table_tasks(self, create_table_tasks, procedure_context) + .await } AlterLogicalTables(alter_table_tasks) => { - handle_alter_logical_table_tasks(self, alter_table_tasks, event_context).await + handle_alter_logical_table_tasks(self, alter_table_tasks, procedure_context) + .await } DropLogicalTables(_) => todo!(), CreateDatabase(create_database_task) => { - handle_create_database_task(self, create_database_task, event_context).await + handle_create_database_task(self, create_database_task, procedure_context).await } DropDatabase(drop_database_task) => { - handle_drop_database_task(self, drop_database_task, event_context).await + handle_drop_database_task(self, drop_database_task, procedure_context).await } AlterDatabase(alter_database_task) => { - handle_alter_database_task(self, alter_database_task, event_context).await + handle_alter_database_task(self, alter_database_task, procedure_context).await } CreateFlow(create_flow_task) => { - handle_create_flow_task(self, create_flow_task, query_context, event_context) - .await + handle_create_flow_task( + self, + create_flow_task, + query_context, + procedure_context, + ) + .await } DropFlow(drop_flow_task) => { - handle_drop_flow_task(self, drop_flow_task, event_context).await + handle_drop_flow_task(self, drop_flow_task, procedure_context).await } CreateView(create_view_task) => { - handle_create_view_task(self, create_view_task, event_context).await + handle_create_view_task(self, create_view_task, procedure_context).await } DropView(drop_view_task) => { - handle_drop_view_task(self, drop_view_task, event_context).await + handle_drop_view_task(self, drop_view_task, procedure_context).await } CommentOn(comment_on_task) => { - handle_comment_on_task(self, comment_on_task, event_context).await + handle_comment_on_task(self, comment_on_task, procedure_context).await } #[cfg(feature = "enterprise")] CreateTrigger(create_trigger_task) => { - handle_create_trigger_task(self, create_trigger_task, query_context).await + handle_create_trigger_task( + self, + create_trigger_task, + query_context, + procedure_context, + ) + .await } #[cfg(feature = "enterprise")] DropTrigger(drop_trigger_task) => { - handle_drop_trigger_task(self, drop_trigger_task, query_context).await + handle_drop_trigger_task( + self, + drop_trigger_task, + query_context, + procedure_context, + ) + .await } } } @@ -947,7 +988,7 @@ impl DdlManager { async fn handle_truncate_table_task( ddl_manager: &DdlManager, truncate_table_task: TruncateTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let table_id = truncate_table_task.table_id; let table_metadata_manager = &ddl_manager.table_metadata_manager(); @@ -972,7 +1013,7 @@ async fn handle_truncate_table_task( ); let (id, _) = ddl_manager - .submit_truncate_table_task(truncate_table_task, table_info_value, event_context) + .submit_truncate_table_task(truncate_table_task, table_info_value, procedure_context) .await?; info!("Table: {table_id} is truncated via procedure_id {id:?}"); @@ -987,7 +1028,7 @@ async fn handle_alter_table_task( ddl_manager: &DdlManager, alter_table_task: AlterTableTask, ddl_options: DdlOptions, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let table_ref = alter_table_task.table_ref(); @@ -1020,7 +1061,7 @@ async fn handle_alter_table_task( ); let (id, _) = ddl_manager - .submit_alter_table_task(table_id, alter_table_task, event_context, ddl_options) + .submit_alter_table_task(table_id, alter_table_task, procedure_context, ddl_options) .await?; info!("Table: {table_id} is altered via procedure_id {id:?}"); @@ -1034,11 +1075,11 @@ async fn handle_alter_table_task( async fn handle_drop_table_task( ddl_manager: &DdlManager, drop_table_task: DropTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let table_id = drop_table_task.table_id; let (id, _) = ddl_manager - .submit_drop_table_task(drop_table_task, event_context) + .submit_drop_table_task(drop_table_task, procedure_context) .await?; info!("Table: {table_id} is dropped via procedure_id {id:?}"); @@ -1052,11 +1093,11 @@ async fn handle_drop_table_task( async fn handle_undrop_table_task( ddl_manager: &DdlManager, undrop_table_task: UndropTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let table_id = undrop_table_task.table_id; let (id, _) = ddl_manager - .submit_undrop_table_task(undrop_table_task, event_context) + .submit_undrop_table_task(undrop_table_task, procedure_context) .await?; info!("Table: {table_id} is undropped via procedure_id {id:?}"); @@ -1070,10 +1111,10 @@ async fn handle_undrop_table_task( async fn handle_purge_dropped_table_task( ddl_manager: &DdlManager, purge_dropped_table_task: PurgeDroppedTableTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, _) = ddl_manager - .submit_purge_dropped_table_task(purge_dropped_table_task, event_context) + .submit_purge_dropped_table_task(purge_dropped_table_task, procedure_context) .await?; info!("Dropped table is purged via procedure_id {id:?}"); @@ -1088,10 +1129,10 @@ async fn handle_create_table_task( ddl_manager: &DdlManager, create_table_task: CreateTableTask, query_context: QueryContext, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, output) = ddl_manager - .submit_create_table_task(create_table_task, query_context, event_context) + .submit_create_table_task(create_table_task, query_context, procedure_context) .await?; let procedure_id = id.to_string(); @@ -1114,7 +1155,7 @@ async fn handle_create_table_task( async fn handle_create_logical_table_tasks( ddl_manager: &DdlManager, create_table_tasks: Vec, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { ensure!( !create_table_tasks.is_empty(), @@ -1130,7 +1171,7 @@ async fn handle_create_logical_table_tasks( let num_logical_tables = create_table_tasks.len(); let (id, output) = ddl_manager - .submit_create_logical_table_tasks(create_table_tasks, physical_table_id, event_context) + .submit_create_logical_table_tasks(create_table_tasks, physical_table_id, procedure_context) .await?; info!( @@ -1159,12 +1200,12 @@ async fn handle_create_logical_table_tasks( async fn handle_create_database_task( ddl_manager: &DdlManager, create_database_task: CreateDatabaseTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let catalog = create_database_task.catalog.clone(); let schema = create_database_task.schema.clone(); let (id, _) = ddl_manager - .submit_create_database(create_database_task, event_context) + .submit_create_database(create_database_task, procedure_context) .await?; let procedure_id = id.to_string(); @@ -1182,10 +1223,10 @@ async fn handle_create_database_task( async fn handle_drop_database_task( ddl_manager: &DdlManager, drop_database_task: DropDatabaseTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, _) = ddl_manager - .submit_drop_database(drop_database_task.clone(), event_context) + .submit_drop_database(drop_database_task.clone(), procedure_context) .await?; let procedure_id = id.to_string(); @@ -1203,10 +1244,10 @@ async fn handle_drop_database_task( async fn handle_alter_database_task( ddl_manager: &DdlManager, alter_database_task: AlterDatabaseTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, _) = ddl_manager - .submit_alter_database(alter_database_task.clone(), event_context) + .submit_alter_database(alter_database_task.clone(), procedure_context) .await?; let procedure_id = id.to_string(); @@ -1225,10 +1266,10 @@ async fn handle_alter_database_task( async fn handle_drop_flow_task( ddl_manager: &DdlManager, drop_flow_task: DropFlowTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, _) = ddl_manager - .submit_drop_flow_task(drop_flow_task.clone(), event_context) + .submit_drop_flow_task(drop_flow_task.clone(), procedure_context) .await?; let procedure_id = id.to_string(); @@ -1248,6 +1289,7 @@ async fn handle_drop_trigger_task( ddl_manager: &DdlManager, drop_trigger_task: DropTriggerTask, query_context: QueryContext, + procedure_context: ProcedureContext, ) -> Result { let Some(m) = ddl_manager.trigger_ddl_manager.as_ref() else { use crate::error::UnsupportedSnafu; @@ -1263,6 +1305,7 @@ async fn handle_drop_trigger_task( ddl_manager.procedure_manager.clone(), ddl_manager.ddl_context.clone(), query_context, + procedure_context, ) .await } @@ -1270,10 +1313,10 @@ async fn handle_drop_trigger_task( async fn handle_drop_view_task( ddl_manager: &DdlManager, drop_view_task: DropViewTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, _) = ddl_manager - .submit_drop_view_task(drop_view_task.clone(), event_context) + .submit_drop_view_task(drop_view_task.clone(), procedure_context) .await?; let procedure_id = id.to_string(); @@ -1293,10 +1336,10 @@ async fn handle_create_flow_task( ddl_manager: &DdlManager, create_flow_task: CreateFlowTask, query_context: QueryContext, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, output) = ddl_manager - .submit_create_flow_task(create_flow_task.clone(), query_context, event_context) + .submit_create_flow_task(create_flow_task.clone(), query_context, procedure_context) .await?; let procedure_id = id.to_string(); @@ -1331,6 +1374,7 @@ async fn handle_create_trigger_task( ddl_manager: &DdlManager, create_trigger_task: CreateTriggerTask, query_context: QueryContext, + procedure_context: ProcedureContext, ) -> Result { let Some(m) = ddl_manager.trigger_ddl_manager.as_ref() else { use crate::error::UnsupportedSnafu; @@ -1346,6 +1390,7 @@ async fn handle_create_trigger_task( ddl_manager.procedure_manager.clone(), ddl_manager.ddl_context.clone(), query_context, + procedure_context, ) .await } @@ -1353,7 +1398,7 @@ async fn handle_create_trigger_task( async fn handle_alter_logical_table_tasks( ddl_manager: &DdlManager, alter_table_tasks: Vec, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { ensure!( !alter_table_tasks.is_empty(), @@ -1373,7 +1418,7 @@ async fn handle_alter_logical_table_tasks( let num_logical_tables = alter_table_tasks.len(); let (id, _) = ddl_manager - .submit_alter_logical_table_tasks(alter_table_tasks, physical_table_id, event_context) + .submit_alter_logical_table_tasks(alter_table_tasks, physical_table_id, procedure_context) .await?; info!( @@ -1392,10 +1437,10 @@ async fn handle_alter_logical_table_tasks( async fn handle_create_view_task( ddl_manager: &DdlManager, create_view_task: CreateViewTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, output) = ddl_manager - .submit_create_view_task(create_view_task, event_context) + .submit_create_view_task(create_view_task, procedure_context) .await?; let procedure_id = id.to_string(); @@ -1418,10 +1463,10 @@ async fn handle_create_view_task( async fn handle_comment_on_task( ddl_manager: &DdlManager, comment_on_task: CommentOnTask, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let (id, _) = ddl_manager - .submit_comment_on_task(comment_on_task.clone(), event_context) + .submit_comment_on_task(comment_on_task.clone(), procedure_context) .await?; let procedure_id = id.to_string(); @@ -1439,16 +1484,22 @@ async fn handle_comment_on_task( #[cfg(test)] mod tests { use std::sync::Arc; + #[cfg(feature = "enterprise")] + use std::sync::Mutex; use std::time::Duration; + #[cfg(feature = "enterprise")] + use common_base::protocol::Channel; use common_error::ext::BoxedError; #[cfg(feature = "enterprise")] use common_error::ext::ErrorExt; #[cfg(feature = "enterprise")] use common_error::status_code::StatusCode; + #[cfg(feature = "enterprise")] + use common_event_recorder::{PersistentEventContext, ProcedureEventInput, TriggerReason}; use common_procedure::local::LocalManager; use common_procedure::test_util::InMemoryPoisonStore; - use common_procedure::{BoxedProcedure, ProcedureManagerRef}; + use common_procedure::{BoxedProcedure, ProcedureContext, ProcedureManagerRef}; use store_api::storage::TableId; use table::table_name::TableName; @@ -1470,13 +1521,16 @@ mod tests { use crate::kv_backend::memory::MemoryKvBackend; use crate::node_manager::{DatanodeManager, DatanodeRef, FlownodeManager, FlownodeRef}; use crate::peer::Peer; - #[cfg(not(feature = "enterprise"))] use crate::procedure_executor::ExecutorContext; use crate::region_keeper::MemoryRegionKeeper; use crate::region_registry::LeaderRegionRegistry; - use crate::rpc::ddl::{CreatorGrantIntent, PersistentEventContext, UndropTableTask}; + #[cfg(feature = "enterprise")] + use crate::rpc::ddl::trigger::{CreateTriggerTask, DropTriggerTask}; + use crate::rpc::ddl::{CreatorGrantIntent, UndropTableTask}; #[cfg(not(feature = "enterprise"))] use crate::rpc::ddl::{DdlTask, PurgeDroppedTableTask, QueryContext, SubmitDdlTaskRequest}; + #[cfg(feature = "enterprise")] + use crate::rpc::ddl::{DdlTask, QueryContext, SubmitDdlTaskRequest}; use crate::sequence::SequenceBuilder; use crate::state_store::KvStateStore; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -1543,6 +1597,50 @@ mod tests { } } + #[cfg(feature = "enterprise")] + #[derive(Default)] + struct RecordingTriggerDdlManager { + procedure_contexts: Mutex>, + } + + #[cfg(feature = "enterprise")] + #[async_trait::async_trait] + impl super::TriggerDdlManager for RecordingTriggerDdlManager { + async fn create_trigger( + &self, + _create_trigger_task: CreateTriggerTask, + _procedure_manager: ProcedureManagerRef, + _ddl_context: DdlContext, + _query_context: QueryContext, + procedure_context: ProcedureContext, + ) -> crate::error::Result { + self.procedure_contexts + .lock() + .unwrap() + .push(procedure_context); + Ok(Default::default()) + } + + async fn drop_trigger( + &self, + _drop_trigger_task: DropTriggerTask, + _procedure_manager: ProcedureManagerRef, + _ddl_context: DdlContext, + _query_context: QueryContext, + procedure_context: ProcedureContext, + ) -> crate::error::Result { + self.procedure_contexts + .lock() + .unwrap() + .push(procedure_context); + Ok(Default::default()) + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + #[test] fn test_generic_loader_captures_configured_committer() { let mut context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); @@ -1678,6 +1776,67 @@ mod tests { ) } + #[cfg(feature = "enterprise")] + #[tokio::test] + async fn test_trigger_ddl_forwards_procedure_context() { + let trigger_ddl_manager = Arc::new(RecordingTriggerDdlManager::default()); + let ddl_manager = build_soft_drop_test_ddl_manager() + .with_trigger_ddl_manager(trigger_ddl_manager.clone()); + let procedure_context = ProcedureContext { + actor: Some("test-user".to_string()), + event_context: Some( + PersistentEventContext::new(TriggerReason::Manual).with_protocol("mysql"), + ), + }; + let executor_context = || ExecutorContext { + query_context: Some(QueryContext { + channel: Channel::Mysql as u8, + ..Default::default() + }), + actor: Some("test-user".to_string()), + event_input: Some(ProcedureEventInput::new(TriggerReason::Manual)), + ..Default::default() + }; + + ddl_manager + .submit_ddl_task( + executor_context(), + SubmitDdlTaskRequest::new(DdlTask::CreateTrigger(CreateTriggerTask { + catalog_name: "greptime".to_string(), + trigger_name: "test_trigger".to_string(), + if_not_exists: false, + sql: "SELECT 1".to_string(), + channels: vec![], + labels: Default::default(), + annotations: Default::default(), + interval: Duration::from_secs(1), + raw_interval_expr: None, + r#for: None, + for_raw_expr: None, + keep_firing_for: None, + keep_firing_for_raw_expr: None, + })), + ) + .await + .unwrap(); + ddl_manager + .submit_ddl_task( + executor_context(), + SubmitDdlTaskRequest::new(DdlTask::DropTrigger(DropTriggerTask { + catalog_name: "greptime".to_string(), + trigger_name: "test_trigger".to_string(), + drop_if_exists: false, + })), + ) + .await + .unwrap(); + + assert_eq!( + *trigger_ddl_manager.procedure_contexts.lock().unwrap(), + vec![procedure_context.clone(), procedure_context] + ); + } + #[cfg(feature = "enterprise")] #[tokio::test] async fn test_submit_undrop_missing_tombstone_returns_table_not_found_directly() { @@ -1686,7 +1845,7 @@ mod tests { let err = ddl_manager .submit_undrop_table_task( UndropTableTask { table_id: 1024 }, - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap_err(); @@ -1703,7 +1862,7 @@ mod tests { let err = ddl_manager .submit_undrop_table_task( UndropTableTask { table_id: 1024 }, - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap_err(); @@ -1712,7 +1871,7 @@ mod tests { let err = ddl_manager .submit_purge_dropped_table_task( PurgeDroppedTableTask { table_id: 1024 }, - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap_err(); @@ -1730,8 +1889,11 @@ mod tests { ] { let err = ddl_manager .submit_ddl_task( - &ExecutorContext::default(), - SubmitDdlTaskRequest::new(QueryContext::default(), task), + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + SubmitDdlTaskRequest::new(task), ) .await .unwrap_err(); diff --git a/src/common/meta/src/procedure_executor.rs b/src/common/meta/src/procedure_executor.rs index b812aeb6aac..a873c26b39a 100644 --- a/src/common/meta/src/procedure_executor.rs +++ b/src/common/meta/src/procedure_executor.rs @@ -15,7 +15,7 @@ use std::sync::Arc; use api::v1::meta::{ProcedureDetailResponse, ReconcileRequest, ReconcileResponse}; -use common_event_recorder::PersistentEventContext; +use common_event_recorder::ProcedureEventInput; use common_procedure::{ProcedureId, ProcedureManagerRef}; use common_telemetry::tracing_context::W3cTrace; use snafu::{OptionExt, ResultExt}; @@ -24,7 +24,7 @@ use crate::ddl_manager::DdlManagerRef; use crate::error::{ ParseProcedureIdSnafu, ProcedureNotFoundSnafu, QueryProcedureSnafu, Result, UnsupportedSnafu, }; -use crate::rpc::ddl::{SubmitDdlTaskRequest, SubmitDdlTaskResponse}; +use crate::rpc::ddl::{QueryContext, SubmitDdlTaskRequest, SubmitDdlTaskResponse}; use crate::rpc::procedure::{ self, GcRegionsRequest, GcResponse, GcTableRequest, ManageRegionFollowerRequest, MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse, @@ -34,8 +34,14 @@ use crate::rpc::procedure::{ #[derive(Debug, Clone, Default)] pub struct ExecutorContext { pub tracing_context: Option, + /// Query execution data available at the frontend/standalone submission boundary. + /// + /// Only DDL serializes the full context. Migration and GC use its typed + /// channel to derive protocol without sending the query context itself. + pub query_context: Option, pub actor: Option, - pub event_context: Option, + /// Caller-supplied event metadata. Protocol is derived by the submission adapter. + pub event_input: Option, } /// The procedure executor that accepts ddl, region migration task etc. @@ -44,7 +50,7 @@ pub trait ProcedureExecutor: Send + Sync { /// Submit a ddl task async fn submit_ddl_task( &self, - ctx: &ExecutorContext, + ctx: ExecutorContext, request: SubmitDdlTaskRequest, ) -> Result; @@ -129,7 +135,7 @@ impl LocalProcedureExecutor { impl ProcedureExecutor for LocalProcedureExecutor { async fn submit_ddl_task( &self, - ctx: &ExecutorContext, + ctx: ExecutorContext, request: SubmitDdlTaskRequest, ) -> Result { self.ddl_manager.submit_ddl_task(ctx, request).await diff --git a/src/common/meta/src/rpc/ddl.rs b/src/common/meta/src/rpc/ddl.rs index d7966804738..47bd537b9ba 100644 --- a/src/common/meta/src/rpc/ddl.rs +++ b/src/common/meta/src/rpc/ddl.rs @@ -45,7 +45,7 @@ use base64::engine::general_purpose; use common_base::protocol::Channel; use common_catalog::{format_full_flow_name, format_full_table_name}; use common_error::ext::BoxedError; -pub use common_event_recorder::{PersistentEventContext, TriggerReason}; +pub use common_event_recorder::TriggerReason; use common_time::{DatabaseTimeToLive, Timestamp}; use prost::Message; use serde::{Deserialize, Serialize}; @@ -68,8 +68,6 @@ use crate::key::table_name::{TableNameKey, TableNameManager}; /// Reserved query-context extension key for the frontend peer address that submitted a DDL request. pub const ORIGIN_FRONTEND_ADDR_EXTENSION_KEY: &str = "__greptime_origin_frontend.addr"; -/// Reserved query-context extension key for the trigger reason supplied by frontend. -pub const TRIGGER_REASON_EXTENSION_KEY: &str = "__greptime_event.trigger_reason"; /// Reserved query-context extension key for the authenticated database creator. pub const CREATE_DATABASE_CREATOR_EXTENSION_KEY: &str = "__greptime_create_database.creator"; /// Internal gRPC metadata key for the authenticated database creator. @@ -336,7 +334,6 @@ impl TryFrom for DdlTask { #[derive(Clone)] pub struct SubmitDdlTaskRequest { - pub query_context: QueryContext, pub wait: bool, pub timeout: Duration, pub task: DdlTask, @@ -344,9 +341,8 @@ pub struct SubmitDdlTaskRequest { impl SubmitDdlTaskRequest { /// The default constructor for [`SubmitDdlTaskRequest`]. - pub fn new(query_context: QueryContext, task: DdlTask) -> Self { + pub fn new(task: DdlTask) -> Self { Self { - query_context, wait: Self::default_wait(), timeout: Self::default_timeout(), task, @@ -377,26 +373,11 @@ impl TryFrom for PbDdlTaskRequest { fn try_from(request: SubmitDdlTaskRequest) -> Result { let SubmitDdlTaskRequest { - mut query_context, wait, timeout, task, } = request; - query_context - .extensions - .remove(CREATE_DATABASE_CREATOR_EXTENSION_KEY); - if let DdlTask::CreateDatabase(CreateDatabaseTask { - creator: Some(creator), - .. - }) = &task - { - query_context.extensions.insert( - CREATE_DATABASE_CREATOR_EXTENSION_KEY.to_string(), - serde_json::to_string(creator).context(error::SerdeJsonSnafu)?, - ); - } - let task = match task { DdlTask::CreateTable(task) => Task::CreateTableTask(task.try_into()?), DdlTask::DropTable(task) => Task::DropTableTask(task.into()), @@ -444,10 +425,12 @@ impl TryFrom for PbDdlTaskRequest { Ok(Self { header: None, - query_context: Some(query_context.into()), + query_context: None, timeout_secs: ddl_timeout_secs(timeout), wait, task: Some(task), + event_context: None, + actor: None, }) } } @@ -1668,22 +1651,6 @@ pub struct QueryContext { pub sst_min_sequences: HashMap, } -/// Builds event context from frontend query context metadata. -pub fn event_context_from_query_context(query_context: &QueryContext) -> PersistentEventContext { - let reason = query_context - .extensions - .get(TRIGGER_REASON_EXTENSION_KEY) - .map(|reason| TriggerReason::from_extension(reason)) - .unwrap_or_default(); - let channel = Channel::from(u32::from(query_context.channel)); - let context = PersistentEventContext::new(reason); - if channel == Channel::Unknown { - context - } else { - context.with_protocol(channel.as_ref()) - } -} - impl QueryContext { /// Get the current catalog pub fn current_catalog(&self) -> &str { @@ -1710,6 +1677,12 @@ impl QueryContext { self.channel } + /// Returns the protocol derived from the typed query channel. + pub fn protocol(&self) -> Option { + let channel = Channel::from(u32::from(self.channel)); + (channel != Channel::Unknown).then(|| channel.as_ref().to_string()) + } + pub fn snapshot_seqs(&self) -> &HashMap { &self.snapshot_seqs } @@ -1922,10 +1895,7 @@ mod tests { #[test] fn test_undrop_table_task_pb_roundtrip() { let expected = UndropTableTask { table_id: 1024 }; - let request = SubmitDdlTaskRequest::new( - QueryContext::default(), - DdlTask::UndropTable(expected.clone()), - ); + let request = SubmitDdlTaskRequest::new(DdlTask::UndropTable(expected.clone())); let pb = PbDdlTaskRequest::try_from(request).unwrap(); let pb_task = pb.task.unwrap(); @@ -1937,10 +1907,7 @@ mod tests { #[test] fn test_purge_dropped_table_task_pb_roundtrip() { let expected = PurgeDroppedTableTask { table_id: 1024 }; - let request = SubmitDdlTaskRequest::new( - QueryContext::default(), - DdlTask::PurgeDroppedTable(expected.clone()), - ); + let request = SubmitDdlTaskRequest::new(DdlTask::PurgeDroppedTable(expected.clone())); let pb = PbDdlTaskRequest::try_from(request).unwrap(); let pb_task = pb.task.unwrap(); @@ -1949,53 +1916,6 @@ mod tests { assert!(matches!(de, DdlTask::PurgeDroppedTable(task) if task == expected)); } - #[test] - fn test_create_database_creator_transport_sanitizes_extension() { - let request = |creator| { - let mut query_context = QueryContext::default(); - query_context.extensions.insert( - CREATE_DATABASE_CREATOR_EXTENSION_KEY.to_string(), - "spoofed".to_string(), - ); - SubmitDdlTaskRequest::new( - query_context, - DdlTask::new_create_database( - "greptime".to_string(), - "test".to_string(), - false, - HashMap::new(), - creator, - ), - ) - }; - let creator = CreatorGrantIntent { - username: "alice".to_string(), - created_at_ns: 42, - }; - - let mut pb = PbDdlTaskRequest::try_from(request(Some(creator.clone()))).unwrap(); - let encoded = pb - .query_context - .as_ref() - .unwrap() - .extensions - .get(CREATE_DATABASE_CREATOR_EXTENSION_KEY) - .unwrap(); - assert_eq!(creator, serde_json::from_str(encoded).unwrap()); - assert!(matches!( - DdlTask::try_from(pb.task.take().unwrap()).unwrap(), - DdlTask::CreateDatabase(CreateDatabaseTask { creator: None, .. }) - )); - - let pb = PbDdlTaskRequest::try_from(request(None)).unwrap(); - assert!( - !pb.query_context - .unwrap() - .extensions - .contains_key(CREATE_DATABASE_CREATOR_EXTENSION_KEY) - ); - } - #[test] fn test_undrop_table_task_json_roundtrip() { let task = UndropTableTask { table_id: 1024 }; @@ -2262,44 +2182,4 @@ mod tests { let reason: TriggerReason = serde_json::from_str("\"future_reason\"").unwrap(); assert_eq!(TriggerReason::Unknown, reason); } - - #[test] - fn test_event_context_serialization() { - let context = PersistentEventContext::new(TriggerReason::Manual).with_protocol("mysql"); - - assert_eq!( - serde_json::json!({ - "reason": "manual", - "protocol": "mysql", - }), - serde_json::to_value(context).unwrap() - ); - - assert_eq!( - serde_json::json!({ "reason": "manual" }), - serde_json::to_value(PersistentEventContext::new(TriggerReason::Manual)).unwrap() - ); - } - - #[test] - fn test_event_context_from_query_context_preserves_extensions() { - let mut query_context = QueryContext::default(); - query_context.extensions.insert( - TRIGGER_REASON_EXTENSION_KEY.to_string(), - TriggerReason::AutoCreate.as_ref().to_string(), - ); - query_context.channel = 4; - - let event_context = event_context_from_query_context(&query_context); - - assert_eq!(event_context.reason, TriggerReason::AutoCreate); - assert_eq!(event_context.protocol.as_deref(), Some("prometheus")); - assert_eq!( - query_context - .extensions - .get(TRIGGER_REASON_EXTENSION_KEY) - .map(String::as_str), - Some("auto_create") - ); - } } diff --git a/src/common/procedure/src/local.rs b/src/common/procedure/src/local.rs index 8fdd3167583..c870cc39a67 100644 --- a/src/common/procedure/src/local.rs +++ b/src/common/procedure/src/local.rs @@ -1062,7 +1062,7 @@ mod tests { use super::*; use crate::error::{self, Error}; - use crate::store::state_store::ObjectStateStore; + use crate::store::state_store::{ObjectStateStore, StateStore}; use crate::test_util::InMemoryPoisonStore; use crate::{ ChildSubmissionOutcome, Context, EventContext, EventTrigger, Procedure, ProcedureEvent, @@ -1721,7 +1721,63 @@ mod tests { } #[tokio::test] - async fn test_recovery_emits_recovered_event() { + async fn test_recovery_loads_legacy_procedure_message() { + let dir = create_temp_dir("legacy_procedure_message_recovery"); + let object_store = test_util::new_object_store(&dir); + let state_store = Arc::new(ObjectStateStore::new(object_store)); + let poison_manager = Arc::new(InMemoryPoisonStore::new()); + let event_recorder = Arc::new(CapturingEventRecorder::default()); + let manager = LocalManager::new( + ManagerConfig { + parent_path: "data/".to_string(), + ..Default::default() + }, + state_store.clone(), + poison_manager, + None, + None, + ); + manager + .event_recorder_handle() + .install(event_recorder.clone()); + manager.manager_ctx.start(); + manager + .register_loader("ProcedureToLoad", ProcedureToLoad::loader()) + .unwrap(); + + let procedure = ProcedureToLoad::new("legacy recovered submission"); + let procedure_id = ProcedureId::random(); + let key = format!("data/procedure/{procedure_id}/0000000000.step"); + // This is the persisted JSON shape written before ProcedureMessage had `context`. + let legacy_message = serde_json::json!({ + "type_name": procedure.type_name(), + "data": procedure.dump().unwrap(), + "parent_id": null, + "step": 0, + }); + state_store + .put(&key, serde_json::to_vec(&legacy_message).unwrap()) + .await + .unwrap(); + + manager.recover().await.unwrap(); + + assert!( + manager + .procedure_state(procedure_id) + .await + .unwrap() + .is_some() + ); + wait_for_trigger(&event_recorder, EventTrigger::Recovered).await; + assert_eq!( + event_recorder.procedure_context(procedure_id), + Some(ProcedureContext::default()) + ); + } + + #[tokio::test] + async fn test_recovery_preserves_procedure_context() { let dir = create_temp_dir("recovery_submission_event"); let object_store = test_util::new_object_store(&dir); let state_store = Arc::new(ObjectStateStore::new(object_store.clone())); @@ -1747,6 +1803,16 @@ mod tests { let procedure = ProcedureToLoad::new("recovered submission"); let procedure_id = ProcedureId::random(); + let mut event_context = + PersistentEventContext::new(TriggerReason::AutoRebalance).with_protocol("mysql"); + event_context.extensions.insert( + "source".to_string(), + serde_json::Value::String("test".to_string()), + ); + let procedure_context = ProcedureContext { + actor: Some("alice".to_string()), + event_context: Some(event_context), + }; ProcedureStore::from_object_store(object_store) .store_procedure_with_context( procedure_id, @@ -1754,9 +1820,7 @@ mod tests { procedure.type_name().to_string(), procedure.dump().unwrap(), None, - ProcedureContext::from_event_context(PersistentEventContext::new( - TriggerReason::AutoRebalance, - )), + procedure_context.clone(), ) .await .unwrap(); @@ -1773,13 +1837,8 @@ mod tests { wait_for_trigger(&event_recorder, EventTrigger::Recovered).await; assert!(!event_recorder.triggers().contains(&EventTrigger::Submitted)); assert_eq!( - event_recorder - .procedure_context(procedure_id) - .unwrap() - .event_context - .unwrap() - .reason, - TriggerReason::AutoRebalance + event_recorder.procedure_context(procedure_id), + Some(procedure_context) ); } diff --git a/src/common/procedure/src/procedure.rs b/src/common/procedure/src/procedure.rs index cd11aea9215..ba6a5d275cb 100644 --- a/src/common/procedure/src/procedure.rs +++ b/src/common/procedure/src/procedure.rs @@ -35,6 +35,9 @@ pub type Output = Arc; /// Context attached to a procedure submission and inherited by its children. #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] pub struct ProcedureContext { + /// Effective user that submitted the procedure. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub actor: Option, /// Context describing why and how the procedure was submitted. #[serde(default, skip_serializing_if = "Option::is_none")] pub event_context: Option, @@ -43,12 +46,13 @@ pub struct ProcedureContext { impl ProcedureContext { pub fn from_event_context(event_context: PersistentEventContext) -> Self { Self { + actor: None, event_context: Some(event_context), } } pub fn is_empty(&self) -> bool { - self.event_context.is_none() + self.actor.is_none() && self.event_context.is_none() } } @@ -498,9 +502,9 @@ impl ProcedureWithId { } } - /// Attaches event context to this procedure submission. - pub fn with_event_context(mut self, event_context: PersistentEventContext) -> Self { - self.context.event_context = Some(event_context); + /// Attaches a complete persisted context to this procedure submission. + pub fn with_context(mut self, context: ProcedureContext) -> Self { + self.context = context; self } } diff --git a/src/common/procedure/src/store.rs b/src/common/procedure/src/store.rs index 4867046f8de..b2f19b2f94c 100644 --- a/src/common/procedure/src/store.rs +++ b/src/common/procedure/src/store.rs @@ -514,6 +514,14 @@ mod tests { r#"{"type_name":"TestMessage","data":"no parent id","parent_id":"9f805a1f-05f7-490c-9f91-bd56e3cc54c1","step":4,"context":{"event_context":{"reason":"auto_rebalance"}}}"# ); + message.context.actor = Some("alice".to_string()); + let json = serde_json::to_string(&message).unwrap(); + assert!(json.contains(r#""actor":"alice""#)); + assert_eq!( + serde_json::from_str::(&json).unwrap(), + message + ); + let legacy: ProcedureMessage = serde_json::from_str( r#"{"type_name":"TestMessage","data":"legacy","parent_id":null,"step":1}"#, ) diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 0c300835e0e..78b701a3442 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -2281,7 +2281,7 @@ mod tests { impl ProcedureExecutor for NoopProcedureExecutor { async fn submit_ddl_task( &self, - _ctx: &ExecutorContext, + _ctx: ExecutorContext, _request: SubmitDdlTaskRequest, ) -> common_meta::error::Result { common_meta::error::UnsupportedSnafu { @@ -2358,7 +2358,7 @@ mod tests { impl ProcedureExecutor for MockProcedureExecutor { async fn submit_ddl_task( &self, - _ctx: &ExecutorContext, + _ctx: ExecutorContext, request: SubmitDdlTaskRequest, ) -> common_meta::error::Result { self.submitted.lock().unwrap().push(request.task.clone()); diff --git a/src/meta-client/Cargo.toml b/src/meta-client/Cargo.toml index 28c03b3f0dd..984f434cc51 100644 --- a/src/meta-client/Cargo.toml +++ b/src/meta-client/Cargo.toml @@ -30,6 +30,7 @@ tokio-stream = { workspace = true, features = ["net"] } tonic.workspace = true [dev-dependencies] +common-event-recorder.workspace = true datatypes.workspace = true futures.workspace = true meta-srv = { workspace = true, features = ["mock"] } diff --git a/src/meta-client/src/client.rs b/src/meta-client/src/client.rs index 4fc6020fffb..e0932af5a69 100644 --- a/src/meta-client/src/client.rs +++ b/src/meta-client/src/client.rs @@ -49,7 +49,10 @@ use common_meta::peer::PeerDiscovery; use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor}; use common_meta::range_stream::PaginationStream; use common_meta::rpc::KeyValue; -use common_meta::rpc::ddl::{SubmitDdlTaskRequest, SubmitDdlTaskResponse}; +use common_meta::rpc::ddl::{ + CREATE_DATABASE_CREATOR_EXTENSION_KEY, CreateDatabaseTask, DdlTask, SubmitDdlTaskRequest, + SubmitDdlTaskResponse, +}; use common_meta::rpc::procedure::{ AddRegionFollowerRequest, AddTableFollowerRequest, GcRegionsRequest, GcResponse, GcTableRequest, ManageRegionFollowerRequest, MigrateRegionRequest, MigrateRegionResponse, @@ -66,7 +69,7 @@ use common_time::util::DefaultSystemTimer; use config::Client as ConfigClient; use futures::TryStreamExt; use heartbeat::{Client as HeartbeatClient, HeartbeatConfig}; -use procedure::Client as ProcedureClient; +use procedure::{Client as ProcedureClient, procedure_actor, procedure_event_context}; use serde::de::DeserializeOwned; use snafu::{OptionExt, ResultExt}; use store::Client as StoreClient; @@ -75,7 +78,7 @@ pub use self::heartbeat::{HeartbeatSender, HeartbeatStream}; use crate::client::ask_leader::{LeaderProviderFactoryImpl, LeaderProviderFactoryRef}; use crate::error::{ ConvertMetaConfigSnafu, ConvertMetaRequestSnafu, ConvertMetaResponseSnafu, Error, - GetFlowStatSnafu, NotStartedSnafu, Result, + GetFlowStatSnafu, MissingQueryContextSnafu, NotStartedSnafu, Result, }; pub type Id = u64; @@ -334,10 +337,10 @@ pub trait RegionFollowerClient: Sync + Send + Debug { impl ProcedureExecutor for MetaClient { async fn submit_ddl_task( &self, - _ctx: &ExecutorContext, + ctx: ExecutorContext, request: SubmitDdlTaskRequest, ) -> MetaResult { - self.submit_ddl_task(request) + MetaClient::submit_ddl_task(self, ctx, request) .await .map_err(BoxedError::new) .context(meta_error::ExternalSnafu) @@ -345,10 +348,10 @@ impl ProcedureExecutor for MetaClient { async fn migrate_region( &self, - _ctx: &ExecutorContext, + ctx: &ExecutorContext, request: MigrateRegionRequest, ) -> MetaResult { - self.migrate_region(request) + self.migrate_region(ctx, request) .await .map_err(BoxedError::new) .context(meta_error::ExternalSnafu) @@ -418,10 +421,10 @@ impl ProcedureExecutor for MetaClient { async fn gc_regions( &self, - _ctx: &ExecutorContext, + ctx: &ExecutorContext, request: GcRegionsRequest, ) -> MetaResult { - self.gc_regions(request) + self.gc_regions(ctx, request) .await .map_err(BoxedError::new) .context(meta_error::ExternalSnafu) @@ -429,10 +432,10 @@ impl ProcedureExecutor for MetaClient { async fn gc_table( &self, - _ctx: &ExecutorContext, + ctx: &ExecutorContext, request: GcTableRequest, ) -> MetaResult { - self.gc_table(request) + self.gc_table(ctx, request) .await .map_err(BoxedError::new) .context(meta_error::ExternalSnafu) @@ -807,10 +810,12 @@ impl MetaClient { /// Submit a region migration task. pub async fn migrate_region( &self, + context: &ExecutorContext, request: MigrateRegionRequest, ) -> Result { self.procedure_client()? .migrate_region( + context, request.region_id, request.from_peer, request.to_peer, @@ -825,28 +830,57 @@ impl MetaClient { } /// Manually trigger GC for specific regions. - pub async fn gc_regions(&self, request: GcRegionsRequest) -> Result { - self.procedure_client()?.gc_regions(request).await + pub async fn gc_regions( + &self, + context: &ExecutorContext, + request: GcRegionsRequest, + ) -> Result { + self.procedure_client()?.gc_regions(context, request).await } /// Manually trigger GC for a table (all its regions). - pub async fn gc_table(&self, request: GcTableRequest) -> Result { - self.procedure_client()?.gc_table(request).await + pub async fn gc_table( + &self, + context: &ExecutorContext, + request: GcTableRequest, + ) -> Result { + self.procedure_client()?.gc_table(context, request).await } - /// Submit a DDL task + /// Submit a DDL task. pub async fn submit_ddl_task( &self, - req: SubmitDdlTaskRequest, + context: ExecutorContext, + request: SubmitDdlTaskRequest, ) -> Result { - let res = self - .procedure_client()? - .submit_ddl_task(req.try_into().context(ConvertMetaRequestSnafu)?) + let event_context = procedure_event_context(&context); + let actor = procedure_actor(&context); + let mut query_context = context.query_context.context(MissingQueryContextSnafu)?; + query_context + .extensions + .remove(CREATE_DATABASE_CREATOR_EXTENSION_KEY); + if let DdlTask::CreateDatabase(CreateDatabaseTask { + creator: Some(creator), + .. + }) = &request.task + { + query_context.extensions.insert( + CREATE_DATABASE_CREATOR_EXTENSION_KEY.to_string(), + serde_json::to_string(creator).context(ConvertMetaConfigSnafu)?, + ); + } + + let mut request: api::v1::meta::DdlTaskRequest = + request.try_into().context(ConvertMetaRequestSnafu)?; + request.query_context = Some(api::v1::QueryContext::from(query_context)); + request.event_context = event_context; + request.actor = actor; + + self.procedure_client()? + .submit_ddl_task(request) .await? .try_into() - .context(ConvertMetaResponseSnafu)?; - - Ok(res) + .context(ConvertMetaResponseSnafu) } pub fn heartbeat_client(&self) -> Result { diff --git a/src/meta-client/src/client/procedure.rs b/src/meta-client/src/client/procedure.rs index d7d5ccfaff6..ccf2c28188f 100644 --- a/src/meta-client/src/client/procedure.rs +++ b/src/meta-client/src/client/procedure.rs @@ -20,11 +20,13 @@ use api::v1::meta::ddl_task_request::Task; use api::v1::meta::procedure_service_client::ProcedureServiceClient; use api::v1::meta::{ DdlTaskRequest, DdlTaskResponse, GcRegionsRequest, GcRegionsResponse, GcTableRequest, - GcTableResponse, MigrateRegionRequest, MigrateRegionResponse, ProcedureDetailRequest, - ProcedureDetailResponse, ProcedureId, ProcedureStateResponse, QueryProcedureRequest, - ReconcileRequest, ReconcileResponse, RequestHeader, ResponseHeader, Role, + GcTableResponse, MigrateRegionRequest, MigrateRegionResponse, ProcedureActor, + ProcedureDetailRequest, ProcedureDetailResponse, ProcedureEventContext, ProcedureId, + ProcedureStateResponse, QueryProcedureRequest, ReconcileRequest, ReconcileResponse, + RequestHeader, ResponseHeader, Role, }; use common_grpc::channel_manager::ChannelManager; +use common_meta::procedure_executor::ExecutorContext; use common_meta::rpc::ddl::{ CREATE_DATABASE_CREATOR_EXTENSION_KEY, CREATE_DATABASE_CREATOR_METADATA_KEY, }; @@ -43,6 +45,33 @@ use crate::client::{Id, LeaderProviderRef, util}; use crate::error; use crate::error::Result; +/// Builds the event context transported by a procedure RPC. +/// +/// The caller can only supply reason/extensions. Protocol is derived here from +/// the trusted, typed query channel held locally in the executor context. +pub(crate) fn procedure_event_context(context: &ExecutorContext) -> Option { + context.event_input.as_ref().map(|input| { + let mut event_context = ProcedureEventContext::from(input); + event_context.protocol = context + .query_context + .as_ref() + .and_then(|query_context| query_context.protocol()) + .unwrap_or_default(); + event_context + }) +} + +/// Builds the optional procedure actor transported by a procedure RPC. +pub(crate) fn procedure_actor(context: &ExecutorContext) -> Option { + context + .actor + .as_deref() + .filter(|username| !username.is_empty()) + .map(|username| ProcedureActor { + username: username.to_string(), + }) +} + #[derive(Clone, Debug)] pub struct Client { inner: Arc>, @@ -92,6 +121,7 @@ impl Client { /// - `timeout`: timeout for downgrading region and upgrading region operations pub async fn migrate_region( &self, + context: &ExecutorContext, region_id: u64, from_peer: u64, to_peer: u64, @@ -99,7 +129,7 @@ impl Client { ) -> Result { let inner = self.inner.read().await; inner - .migrate_region(region_id, from_peer, to_peer, timeout) + .migrate_region(context, region_id, from_peer, to_peer, timeout) .await } @@ -114,14 +144,22 @@ impl Client { inner.list_procedures().await } - pub async fn gc_regions(&self, request: MetaGcRegionsRequest) -> Result { + pub async fn gc_regions( + &self, + context: &ExecutorContext, + request: MetaGcRegionsRequest, + ) -> Result { let inner = self.inner.read().await; - inner.gc_regions(request).await + inner.gc_regions(context, request).await } - pub async fn gc_table(&self, request: MetaGcTableRequest) -> Result { + pub async fn gc_table( + &self, + context: &ExecutorContext, + request: MetaGcTableRequest, + ) -> Result { let inner = self.inner.read().await; - inner.gc_table(request).await + inner.gc_table(context, request).await } } @@ -225,6 +263,7 @@ impl Inner { async fn migrate_region( &self, + context: &ExecutorContext, region_id: u64, from_peer: u64, to_peer: u64, @@ -235,6 +274,8 @@ impl Inner { from_peer, to_peer, timeout_secs: timeout.as_secs() as u32, + event_context: procedure_event_context(context), + actor: procedure_actor(context), ..Default::default() }; @@ -278,7 +319,11 @@ impl Inner { .await } - async fn gc_regions(&self, request: MetaGcRegionsRequest) -> Result { + async fn gc_regions( + &self, + context: &ExecutorContext, + request: MetaGcRegionsRequest, + ) -> Result { let timeout = request.timeout; let req = GcRegionsRequest { header: Some(RequestHeader { @@ -290,6 +335,8 @@ impl Inner { region_ids: request.region_ids, full_file_listing: request.full_file_listing, timeout_secs: gc_timeout_secs(timeout), + event_context: procedure_event_context(context), + actor: procedure_actor(context), }; let resp: GcRegionsResponse = self @@ -315,7 +362,11 @@ impl Inner { }) } - async fn gc_table(&self, request: MetaGcTableRequest) -> Result { + async fn gc_table( + &self, + context: &ExecutorContext, + request: MetaGcTableRequest, + ) -> Result { let timeout = request.timeout; let req = GcTableRequest { header: Some(RequestHeader { @@ -329,6 +380,8 @@ impl Inner { table_name: request.table_name, full_file_listing: request.full_file_listing, timeout_secs: gc_timeout_secs(timeout), + event_context: procedure_event_context(context), + actor: procedure_actor(context), }; let resp: GcTableResponse = self @@ -460,11 +513,14 @@ mod tests { ReconcileResponse, ResponseHeader, Role, }; use async_trait::async_trait; + use common_base::protocol::Channel; use common_error::status_code::StatusCode; + use common_event_recorder::{PersistentEventContext, ProcedureEventInput}; + use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor}; use common_meta::rpc::ddl::{ CREATE_DATABASE_CREATOR_EXTENSION_KEY, CREATE_DATABASE_CREATOR_METADATA_KEY, CommentObjectType, CommentOnTask, CreatorGrantIntent, DdlTask, QueryContext, - SubmitDdlTaskRequest, + SubmitDdlTaskRequest, TriggerReason, }; use common_telemetry::common_error::ext::ErrorExt; use common_telemetry::info; @@ -474,8 +530,8 @@ mod tests { use tonic::codec::CompressionEncoding; use tonic::{Request, Response, Status}; - use super::gc_timeout_secs; use crate::client::MetaClientBuilder; + use crate::client::procedure::{gc_timeout_secs, procedure_actor, procedure_event_context}; #[test] fn test_gc_timeout_secs() { @@ -486,6 +542,57 @@ mod tests { assert_eq!(gc_timeout_secs(Some(Duration::from_secs(10))), 10); } + #[test] + fn test_procedure_event_context_derives_protocol_from_query_context() { + let context = ExecutorContext { + query_context: Some(QueryContext { + channel: Channel::Postgres as u8, + ..Default::default() + }), + event_input: Some(ProcedureEventInput::new(TriggerReason::Manual)), + ..Default::default() + }; + + assert_eq!( + procedure_event_context(&context).map(PersistentEventContext::from), + Some(PersistentEventContext::new(TriggerReason::Manual).with_protocol("postgres")) + ); + + let automatic_context = ExecutorContext { + event_input: Some(ProcedureEventInput::new(TriggerReason::ScheduledGc)), + ..Default::default() + }; + assert_eq!( + procedure_event_context(&automatic_context).map(PersistentEventContext::from), + Some(PersistentEventContext::new(TriggerReason::ScheduledGc)) + ); + + let unknown_channel_context = ExecutorContext { + query_context: Some(QueryContext::default()), + event_input: Some(ProcedureEventInput::new(TriggerReason::Manual)), + ..Default::default() + }; + assert_eq!( + procedure_event_context(&unknown_channel_context).map(PersistentEventContext::from), + Some(PersistentEventContext::new(TriggerReason::Manual)) + ); + } + + #[test] + fn test_procedure_actor() { + let context = ExecutorContext { + actor: Some(String::new()), + ..Default::default() + }; + assert!(procedure_actor(&context).is_none()); + + let context = ExecutorContext { + actor: Some("alice".to_string()), + ..Default::default() + }; + assert_eq!(procedure_actor(&context).unwrap().username, "alice"); + } + #[derive(Clone)] struct MockHeartbeat { leader_addr: String, @@ -621,19 +728,28 @@ mod tests { username: "alice".to_string(), created_at_ns: 42, }; - client - .submit_ddl_task(SubmitDdlTaskRequest::new( - QueryContext::default(), - DdlTask::new_create_database( - "greptime".to_string(), - "metrics".to_string(), - false, - Default::default(), - Some(creator.clone()), - ), - )) - .await - .unwrap(); + let executor_context = |actor: String| ExecutorContext { + query_context: Some(QueryContext { + channel: Channel::Postgres as u8, + ..Default::default() + }), + actor: Some(actor), + event_input: Some(ProcedureEventInput::new(TriggerReason::Manual)), + ..Default::default() + }; + ProcedureExecutor::submit_ddl_task( + &client, + executor_context("effective-user".to_string()), + SubmitDdlTaskRequest::new(DdlTask::new_create_database( + "greptime".to_string(), + "metrics".to_string(), + false, + Default::default(), + Some(creator.clone()), + )), + ) + .await + .unwrap(); let request = request_rx.recv().await.unwrap(); let encoded = serde_json::to_string(&creator).unwrap(); @@ -647,9 +763,36 @@ mod tests { .as_ref(), encoded.as_bytes() ); - let extensions = &request.into_inner().query_context.unwrap().extensions; + let request = request.into_inner(); + assert_eq!(request.actor.unwrap().username, "effective-user"); + assert_eq!( + PersistentEventContext::from(request.event_context.unwrap()), + PersistentEventContext::new(TriggerReason::Manual).with_protocol("postgres") + ); + let extensions = &request.query_context.unwrap().extensions; assert_eq!(extensions[CREATE_DATABASE_CREATOR_EXTENSION_KEY], encoded); + ProcedureExecutor::submit_ddl_task( + &client, + executor_context(String::new()), + SubmitDdlTaskRequest::new(DdlTask::new_drop_database( + "greptime".to_string(), + "metrics".to_string(), + false, + )), + ) + .await + .unwrap(); + assert!( + request_rx + .recv() + .await + .unwrap() + .into_inner() + .actor + .is_none() + ); + server_handle.abort(); } @@ -689,22 +832,28 @@ mod tests { .build(); client.start(&[addr_str.as_str()]).await.unwrap(); - let mut request = SubmitDdlTaskRequest::new( - QueryContext::default(), - DdlTask::new_comment_on(CommentOnTask { - catalog_name: "greptime".to_string(), - schema_name: "public".to_string(), - object_type: CommentObjectType::Table, - object_name: "test_table".to_string(), - column_name: None, - object_id: None, - comment: Some("timeout".to_string()), - }), - ); + let mut request = SubmitDdlTaskRequest::new(DdlTask::new_comment_on(CommentOnTask { + catalog_name: "greptime".to_string(), + schema_name: "public".to_string(), + object_type: CommentObjectType::Table, + object_name: "test_table".to_string(), + column_name: None, + object_id: None, + comment: Some("timeout".to_string()), + })); request.timeout = Duration::from_secs(1); let now = Instant::now(); - let err = client.submit_ddl_task(request).await.unwrap_err(); + let err = client + .submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + request, + ) + .await + .unwrap_err(); let elapsed = now.elapsed(); // The request should be cancelled within 1 second. assert!(elapsed < Duration::from_secs(2)); diff --git a/src/meta-client/src/error.rs b/src/meta-client/src/error.rs index 5121bba847a..4623cfee33b 100644 --- a/src/meta-client/src/error.rs +++ b/src/meta-client/src/error.rs @@ -67,6 +67,12 @@ pub enum Error { location: Location, }, + #[snafu(display("Procedure submission requires query context"))] + MissingQueryContext { + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to send heartbeat: {}", err_msg))] SendHeartbeat { err_msg: String, @@ -141,6 +147,7 @@ impl ErrorExt for Error { | Error::NoLeader { .. } | Error::AskLeaderTimeout { .. } | Error::NotStarted { .. } + | Error::MissingQueryContext { .. } | Error::SendHeartbeat { .. } | Error::CreateHeartbeatStream { .. } | Error::CreateChannel { .. } diff --git a/src/meta-srv/src/gc/ctx.rs b/src/meta-srv/src/gc/ctx.rs index c9cc8049867..4b46a941e25 100644 --- a/src/meta-srv/src/gc/ctx.rs +++ b/src/meta-srv/src/gc/ctx.rs @@ -29,10 +29,9 @@ use common_meta::key::DroppedTableName; use common_meta::key::TableMetadataManagerRef; use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; -use common_meta::rpc::ddl::PersistentEventContext; #[cfg(feature = "enterprise")] use common_meta::rpc::ddl::PurgeDroppedTableTask; -use common_procedure::{ProcedureManagerRef, ProcedureWithId, watcher}; +use common_procedure::{ProcedureContext, ProcedureManagerRef, ProcedureWithId, watcher}; use common_telemetry::debug; use snafu::{OptionExt as _, ResultExt as _}; use store_api::storage::{GcReport, RegionId}; @@ -68,7 +67,7 @@ pub(crate) trait SchedulerCtx: Send + Sync { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result; #[cfg(feature = "enterprise")] @@ -256,14 +255,14 @@ impl SchedulerCtx for DefaultGcSchedulerCtx { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { self.gc_regions_inner( region_ids, full_file_listing, timeout, region_routes_override, - event_context, + procedure_context, ) .await } @@ -310,7 +309,7 @@ impl DefaultGcSchedulerCtx { full_file_listing: bool, timeout: Duration, region_routes_override: Region2Peers, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { debug!( "Sending GC instruction for {} regions (full_file_listing: {})", @@ -328,7 +327,7 @@ impl DefaultGcSchedulerCtx { region_routes_override, ); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); let id = procedure_with_id.id; diff --git a/src/meta-srv/src/gc/handler.rs b/src/meta-srv/src/gc/handler.rs index 16b29a39c7e..eda87fbb144 100644 --- a/src/meta-srv/src/gc/handler.rs +++ b/src/meta-srv/src/gc/handler.rs @@ -18,7 +18,7 @@ use std::time::Instant; use common_catalog::consts::MITO_ENGINE; use common_meta::datanode::{RegionManifestInfo, RegionStat}; use common_meta::peer::Peer; -use common_meta::rpc::ddl::PersistentEventContext; +use common_procedure::ProcedureContext; use common_telemetry::tracing::Instrument as _; use common_telemetry::{debug, error, info, warn}; use futures::StreamExt; @@ -39,7 +39,7 @@ use crate::metrics::METRIC_META_GC_CANDIDATE_REGIONS; impl GcScheduler { pub(crate) async fn trigger_gc( &self, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { let start_time = Instant::now(); info!("Starting GC cycle"); @@ -111,7 +111,7 @@ impl GcScheduler { datanode_to_candidates, dropped_assignment.force_full_listing, dropped_assignment.region_routes_override, - event_context, + procedure_context, ) .instrument(common_telemetry::tracing::info_span!( "meta_gc_dispatch_to_datanodes" @@ -215,7 +215,7 @@ impl GcScheduler { datanode_to_candidates: HashMap>, force_full_listing_by_peer: HashMap>, region_routes_override_by_peer: HashMap, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> GcJobReport { let mut per_datanode_reports = HashMap::new(); let mut failed_datanodes: HashMap<_, Vec<_>> = HashMap::new(); @@ -237,7 +237,7 @@ impl GcScheduler { .get(&peer) .cloned() .unwrap_or_default(); - let event_context = event_context.clone(); + let procedure_context = procedure_context.clone(); async move { ( peer, @@ -247,7 +247,7 @@ impl GcScheduler { candidates, force_full_listing, region_routes_override, - event_context, + procedure_context, ) .await, ) @@ -286,7 +286,7 @@ impl GcScheduler { candidates: Vec<(TableId, GcCandidate)>, force_full_listing: HashSet, region_routes_override: Region2Peers, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { info!( "Starting GC for datanode {} with {} candidate regions", @@ -336,7 +336,7 @@ impl GcScheduler { false, self.config.mailbox_timeout, region_routes_override.clone(), - event_context.clone(), + procedure_context.clone(), ) .instrument(common_telemetry::tracing::info_span!( "meta_gc_call_datanode", @@ -369,7 +369,7 @@ impl GcScheduler { true, self.config.mailbox_timeout, region_routes_override, - event_context, + procedure_context, ) .instrument(common_telemetry::tracing::info_span!( "meta_gc_call_datanode", diff --git a/src/meta-srv/src/gc/mock.rs b/src/meta-srv/src/gc/mock.rs index 94749dfb469..dcf913462c2 100644 --- a/src/meta-srv/src/gc/mock.rs +++ b/src/meta-srv/src/gc/mock.rs @@ -26,12 +26,13 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use common_catalog::consts::MITO_ENGINE; +use common_event_recorder::PersistentEventContext; use common_meta::datanode::{RegionManifestInfo, RegionStat}; use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; use common_meta::peer::Peer; -use common_meta::rpc::ddl::PersistentEventContext; use common_meta::rpc::router::{Region, RegionRoute}; +use common_procedure::ProcedureContext; use common_telemetry::debug; use ordered_float::OrderedFloat; use store_api::region_engine::RegionRole; @@ -214,10 +215,13 @@ impl SchedulerCtx for MockSchedulerCtx { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - event_context: PersistentEventContext, + procedure_context: ProcedureContext, ) -> Result { *self.gc_regions_calls.lock().unwrap() += 1; - self.gc_event_contexts.lock().unwrap().push(event_context); + self.gc_event_contexts + .lock() + .unwrap() + .push(procedure_context.event_context.unwrap_or_default()); // Check per-region error injection first (for any region) for ®ion_id in region_ids { diff --git a/src/meta-srv/src/gc/mock/basic.rs b/src/meta-srv/src/gc/mock/basic.rs index ed051cc839e..70f7a10b1a1 100644 --- a/src/meta-srv/src/gc/mock/basic.rs +++ b/src/meta-srv/src/gc/mock/basic.rs @@ -16,8 +16,10 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex}; use std::time::Instant; +use common_event_recorder::PersistentEventContext; use common_meta::peer::Peer; -use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; +use common_meta::rpc::ddl::TriggerReason; +use common_procedure::ProcedureContext; use common_telemetry::init_default_ut_logging; use store_api::region_engine::RegionRole; use store_api::storage::{FileId, FileRefsManifest, GcReport, RegionId}; @@ -36,7 +38,7 @@ async fn test_parallel_process_datanodes_empty() { HashMap::new(), HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; @@ -105,7 +107,7 @@ async fn test_parallel_process_datanodes_with_candidates() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; @@ -245,7 +247,18 @@ async fn test_handle_manual_gc_without_regions_records_manual_event_context() { last_tracker_cleanup: Arc::new(tokio::sync::Mutex::new(Instant::now())), }; - scheduler.handle_manual_gc(None, None, None).await.unwrap(); + scheduler + .handle_manual_gc( + None, + None, + None, + ProcedureContext { + event_context: Some(PersistentEventContext::new(TriggerReason::Manual)), + ..Default::default() + }, + ) + .await + .unwrap(); assert_eq!( ctx.gc_event_contexts.lock().unwrap().as_slice(), diff --git a/src/meta-srv/src/gc/mock/concurrent.rs b/src/meta-srv/src/gc/mock/concurrent.rs index 5bce9a1be8b..bfb7ab1d26c 100644 --- a/src/meta-srv/src/gc/mock/concurrent.rs +++ b/src/meta-srv/src/gc/mock/concurrent.rs @@ -18,8 +18,8 @@ use std::time::{Duration, Instant}; use common_meta::key::table_route::PhysicalTableRouteValue; use common_meta::peer::Peer; -use common_meta::rpc::ddl::PersistentEventContext; use common_meta::rpc::router::{Region, RegionRoute}; +use common_procedure::ProcedureContext; use common_telemetry::{info, init_default_ut_logging}; use store_api::region_engine::RegionRole; use store_api::storage::{FileId, FileRefsManifest, GcReport, RegionId}; @@ -102,7 +102,7 @@ async fn test_concurrent_table_processing_limits() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; @@ -190,7 +190,7 @@ async fn test_datanode_processes_tables_with_partial_gc_failures() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; @@ -296,7 +296,7 @@ async fn test_region_gc_concurrency_limit() { candidates.into_iter().map(|c| (table_id, c)).collect(), HashSet::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap(); @@ -413,7 +413,7 @@ async fn test_region_gc_concurrency_with_partial_failures() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; @@ -557,7 +557,7 @@ async fn test_region_gc_concurrency_with_retryable_errors() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; diff --git a/src/meta-srv/src/gc/mock/full_list.rs b/src/meta-srv/src/gc/mock/full_list.rs index 8193a9094e1..843e45134da 100644 --- a/src/meta-srv/src/gc/mock/full_list.rs +++ b/src/meta-srv/src/gc/mock/full_list.rs @@ -17,7 +17,7 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use common_meta::peer::Peer; -use common_meta::rpc::ddl::PersistentEventContext; +use common_procedure::ProcedureContext; use common_telemetry::init_default_ut_logging; use store_api::region_engine::RegionRole; use store_api::storage::{FileId, FileRefsManifest, GcReport, RegionId}; @@ -83,7 +83,7 @@ async fn test_full_file_listing_first_time_gc() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap(); @@ -157,7 +157,7 @@ async fn test_full_file_listing_interval_enforcement() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap(); @@ -184,7 +184,7 @@ async fn test_full_file_listing_interval_enforcement() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap(); @@ -260,7 +260,7 @@ async fn test_full_file_listing_no_interval_passed() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap(); @@ -284,7 +284,7 @@ async fn test_full_file_listing_no_interval_passed() { vec![(table_id, mock_candidate(region_id))], HashSet::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await .unwrap(); diff --git a/src/meta-srv/src/gc/mock/misc.rs b/src/meta-srv/src/gc/mock/misc.rs index 575a5c5fc08..fabf5bcae4b 100644 --- a/src/meta-srv/src/gc/mock/misc.rs +++ b/src/meta-srv/src/gc/mock/misc.rs @@ -17,7 +17,7 @@ use std::sync::{Arc, Mutex}; use std::time::Instant; use common_meta::peer::Peer; -use common_meta::rpc::ddl::PersistentEventContext; +use common_procedure::ProcedureContext; use common_telemetry::init_default_ut_logging; use store_api::storage::{FileRefsManifest, GcReport, RegionId}; @@ -76,7 +76,7 @@ async fn test_empty_file_refs_manifest() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; @@ -169,7 +169,7 @@ async fn test_multiple_regions_per_table() { datanode_to_candidates, HashMap::new(), HashMap::new(), - PersistentEventContext::default(), + ProcedureContext::default(), ) .await; diff --git a/src/meta-srv/src/gc/scheduler.rs b/src/meta-srv/src/gc/scheduler.rs index b89c0d06be9..857d4deb57c 100644 --- a/src/meta-srv/src/gc/scheduler.rs +++ b/src/meta-srv/src/gc/scheduler.rs @@ -16,9 +16,11 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::{Duration, Instant}; +use common_event_recorder::PersistentEventContext; use common_meta::DatanodeId; use common_meta::key::runtime_switch::RuntimeSwitchManagerRef; -use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; +use common_meta::rpc::ddl::TriggerReason; +use common_procedure::ProcedureContext; use common_telemetry::tracing::Instrument as _; use common_telemetry::{error, info}; use snafu::ResultExt; @@ -104,6 +106,7 @@ pub enum Event { full_file_listing: Option, /// Optional override for timeout. If None, uses scheduler config. timeout: Option, + procedure_context: ProcedureContext, }, } @@ -182,12 +185,13 @@ impl GcScheduler { region_ids, full_file_listing, timeout, + procedure_context, } => { info!("Received manually gc request"); let span = common_telemetry::tracing::info_span!("meta_gc_tick", trigger = "manual"); let result = self - .handle_manual_gc(region_ids, full_file_listing, timeout) + .handle_manual_gc(region_ids, full_file_listing, timeout, procedure_context) .instrument(span) .await; if let Err(e) = &result { @@ -218,7 +222,10 @@ impl GcScheduler { } let span = common_telemetry::tracing::info_span!("meta_gc_handle_tick"); let report = self - .trigger_gc(PersistentEventContext::new(TriggerReason::ScheduledGc)) + .trigger_gc(ProcedureContext { + actor: None, + event_context: Some(PersistentEventContext::new(TriggerReason::ScheduledGc)), + }) .instrument(span) .await?; @@ -294,6 +301,7 @@ impl GcScheduler { region_ids: Option>, full_file_listing: Option, timeout: Option, + procedure_context: ProcedureContext, ) -> Result { info!("Start to handle manual gc request"); @@ -304,9 +312,7 @@ impl GcScheduler { // No specific regions, use default tick behavior let Some(regions) = region_ids else { - let report = self - .trigger_gc(PersistentEventContext::new(TriggerReason::Manual)) - .await?; + let report = self.trigger_gc(procedure_context).await?; info!("Finished manual gc request"); return Ok(report); }; @@ -353,7 +359,7 @@ impl GcScheduler { full_listing, gc_timeout, Region2Peers::new(), - PersistentEventContext::new(TriggerReason::Manual), + procedure_context.clone(), ) .await?; combined_report.merge(report); @@ -367,7 +373,7 @@ impl GcScheduler { true, gc_timeout, dropped_routes_override, - PersistentEventContext::new(TriggerReason::Manual), + procedure_context, ) .await?; combined_report.merge(report); @@ -490,7 +496,7 @@ mod tests { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - _event_context: PersistentEventContext, + _procedure_context: ProcedureContext, ) -> Result { self.gc_regions_calls.fetch_add(1, Ordering::Relaxed); panic!("gc_regions should not be called in maintenance mode") @@ -564,7 +570,7 @@ mod tests { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - _event_context: PersistentEventContext, + _procedure_context: ProcedureContext, ) -> Result { Ok(GcReport::default()) } @@ -861,7 +867,10 @@ mod tests { *ctx.dropped_tables.lock().unwrap() = vec![dropped_table(1, Some(i64::MIN))]; let scheduler = soft_drop_scheduler(ctx.clone()); - scheduler.handle_manual_gc(None, None, None).await.unwrap(); + scheduler + .handle_manual_gc(None, None, None, ProcedureContext::default()) + .await + .unwrap(); assert!(ctx.purge_attempts.lock().unwrap().is_empty()); } @@ -898,7 +907,7 @@ mod tests { _full_file_listing: bool, _timeout: Duration, _region_routes_override: Region2Peers, - _event_context: PersistentEventContext, + _procedure_context: ProcedureContext, ) -> Result { crate::error::UnexpectedSnafu { violated: "mock gc failure".to_string(), @@ -949,6 +958,7 @@ mod tests { Some(vec![RegionId::new(1, 0)]), Some(false), Some(Duration::from_secs(1)), + ProcedureContext::default(), ) .await; @@ -977,6 +987,7 @@ mod tests { Some(vec![RegionId::new(1, 0)]), Some(false), Some(Duration::from_secs(1)), + ProcedureContext::default(), ) .await; diff --git a/src/meta-srv/src/procedure/region_migration.rs b/src/meta-srv/src/procedure/region_migration.rs index ad1553c92d7..75ab5c79397 100644 --- a/src/meta-srv/src/procedure/region_migration.rs +++ b/src/meta-srv/src/procedure/region_migration.rs @@ -32,7 +32,7 @@ use std::fmt::{Debug, Display}; use std::time::Duration; use common_error::ext::BoxedError; -use common_event_recorder::Event; +use common_event_recorder::{Event, PersistentEventContext}; use common_meta::cache_invalidator::CacheInvalidatorRef; use common_meta::ddl::RegionFailureDetectorControllerRef; use common_meta::instruction::CacheIdent; @@ -45,7 +45,7 @@ use common_meta::kv_backend::{KvBackendRef, ResettableKvBackendRef}; use common_meta::lock_key::{CatalogLock, RegionLock, SchemaLock, TableLock}; use common_meta::peer::Peer; use common_meta::region_keeper::{MemoryRegionKeeperRef, OperatingRegionGuard}; -use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; +use common_meta::rpc::ddl::TriggerReason; use common_procedure::error::{ Error as ProcedureError, FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu, }; diff --git a/src/meta-srv/src/procedure/region_migration/manager.rs b/src/meta-srv/src/procedure/region_migration/manager.rs index e1d719d9747..c21de506243 100644 --- a/src/meta-srv/src/procedure/region_migration/manager.rs +++ b/src/meta-srv/src/procedure/region_migration/manager.rs @@ -18,12 +18,14 @@ use std::fmt::Display; use std::sync::{Arc, RwLock}; use std::time::Duration; +use common_event_recorder::PersistentEventContext; use common_meta::key::table_info::TableInfoValue; use common_meta::key::table_route::TableRouteValue; use common_meta::peer::Peer; -use common_meta::rpc::ddl::PersistentEventContext; use common_meta::rpc::router::RegionRoute; -use common_procedure::{ProcedureId, ProcedureManagerRef, ProcedureWithId, watcher}; +use common_procedure::{ + ProcedureContext, ProcedureId, ProcedureManagerRef, ProcedureWithId, watcher, +}; use common_telemetry::{error, info, warn}; use serde::{Deserialize, Serialize}; use snafu::{OptionExt, ResultExt, ensure}; @@ -458,8 +460,20 @@ impl RegionMigrationManager { } let submitting_region_ids = task.region_ids.clone(); + // TODO(weny): Pass the actor when region migration task submission propagates it. + let procedure_context = ProcedureContext { + actor: None, + event_context: Some(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + }; let procedure_id = self - .submit_procedure_inner(task, procedure_guards, catalog_and_schema) + .submit_procedure_inner( + procedure_context, + task, + procedure_guards, + catalog_and_schema, + ) .await?; Ok(SubmitRegionMigrationTaskResult { migrated, @@ -475,11 +489,11 @@ impl RegionMigrationManager { async fn submit_procedure_inner( &self, + procedure_context: ProcedureContext, task: RegionMigrationTaskBatch, procedure_guards: Vec, catalog_and_schema: Vec<(String, String)>, ) -> Result { - let event_context = PersistentEventContext::new(task.trigger_reason.to_trigger_reason()); let procedure = RegionMigrationProcedure::new( PersistentContext::new( catalog_and_schema, @@ -493,7 +507,7 @@ impl RegionMigrationManager { procedure_guards, ); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); let procedure_id = procedure_with_id.id; info!("Starting region migration procedure {procedure_id} for {task}"); let procedure_manager = self.procedure_manager.clone(); @@ -529,8 +543,14 @@ impl RegionMigrationManager { /// Submits a new region migration procedure. pub async fn submit_procedure( &self, + procedure_context: ProcedureContext, mut task: RegionMigrationProcedureTask, ) -> Result> { + if let Some(event_context) = procedure_context.event_context.as_ref() { + task.trigger_reason = + RegionMigrationTriggerReason::from_trigger_reason(event_context.reason); + } + let Some(guard) = self.insert_running_procedure(&task) else { return error::MigrationRunningSnafu { region_id: task.region_id, @@ -577,7 +597,6 @@ impl RegionMigrationManager { timeout, trigger_reason, } = task.clone(); - let event_context = PersistentEventContext::new(trigger_reason.to_trigger_reason()); let procedure = RegionMigrationProcedure::new( PersistentContext::new( vec![(catalog_name, schema_name)], @@ -591,7 +610,7 @@ impl RegionMigrationManager { vec![guard], ); let procedure_with_id = - ProcedureWithId::with_random_id(Box::new(procedure)).with_event_context(event_context); + ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); let procedure_id = procedure_with_id.id; info!("Starting region migration procedure {procedure_id} for {task}"); let procedure_manager = self.procedure_manager.clone(); @@ -655,7 +674,15 @@ mod test { .unwrap() .insert(region_id, task.clone()); - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::MigrationRunning { .. }); } @@ -673,7 +700,15 @@ mod test { trigger_reason: RegionMigrationTriggerReason::Manual, }; - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::InvalidArguments { .. }); } @@ -691,7 +726,15 @@ mod test { trigger_reason: RegionMigrationTriggerReason::Manual, }; - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::TableRouteNotFound { .. }); } @@ -719,7 +762,15 @@ mod test { env.create_physical_table_metadata(table_info, region_routes) .await; - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::RegionRouteNotFound { .. }); } @@ -747,7 +798,15 @@ mod test { env.create_physical_table_metadata(table_info, region_routes) .await; - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::LeaderPeerChanged { .. }); assert_eq!( err.to_string(), @@ -780,7 +839,15 @@ mod test { env.create_physical_table_metadata(table_info, region_routes) .await; - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::InvalidArguments { .. }); assert_eq!( err.to_string(), @@ -813,7 +880,15 @@ mod test { env.create_physical_table_metadata(table_info, region_routes) .await; - let err = manager.submit_procedure(task).await.unwrap_err(); + let err = manager + .submit_procedure( + ProcedureContext::from_event_context(PersistentEventContext::new( + task.trigger_reason.to_trigger_reason(), + )), + task, + ) + .await + .unwrap_err(); assert_matches!(err, error::Error::RegionMigrated { .. }); } diff --git a/src/meta-srv/src/service/procedure.rs b/src/meta-srv/src/service/procedure.rs index 1b5c4c42dbc..9e226bc18fe 100644 --- a/src/meta-srv/src/service/procedure.rs +++ b/src/meta-srv/src/service/procedure.rs @@ -18,10 +18,12 @@ use api::v1::meta::reconcile_request::Target; use api::v1::meta::{ DdlTaskRequest as PbDdlTaskRequest, DdlTaskResponse as PbDdlTaskResponse, GcRegionsRequest, GcRegionsResponse, GcStats, GcTableRequest, GcTableResponse, MigrateRegionRequest, - MigrateRegionResponse, ProcedureDetailRequest, ProcedureDetailResponse, ProcedureStateResponse, + MigrateRegionResponse, ProcedureActor, ProcedureDetailRequest, ProcedureDetailResponse, + ProcedureEventContext as PbProcedureEventContext, ProcedureStateResponse, QueryProcedureRequest, ReconcileCatalog, ReconcileDatabase, ReconcileRequest, ReconcileResponse, ReconcileTable, ResolveStrategy, Role, procedure_service_server, }; +use common_event_recorder::{PersistentEventContext, ProcedureEventInput}; use common_meta::key::TableMetadataManagerRef; use common_meta::key::table_name::TableNameKey; use common_meta::procedure_executor::ExecutorContext; @@ -33,6 +35,7 @@ use common_meta::rpc::procedure::{ self, GcRegionsRequest as MetaGcRegionsRequest, GcResponse, GcTableRequest as MetaGcTableRequest, }; +use common_procedure::ProcedureContext; use snafu::{OptionExt, ResultExt}; use store_api::storage::RegionId; use table::table_reference::TableReference; @@ -47,6 +50,45 @@ use crate::procedure::region_migration::manager::{ use crate::service::GrpcResult; use crate::{check_leader, error, gc}; +struct ProcedureSubmission { + actor: Option, + event_context: Option, +} + +impl TryFrom<(Option, Option)> for ProcedureSubmission { + type Error = Status; + + fn try_from( + (actor, event_context): (Option, Option), + ) -> Result { + let actor = actor + .map(|actor| { + if actor.username.is_empty() { + Err(Status::invalid_argument( + "Procedure actor must not be empty", + )) + } else { + Ok(actor.username) + } + }) + .transpose()?; + + Ok(Self { + actor, + event_context, + }) + } +} + +impl From for ProcedureContext { + fn from(context: ProcedureSubmission) -> Self { + Self { + actor: context.actor, + event_context: context.event_context.map(PersistentEventContext::from), + } + } +} + #[async_trait::async_trait] impl procedure_service_server::ProcedureService for Metasrv { async fn query( @@ -89,9 +131,15 @@ impl procedure_service_server::ProcedureService for Metasrv { task, wait, timeout_secs, + event_context, + actor, } = request; let header = header.context(error::MissingRequestHeaderSnafu)?; + let ProcedureSubmission { + actor, + event_context, + } = ProcedureSubmission::try_from((actor, event_context))?; let mut query_context = query_context .context(error::MissingRequiredParameterSnafu { param: "query_context", @@ -102,15 +150,17 @@ impl procedure_service_server::ProcedureService for Metasrv { .try_into() .context(error::ConvertProtoDataSnafu)?; restore_create_database_creator(&metadata, header.role, &mut task, &mut query_context)?; + let executor_context = ExecutorContext { + tracing_context: Some(header.tracing_context), + query_context: Some(query_context), + actor, + event_input: event_context.map(ProcedureEventInput::from), + }; let resp = self .ddl_manager() .submit_ddl_task( - &ExecutorContext { - tracing_context: Some(header.tracing_context), - ..Default::default() - }, + executor_context, SubmitDdlTaskRequest { - query_context, wait, timeout: Duration::from_secs(timeout_secs.into()), task, @@ -135,9 +185,13 @@ impl procedure_service_server::ProcedureService for Metasrv { from_peer, to_peer, timeout_secs, + event_context, + actor, } = request.into_inner(); let _header = header.context(error::MissingRequestHeaderSnafu)?; + let procedure_context = + ProcedureContext::from(ProcedureSubmission::try_from((actor, event_context))?); let from_peer = self .lookup_datanode_peer(from_peer) .await? @@ -149,13 +203,16 @@ impl procedure_service_server::ProcedureService for Metasrv { let pid = self .region_migration_manager() - .submit_procedure(RegionMigrationProcedureTask { - region_id: region_id.into(), - from_peer, - to_peer, - timeout: Duration::from_secs(timeout_secs.into()), - trigger_reason: RegionMigrationTriggerReason::Manual, - }) + .submit_procedure( + procedure_context, + RegionMigrationProcedureTask { + region_id: region_id.into(), + from_peer, + to_peer, + timeout: Duration::from_secs(timeout_secs.into()), + trigger_reason: RegionMigrationTriggerReason::Manual, + }, + ) .await? .map(procedure::pid_to_pb_pid); @@ -266,16 +323,23 @@ impl procedure_service_server::ProcedureService for Metasrv { region_ids, full_file_listing, timeout_secs, + event_context, + actor, } = request.into_inner(); let _header = header.context(error::MissingRequestHeaderSnafu)?; + let procedure_context = + ProcedureContext::from(ProcedureSubmission::try_from((actor, event_context))?); let response = self - .handle_gc_regions(MetaGcRegionsRequest { - region_ids, - full_file_listing, - timeout: Self::normalize_gc_timeout(Duration::from_secs(timeout_secs as u64)), - }) + .handle_gc_regions( + procedure_context, + MetaGcRegionsRequest { + region_ids, + full_file_listing, + timeout: Self::normalize_gc_timeout(Duration::from_secs(timeout_secs as u64)), + }, + ) .await?; Ok(Response::new(gc_response_to_regions_pb(response))) @@ -291,18 +355,25 @@ impl procedure_service_server::ProcedureService for Metasrv { table_name, full_file_listing, timeout_secs, + event_context, + actor, } = request.into_inner(); let _header = header.context(error::MissingRequestHeaderSnafu)?; + let procedure_context = + ProcedureContext::from(ProcedureSubmission::try_from((actor, event_context))?); let response = self - .handle_gc_table(MetaGcTableRequest { - catalog_name, - schema_name, - table_name, - full_file_listing, - timeout: Self::normalize_gc_timeout(Duration::from_secs(timeout_secs as u64)), - }) + .handle_gc_table( + procedure_context, + MetaGcTableRequest { + catalog_name, + schema_name, + table_name, + full_file_listing, + timeout: Self::normalize_gc_timeout(Duration::from_secs(timeout_secs as u64)), + }, + ) .await?; Ok(Response::new(gc_response_to_table_pb(response))) @@ -371,17 +442,30 @@ impl Metasrv { } } - async fn handle_gc_regions(&self, request: MetaGcRegionsRequest) -> error::Result { + async fn handle_gc_regions( + &self, + procedure_context: ProcedureContext, + request: MetaGcRegionsRequest, + ) -> error::Result { let region_ids: Vec = request .region_ids .into_iter() .map(RegionId::from_u64) .collect(); - self.trigger_gc_for_regions(region_ids, request.full_file_listing, request.timeout) - .await + self.trigger_gc_for_regions( + procedure_context, + region_ids, + request.full_file_listing, + request.timeout, + ) + .await } - async fn handle_gc_table(&self, request: MetaGcTableRequest) -> error::Result { + async fn handle_gc_table( + &self, + procedure_context: ProcedureContext, + request: MetaGcTableRequest, + ) -> error::Result { let table_name_key = TableNameKey::new( &request.catalog_name, &request.schema_name, @@ -406,13 +490,19 @@ impl Metasrv { .context(TableMetadataManagerSnafu)?; let region_ids: Vec = route.region_routes.iter().map(|r| r.region.id).collect(); - self.trigger_gc_for_regions(region_ids, request.full_file_listing, request.timeout) - .await + self.trigger_gc_for_regions( + procedure_context, + region_ids, + request.full_file_listing, + request.timeout, + ) + .await } /// Triggers manual GC for specified regions and returns the GC response. async fn trigger_gc_for_regions( &self, + procedure_context: ProcedureContext, region_ids: Vec, full_file_listing: bool, timeout: Option, @@ -429,6 +519,7 @@ impl Metasrv { region_ids: Some(region_ids), full_file_listing: Some(full_file_listing), timeout, + procedure_context, }) .await .map_err(|_| { @@ -502,14 +593,15 @@ fn gc_response_to_table_pb(resp: GcResponse) -> GcTableResponse { mod tests { use std::time::Duration; - use api::v1::meta::Role; + use api::v1::meta::{ProcedureActor, ProcedureEventContext, Role}; use common_meta::rpc::ddl::{ CREATE_DATABASE_CREATOR_EXTENSION_KEY, CREATE_DATABASE_CREATOR_METADATA_KEY, CreatorGrantIntent, DdlTask, QueryContext, }; + use common_procedure::ProcedureContext; use tonic::metadata::{MetadataMap, MetadataValue}; - use super::{Metasrv, restore_create_database_creator}; + use super::{Metasrv, ProcedureSubmission, restore_create_database_creator}; fn create_database_task() -> DdlTask { DdlTask::new_create_database( @@ -646,4 +738,41 @@ mod tests { assert_eq!(task.creator, Some(creator)); assert!(trusted.extensions.is_empty()); } + + #[test] + fn test_procedure_submission_rejects_empty_actor() { + let submission = ProcedureSubmission::try_from(( + Some(ProcedureActor { + username: "alice".to_string(), + }), + Some(ProcedureEventContext { + reason: "manual".to_string(), + protocol: "postgres".to_string(), + extensions: Default::default(), + }), + )) + .unwrap(); + assert_eq!(submission.actor.as_deref(), Some("alice")); + assert_eq!( + submission + .event_context + .as_ref() + .map(|context| context.protocol.as_str()), + Some("postgres") + ); + + assert!( + ProcedureSubmission::try_from(( + Some(ProcedureActor { + username: String::new(), + }), + None, + )) + .is_err() + ); + + let procedure_context = + ProcedureContext::from(ProcedureSubmission::try_from((None, None)).unwrap()); + assert_eq!(procedure_context.event_context, None); + } } diff --git a/src/operator/src/procedure.rs b/src/operator/src/procedure.rs index 1a849afaf7a..03fb4f909c9 100644 --- a/src/operator/src/procedure.rs +++ b/src/operator/src/procedure.rs @@ -27,11 +27,12 @@ use common_meta::rpc::procedure::{ }; use common_query::error as query_error; use common_query::error::Result as QueryResult; +use session::context::QueryContextRef; use snafu::ResultExt; use table::table_name::TableName; use crate::error; -use crate::utils::to_meta_query_context_with_trigger_reason; +use crate::utils::to_executor_context; /// The operator for procedures which implements [`ProcedureServiceHandler`]. #[derive(Clone)] @@ -59,8 +60,8 @@ impl ProcedureServiceOperator { impl ProcedureServiceHandler for ProcedureServiceOperator { async fn purge_table( &self, + query_ctx: QueryContextRef, table_name: TableName, - query_ctx: session::context::QueryContextRef, ) -> QueryResult<()> { let dropped = self .table_metadata_manager @@ -76,24 +77,25 @@ impl ProcedureServiceHandler for ProcedureServiceOperator { }) .map_err(BoxedError::new) .context(query_error::ProcedureServiceSnafu)?; - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_ctx, TriggerReason::Manual); - let request = SubmitDdlTaskRequest::new( - meta_query_context, - DdlTask::new_purge_dropped_table(dropped.table_id), - ); + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_purge_dropped_table(dropped.table_id)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .map_err(BoxedError::new) .context(query_error::ProcedureServiceSnafu)?; Ok(()) } - async fn migrate_region(&self, request: MigrateRegionRequest) -> QueryResult> { + async fn migrate_region( + &self, + query_ctx: QueryContextRef, + request: MigrateRegionRequest, + ) -> QueryResult> { + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); Ok(self .procedure_executor - .migrate_region(&ExecutorContext::default(), request) + .migrate_region(&executor_context, request) .await .map_err(BoxedError::new) .context(query_error::ProcedureServiceSnafu)? @@ -135,17 +137,27 @@ impl ProcedureServiceHandler for ProcedureServiceOperator { &self.catalog_manager } - async fn gc_regions(&self, request: MetaGcRegionsRequest) -> QueryResult { + async fn gc_regions( + &self, + query_ctx: QueryContextRef, + request: MetaGcRegionsRequest, + ) -> QueryResult { + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); self.procedure_executor - .gc_regions(&ExecutorContext::default(), request) + .gc_regions(&executor_context, request) .await .map_err(BoxedError::new) .context(query_error::ProcedureServiceSnafu) } - async fn gc_table(&self, request: MetaGcTableRequest) -> QueryResult { + async fn gc_table( + &self, + query_ctx: QueryContextRef, + request: MetaGcTableRequest, + ) -> QueryResult { + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); self.procedure_executor - .gc_table(&ExecutorContext::default(), request) + .gc_table(&executor_context, request) .await .map_err(BoxedError::new) .context(query_error::ProcedureServiceSnafu) @@ -162,10 +174,10 @@ mod tests { use common_meta::key::table_route::TableRouteValue; use common_meta::key::test_utils::new_test_table_info_with_name; use common_meta::kv_backend::memory::MemoryKvBackend; - use common_meta::procedure_executor::{ProcedureExecutor, ProcedureExecutorRef}; - use common_meta::rpc::ddl::{ - DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse, TRIGGER_REASON_EXTENSION_KEY, + use common_meta::procedure_executor::{ + ExecutorContext, ProcedureExecutor, ProcedureExecutorRef, }; + use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse}; use common_meta::rpc::procedure::{MigrateRegionResponse, ProcedureStateResponse}; use session::context::QueryContextBuilder; use table::table_name::TableName; @@ -175,6 +187,7 @@ mod tests { #[derive(Default)] struct RecordingProcedureExecutor { requests: Mutex>, + contexts: Mutex>, fail: bool, } @@ -182,9 +195,10 @@ mod tests { impl ProcedureExecutor for RecordingProcedureExecutor { async fn submit_ddl_task( &self, - _: &ExecutorContext, + context: ExecutorContext, request: SubmitDdlTaskRequest, ) -> common_meta::error::Result { + self.contexts.lock().unwrap().push(context); self.requests.lock().unwrap().push(request); if self.fail { return common_meta::error::UnsupportedSnafu { @@ -273,20 +287,19 @@ mod tests { .build() .into(); - operator.purge_table(name.clone(), query_ctx).await.unwrap(); + operator.purge_table(query_ctx, name.clone()).await.unwrap(); let requests = executor.requests.lock().unwrap(); assert!( matches!(&requests[0].task, DdlTask::PurgeDroppedTable(task) if task.table_id == 42) ); - assert_eq!(requests[0].query_context.current_catalog, "catalog"); - assert_eq!(requests[0].query_context.current_schema, "schema"); + let contexts = executor.contexts.lock().unwrap(); + let query_context = contexts[0].query_context.as_ref().unwrap(); + assert_eq!(query_context.current_catalog, "catalog"); + assert_eq!(query_context.current_schema, "schema"); assert_eq!( - Some(&TriggerReason::Manual.as_ref().to_string()), - requests[0] - .query_context - .extensions - .get(TRIGGER_REASON_EXTENSION_KEY) + contexts[0].event_input.as_ref().map(|input| input.reason), + Some(TriggerReason::Manual) ); } @@ -303,12 +316,12 @@ mod tests { ); let error = operator .purge_table( - TableName::new("catalog", "schema", "missing"), QueryContextBuilder::default() .current_catalog("catalog".to_string()) .current_schema("schema".to_string()) .build() .into(), + TableName::new("catalog", "schema", "missing"), ) .await .unwrap_err(); @@ -326,12 +339,12 @@ mod tests { assert!( operator .purge_table( - name, QueryContextBuilder::default() .current_catalog("catalog".to_string()) .current_schema("schema".to_string()) .build() - .into() + .into(), + name, ) .await .is_err() diff --git a/src/operator/src/statement/comment.rs b/src/operator/src/statement/comment.rs index 824287fc80b..c55e1be9f74 100644 --- a/src/operator/src/statement/comment.rs +++ b/src/operator/src/statement/comment.rs @@ -15,7 +15,6 @@ use api::v1::CommentOnExpr; use common_error::ext::BoxedError; use common_meta::cache_invalidator::Context; -use common_meta::procedure_executor::ExecutorContext; use common_meta::rpc::ddl::{ CommentObjectType, CommentOnTask, DdlTask, SubmitDdlTaskRequest, TriggerReason, }; @@ -30,7 +29,7 @@ use crate::error::{ self, ExecuteDdlSnafu, ExternalSnafu, InvalidSqlSnafu, Result, TableMetadataManagerSnafu, }; use crate::statement::StatementExecutor; -use crate::utils::to_meta_query_context_with_trigger_reason; +use crate::utils::to_executor_context; impl StatementExecutor { /// Adds a comment to a database object (table, column, or flow). @@ -54,13 +53,11 @@ impl StatementExecutor { .context(TableMetadataManagerSnafu)?; let cache_idents = comment_on_task.cache_idents(); - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_ctx, TriggerReason::Manual); - let request = - SubmitDdlTaskRequest::new(meta_query_context, DdlTask::new_comment_on(comment_on_task)); + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_comment_on(comment_on_task)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(ExecuteDdlSnafu)?; @@ -88,13 +85,11 @@ impl StatementExecutor { .context(TableMetadataManagerSnafu)?; let cache_idents = comment_on_task.cache_idents(); - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_ctx, TriggerReason::Manual); - let request = - SubmitDdlTaskRequest::new(meta_query_context, DdlTask::new_comment_on(comment_on_task)); + let executor_context = to_executor_context(query_ctx, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_comment_on(comment_on_task)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(ExecuteDdlSnafu)?; diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 256b849123e..3c10186ca4e 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -47,7 +47,6 @@ use common_meta::instruction::CacheIdent; #[cfg(feature = "enterprise")] use common_meta::key::TableMetadataManagerRef; use common_meta::key::schema_name::{SchemaName, SchemaNameKey}; -use common_meta::procedure_executor::ExecutorContext; #[cfg(feature = "enterprise")] use common_meta::procedure_executor::ProcedureExecutorRef; #[cfg(feature = "enterprise")] @@ -119,10 +118,7 @@ use crate::error::{ use crate::expr_helper::{self, RepartitionRequest, RepartitionSource}; use crate::statement::StatementExecutor; use crate::statement::show::create_partitions_stmt; -use crate::utils::{ - to_meta_query_context_with_origin_frontend_and_trigger_reason, - to_meta_query_context_with_trigger_reason, -}; +use crate::utils::{to_executor_context, to_executor_context_with_origin_frontend}; #[derive(Debug, Clone, Copy)] struct DdlSubmitOptions { @@ -729,13 +725,11 @@ impl StatementExecutor { }) .context(error::InvalidExprSnafu)?; - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_create_trigger(task), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_create_trigger(task)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -793,13 +787,11 @@ impl StatementExecutor { create_flow: Some(expr), }) .context(error::InvalidExprSnafu)?; - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_create_flow(task), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_create_flow(task)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -1109,14 +1101,12 @@ impl StatementExecutor { table_type: TableType::View, }; - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(ctx, TriggerReason::Manual), - DdlTask::new_create_view(expr, view_info.clone()), - ); + let executor_context = to_executor_context(ctx, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_create_view(expr, view_info.clone())); let resp = self .procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu)?; @@ -1195,13 +1185,11 @@ impl StatementExecutor { expr: DropFlowTask, query_context: QueryContextRef, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_drop_flow(expr), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_drop_flow(expr)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -1230,13 +1218,11 @@ impl StatementExecutor { expr: DropTriggerTask, query_context: QueryContextRef, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_drop_trigger(expr), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_drop_trigger(expr)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -1311,13 +1297,11 @@ impl StatementExecutor { expr: DropViewTask, query_context: QueryContextRef, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_drop_view(expr), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_drop_view(expr)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -1780,17 +1764,13 @@ impl StatementExecutor { source: Some(source), ..Default::default() }; - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_context.clone(), TriggerReason::Manual); - let mut req = SubmitDdlTaskRequest::new( - meta_query_context, - DdlTask::new_alter_table(AlterTableExpr { - catalog_name: request.catalog_name.clone(), - schema_name: request.schema_name.clone(), - table_name: request.table_name.clone(), - kind: Some(Kind::Repartition(repartition)), - }), - ); + let executor_context = to_executor_context(query_context.clone(), TriggerReason::Manual); + let mut req = SubmitDdlTaskRequest::new(DdlTask::new_alter_table(AlterTableExpr { + catalog_name: request.catalog_name.clone(), + schema_name: request.schema_name.clone(), + table_name: request.table_name.clone(), + kind: Some(Kind::Repartition(repartition)), + })); req.wait = ddl_options.wait; req.timeout = ddl_options.timeout; @@ -1806,7 +1786,7 @@ impl StatementExecutor { let response = self .procedure_executor - .submit_ddl_task(&ExecutorContext::default(), req) + .submit_ddl_task(executor_context, req) .await .context(error::ExecuteDdlSnafu)?; @@ -1888,12 +1868,11 @@ impl StatementExecutor { .await .context(TableMetadataManagerSnafu)?; - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_context, trigger_reason); + let executor_context = to_executor_context(query_context, trigger_reason); let (req, invalidate_keys) = if physical_table_id == table_id { // This is physical table - let req = SubmitDdlTaskRequest::new(meta_query_context, DdlTask::new_alter_table(expr)); + let req = SubmitDdlTaskRequest::new(DdlTask::new_alter_table(expr)); let invalidate_keys = vec![ CacheIdent::TableId(table_id), @@ -1903,10 +1882,7 @@ impl StatementExecutor { (req, invalidate_keys) } else { // This is logical table - let req = SubmitDdlTaskRequest::new( - meta_query_context, - DdlTask::new_alter_logical_tables(vec![expr]), - ); + let req = SubmitDdlTaskRequest::new(DdlTask::new_alter_logical_tables(vec![expr])); let mut invalidate_keys = vec![ CacheIdent::TableId(physical_table_id), @@ -1934,7 +1910,7 @@ impl StatementExecutor { }; self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), req) + .submit_ddl_task(executor_context, req) .await .context(error::ExecuteDdlSnafu)?; @@ -2025,17 +2001,19 @@ impl StatementExecutor { .map(|expr| expr.as_pb_partition().context(PartitionExprToPbSnafu)) .collect::>>()?; - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_origin_frontend_and_trigger_reason( - query_context, - &self.origin_frontend_addr, - trigger_reason, - ), - DdlTask::new_create_table(create_table, partitions, table_info), + let executor_context = to_executor_context_with_origin_frontend( + query_context, + &self.origin_frontend_addr, + trigger_reason, ); + let request = SubmitDdlTaskRequest::new(DdlTask::new_create_table( + create_table, + partitions, + table_info, + )); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2046,17 +2024,15 @@ impl StatementExecutor { query_context: QueryContextRef, trigger_reason: TriggerReason, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_origin_frontend_and_trigger_reason( - query_context, - &self.origin_frontend_addr, - trigger_reason, - ), - DdlTask::new_create_logical_tables(tables_data), + let executor_context = to_executor_context_with_origin_frontend( + query_context, + &self.origin_frontend_addr, + trigger_reason, ); + let request = SubmitDdlTaskRequest::new(DdlTask::new_create_logical_tables(tables_data)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2067,15 +2043,11 @@ impl StatementExecutor { query_context: QueryContextRef, trigger_reason: TriggerReason, ) -> Result { - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_context, trigger_reason); - let request = SubmitDdlTaskRequest::new( - meta_query_context, - DdlTask::new_alter_logical_tables(tables_data), - ); + let executor_context = to_executor_context(query_context, trigger_reason); + let request = SubmitDdlTaskRequest::new(DdlTask::new_alter_logical_tables(tables_data)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2087,19 +2059,17 @@ impl StatementExecutor { drop_if_exists: bool, query_context: QueryContextRef, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_drop_table( - table_name.catalog_name.clone(), - table_name.schema_name.clone(), - table_name.table_name.clone(), - table_id, - drop_if_exists, - ), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_drop_table( + table_name.catalog_name.clone(), + table_name.schema_name.clone(), + table_name.table_name.clone(), + table_id, + drop_if_exists, + )); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2111,13 +2081,12 @@ impl StatementExecutor { drop_if_exists: bool, query_context: QueryContextRef, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_drop_database(catalog, schema, drop_if_exists), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = + SubmitDdlTaskRequest::new(DdlTask::new_drop_database(catalog, schema, drop_if_exists)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2127,13 +2096,11 @@ impl StatementExecutor { alter_expr: AlterDatabaseExpr, query_context: QueryContextRef, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_alter_database(alter_expr), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_alter_database(alter_expr)); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2145,21 +2112,17 @@ impl StatementExecutor { time_ranges: Vec<(Timestamp, Timestamp)>, query_context: QueryContextRef, ) -> Result { - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual); - let request = SubmitDdlTaskRequest::new( - meta_query_context, - DdlTask::new_truncate_table( - table_name.catalog_name.clone(), - table_name.schema_name.clone(), - table_name.table_name.clone(), - table_id, - time_ranges, - ), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_truncate_table( + table_name.catalog_name.clone(), + table_name.schema_name.clone(), + table_name.table_name.clone(), + table_id, + time_ranges, + )); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2239,13 +2202,17 @@ impl StatementExecutor { query_context: QueryContextRef, creator: Option, ) -> Result { - let request = SubmitDdlTaskRequest::new( - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual), - DdlTask::new_create_database(catalog, database, create_if_not_exists, options, creator), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_create_database( + catalog, + database, + create_if_not_exists, + options, + creator, + )); self.procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu) } @@ -2876,14 +2843,10 @@ async fn execute_undrop_table( table_name: table_name.to_string(), })?; - let meta_query_context = - to_meta_query_context_with_trigger_reason(query_context, TriggerReason::Manual); - let request = SubmitDdlTaskRequest::new( - meta_query_context, - DdlTask::new_undrop_table(dropped.table_id), - ); + let executor_context = to_executor_context(query_context, TriggerReason::Manual); + let request = SubmitDdlTaskRequest::new(DdlTask::new_undrop_table(dropped.table_id)); procedure_executor - .submit_ddl_task(&ExecutorContext::default(), request) + .submit_ddl_task(executor_context, request) .await .context(error::ExecuteDdlSnafu)?; @@ -2931,9 +2894,7 @@ mod test { ExecutorContext, ProcedureExecutor, ProcedureExecutorRef, }; #[cfg(feature = "enterprise")] - use common_meta::rpc::ddl::{ - DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse, TRIGGER_REASON_EXTENSION_KEY, - }; + use common_meta::rpc::ddl::{DdlTask, SubmitDdlTaskRequest, SubmitDdlTaskResponse}; #[cfg(feature = "enterprise")] use common_meta::rpc::procedure::{ MigrateRegionRequest, MigrateRegionResponse, ProcedureStateResponse, @@ -2951,6 +2912,7 @@ mod test { #[derive(Default)] struct RecordingProcedureExecutor { requests: Mutex>, + contexts: Mutex>, fail: bool, } @@ -2959,9 +2921,10 @@ mod test { impl ProcedureExecutor for RecordingProcedureExecutor { async fn submit_ddl_task( &self, - _ctx: &ExecutorContext, + ctx: ExecutorContext, request: SubmitDdlTaskRequest, ) -> common_meta::error::Result { + self.contexts.lock().unwrap().push(ctx); self.requests.lock().unwrap().push(request); if self.fail { return common_meta::error::UnsupportedSnafu { @@ -3078,12 +3041,10 @@ mod test { let requests = procedure.requests.lock().unwrap(); assert!(matches!(&requests[0].task, DdlTask::UndropTable(task) if task.table_id == 42)); + let contexts = procedure.contexts.lock().unwrap(); assert_eq!( - requests[0] - .query_context - .extensions - .get(TRIGGER_REASON_EXTENSION_KEY), - Some(&TriggerReason::Manual.as_ref().to_string()) + contexts[0].event_input.as_ref().map(|input| input.reason), + Some(TriggerReason::Manual) ); assert_eq!( cache.invalidations.lock().unwrap()[0], diff --git a/src/operator/src/utils.rs b/src/operator/src/utils.rs index 03649713b0b..160d7120726 100644 --- a/src/operator/src/utils.rs +++ b/src/operator/src/utils.rs @@ -14,9 +14,9 @@ use std::sync::{Arc, RwLock}; -use common_meta::rpc::ddl::{ - ORIGIN_FRONTEND_ADDR_EXTENSION_KEY, TRIGGER_REASON_EXTENSION_KEY, TriggerReason, -}; +use common_event_recorder::ProcedureEventInput; +use common_meta::procedure_executor::ExecutorContext; +use common_meta::rpc::ddl::{ORIGIN_FRONTEND_ADDR_EXTENSION_KEY, TriggerReason}; use common_time::Timezone; use session::context::{QueryContextBuilder, QueryContextRef}; use snafu::ResultExt; @@ -37,31 +37,34 @@ pub fn to_meta_query_context( } } -/// Converts a session query context and records the DDL trigger reason. -pub fn to_meta_query_context_with_trigger_reason( +/// Builds the execution envelope for a frontend-originated procedure. +pub(crate) fn to_executor_context( query_context: QueryContextRef, trigger_reason: TriggerReason, -) -> common_meta::rpc::ddl::QueryContext { - let mut meta_query_context = to_meta_query_context(query_context); - meta_query_context.extensions.insert( - TRIGGER_REASON_EXTENSION_KEY.to_string(), - trigger_reason.as_ref().to_string(), - ); - meta_query_context +) -> ExecutorContext { + let actor = query_context.current_user().username().to_string(); + ExecutorContext { + query_context: Some(to_meta_query_context(query_context)), + actor: Some(actor), + event_input: Some(ProcedureEventInput::new(trigger_reason)), + ..Default::default() + } } -pub fn to_meta_query_context_with_origin_frontend_and_trigger_reason( +/// Builds a frontend-originated procedure envelope with the trusted frontend origin. +pub(crate) fn to_executor_context_with_origin_frontend( query_context: QueryContextRef, origin_frontend_addr: &str, trigger_reason: TriggerReason, -) -> common_meta::rpc::ddl::QueryContext { - let mut meta_query_context = - to_meta_query_context_with_trigger_reason(query_context, trigger_reason); - meta_query_context.extensions.insert( - ORIGIN_FRONTEND_ADDR_EXTENSION_KEY.to_string(), - origin_frontend_addr.to_string(), - ); - meta_query_context +) -> ExecutorContext { + let mut executor_context = to_executor_context(query_context, trigger_reason); + if let Some(query_context) = &mut executor_context.query_context { + query_context.extensions.insert( + ORIGIN_FRONTEND_ADDR_EXTENSION_KEY.to_string(), + origin_frontend_addr.to_string(), + ); + } + executor_context } pub fn try_to_session_query_context( @@ -87,14 +90,12 @@ mod tests { use std::collections::HashMap; use std::sync::{Arc, RwLock}; - use common_meta::rpc::ddl::{ - ORIGIN_FRONTEND_ADDR_EXTENSION_KEY, TRIGGER_REASON_EXTENSION_KEY, TriggerReason, - }; + use common_meta::rpc::ddl::{ORIGIN_FRONTEND_ADDR_EXTENSION_KEY, TriggerReason}; use common_time::Timezone; use session::context::QueryContextBuilder; use super::{ - to_meta_query_context, to_meta_query_context_with_origin_frontend_and_trigger_reason, + to_executor_context_with_origin_frontend, to_meta_query_context, try_to_session_query_context, }; @@ -122,7 +123,7 @@ mod tests { } #[test] - fn test_meta_query_context_with_origin_frontend_overrides_reserved_key() { + fn test_executor_context_with_origin_frontend_overrides_reserved_key() { let session_ctx = Arc::new( QueryContextBuilder::default() .set_extension( @@ -132,11 +133,12 @@ mod tests { .build(), ); - let meta_ctx = to_meta_query_context_with_origin_frontend_and_trigger_reason( + let executor_context = to_executor_context_with_origin_frontend( session_ctx, "127.0.0.1:4000", TriggerReason::Manual, ); + let meta_ctx = executor_context.query_context.as_ref().unwrap(); assert_eq!( meta_ctx @@ -146,11 +148,8 @@ mod tests { Some("127.0.0.1:4000") ); assert_eq!( - meta_ctx - .extensions - .get(TRIGGER_REASON_EXTENSION_KEY) - .map(String::as_str), - Some("manual") + executor_context.event_input.unwrap().reason, + TriggerReason::Manual ); } } diff --git a/tests-integration/tests/gc_event.rs b/tests-integration/tests/gc_event.rs index 9ecedc77775..54a0787d12a 100644 --- a/tests-integration/tests/gc_event.rs +++ b/tests-integration/tests/gc_event.rs @@ -15,8 +15,8 @@ use std::sync::Arc; use std::time::Duration; -use common_meta::rpc::ddl::{PersistentEventContext, TriggerReason}; -use common_procedure::{ProcedureId, ProcedureWithId, watcher}; +use common_event_recorder::{PersistentEventContext, TriggerReason}; +use common_procedure::{ProcedureContext, ProcedureId, ProcedureWithId, watcher}; use common_test_util::temp_dir::create_temp_dir; use meta_srv::gc::{BatchGcProcedure, GcSchedulerOptions}; use mito2::gc::GcConfig; @@ -137,9 +137,9 @@ async fn test_batch_gc_event() { .submit(ProcedureWithId { id: procedure_id, procedure: Box::new(procedure), - context: common_procedure::ProcedureContext::from_event_context( - PersistentEventContext::new(TriggerReason::Manual), - ), + context: ProcedureContext::from_event_context(PersistentEventContext::new( + TriggerReason::Manual, + )), }) .await .unwrap(); diff --git a/tests-integration/tests/region_migration.rs b/tests-integration/tests/region_migration.rs index 3d27dabc3f7..d67180249c9 100644 --- a/tests-integration/tests/region_migration.rs +++ b/tests-integration/tests/region_migration.rs @@ -23,10 +23,11 @@ use common_event_recorder::event_table::{ }; use common_event_recorder::{ DEFAULT_EVENTS_TABLE_NAME, DEFAULT_FLUSH_INTERVAL_SECONDS, EVENTS_TABLE_TIMESTAMP_COLUMN_NAME, - EVENTS_TABLE_TYPE_COLUMN_NAME, + EVENTS_TABLE_TYPE_COLUMN_NAME, PersistentEventContext, TriggerReason, }; use common_meta::key::{RegionDistribution, RegionRoleSet, TableMetadataManagerRef}; use common_meta::peer::Peer; +use common_procedure::ProcedureContext; use common_procedure::event::{ EVENTS_TABLE_PROCEDURE_ID_COLUMN_NAME, EVENTS_TABLE_PROCEDURE_STATE_COLUMN_NAME, }; @@ -183,13 +184,18 @@ pub async fn test_region_migration(store_type: StorageType, endpoints: Vec Date: Thu, 13 Aug 2026 07:52:44 +0000 Subject: [PATCH 39/92] perf(mito2): optimize dictionary primary key sorting (#8767) Signed-off-by: lyang24 --- src/mito2/benches/memtable_bench.rs | 117 +++++++++++- src/mito2/src/memtable/bulk/part.rs | 285 +++++++++++++++++++++++++++- 2 files changed, 394 insertions(+), 8 deletions(-) diff --git a/src/mito2/benches/memtable_bench.rs b/src/mito2/benches/memtable_bench.rs index c9151eea4ca..80331941865 100644 --- a/src/mito2/benches/memtable_bench.rs +++ b/src/mito2/benches/memtable_bench.rs @@ -22,10 +22,17 @@ use std::sync::Arc; -use criterion::{Criterion, criterion_group, criterion_main}; +use common_recordbatch::DfRecordBatch; +use criterion::{Criterion, Throughput, criterion_group, criterion_main}; +use datatypes::arrow::array::{ + ArrayRef, BinaryDictionaryBuilder, Int64Array, TimestampMillisecondArray, UInt8Array, + UInt64Array, +}; +use datatypes::arrow::compute::{SortColumn, SortOptions}; +use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit, UInt32Type}; use mito_codec::row_converter::DensePrimaryKeyCodec; use mito2::memtable::bulk::context::BulkIterContext; -use mito2::memtable::bulk::part::BulkPartConverter; +use mito2::memtable::bulk::part::{BulkPartConverter, sort_primary_key_record_batch}; use mito2::memtable::bulk::part_reader::BulkPartBatchIter; use mito2::memtable::bulk::{BulkMemtable, BulkMemtableConfig}; use mito2::memtable::time_series::TimeSeriesMemtable; @@ -39,6 +46,111 @@ use mito2::test_util::memtable_util; const DEFAULT_BATCH_SIZE: usize = 8 * 1024; +fn primary_key_sort_batch( + series_count: usize, + samples_per_series: usize, + sorted_within_series: bool, +) -> DfRecordBatch { + let num_rows = series_count * samples_per_series; + let primary_keys = (0..series_count) + .map(|series| format!("series_{series:08}").into_bytes()) + .collect::>(); + let mut primary_key_builder = BinaryDictionaryBuilder::::new(); + let mut timestamps = Vec::with_capacity(num_rows); + let mut sequences = Vec::with_capacity(num_rows); + + for series_order in 0..series_count { + // 17 is coprime to the power-of-two cardinalities used below and gives a deterministic + // non-lexical series order. + let series = series_order * 17 % series_count; + for sample_order in 0..samples_per_series { + let sample = if sorted_within_series { + sample_order + } else { + sample_order * 37 % samples_per_series + }; + primary_key_builder + .append(primary_keys[series].as_slice()) + .unwrap(); + timestamps.push(sample as i64); + sequences.push((series * samples_per_series + sample_order) as u64); + } + } + + let columns: Vec = vec![ + Arc::new(Int64Array::from_iter_values(0..num_rows as i64)), + Arc::new(TimestampMillisecondArray::from(timestamps)), + Arc::new(primary_key_builder.finish()), + Arc::new(UInt64Array::from(sequences)), + Arc::new(UInt8Array::from_value(1, num_rows)), + ]; + let schema = Arc::new(Schema::new(vec![ + Field::new("value", DataType::Int64, false), + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new( + "__primary_key", + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Binary)), + false, + ), + Field::new("__sequence", DataType::UInt64, false), + Field::new("__op_type", DataType::UInt8, false), + ])); + DfRecordBatch::try_new(schema, columns).unwrap() +} + +fn lexsort_primary_key_record_batch(batch: &DfRecordBatch) -> DfRecordBatch { + let total_columns = batch.num_columns(); + let sort_columns = vec![ + SortColumn { + values: batch.column(total_columns - 3).clone(), + options: Some(SortOptions { + descending: false, + nulls_first: true, + }), + }, + SortColumn { + values: batch.column(total_columns - 4).clone(), + options: Some(SortOptions { + descending: false, + nulls_first: true, + }), + }, + SortColumn { + values: batch.column(total_columns - 2).clone(), + options: Some(SortOptions { + descending: true, + nulls_first: true, + }), + }, + ]; + let indices = datatypes::arrow::compute::lexsort_to_indices(&sort_columns, None).unwrap(); + datatypes::arrow::compute::take_record_batch(batch, &indices).unwrap() +} + +fn primary_key_sort(c: &mut Criterion) { + let mut group = c.benchmark_group("primary_key_sort"); + group.sample_size(20); + + for (name, series_count, samples_per_series, sorted_within_series) in [ + ("sorted_runs", 512, 16, true), + ("shuffled_runs", 512, 16, false), + ("high_cardinality", 8192, 1, true), + ] { + let batch = primary_key_sort_batch(series_count, samples_per_series, sorted_within_series); + group.throughput(Throughput::Elements(batch.num_rows() as u64)); + group.bench_function(format!("{name}/arrow_lexsort"), |b| { + b.iter(|| lexsort_primary_key_record_batch(&batch)); + }); + group.bench_function(format!("{name}/rank_scatter"), |b| { + b.iter(|| sort_primary_key_record_batch(&batch).unwrap()); + }); + } +} + /// Writes rows. fn write_rows(c: &mut Criterion) { let metadata = Arc::new(memtable_util::metadata_with_primary_key(vec![1, 0], true)); @@ -391,6 +503,7 @@ fn bulk_part_record_batch_iter_filter(c: &mut Criterion) { criterion_group!( benches, + primary_key_sort, write_rows, full_scan, filter_1_host, diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index 20482794079..cc2c37a7853 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -30,11 +30,13 @@ use datafusion_common::pruning::PruningStatistics; use datafusion_expr::utils::expr_to_columns; use datatypes::arrow; use datatypes::arrow::array::{ - Array, ArrayRef, BinaryArray, BooleanArray, StringDictionaryBuilder, UInt8Array, UInt64Array, + Array, ArrayRef, BinaryArray, BooleanArray, DictionaryArray, StringDictionaryBuilder, + TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, + TimestampSecondArray, UInt8Array, UInt32Array, UInt64Array, }; use datatypes::arrow::compute::{SortColumn, SortOptions, concat_batches}; use datatypes::arrow::datatypes::{ - DataType as ArrowDataType, Field, Schema, SchemaRef, UInt32Type, + DataType as ArrowDataType, Field, Schema, SchemaRef, TimeUnit, UInt32Type, }; use datatypes::data_type::DataType; use datatypes::extension::json::is_json2_extension_type; @@ -771,6 +773,147 @@ fn new_primary_key_column_builders( /// Sorts the record batch with primary key format. pub fn sort_primary_key_record_batch(batch: &RecordBatch) -> Result { + let indices = if let Some(indices) = try_sort_primary_key_indices(batch)? { + indices + } else { + lexsort_primary_key_indices(batch)? + }; + + datatypes::arrow::compute::take_record_batch(batch, &indices).context(ComputeArrowSnafu) +} + +/// Sorts the known, non-null primary-key layout without comparing encoded keys for every row. +/// +/// Returns `None` if the batch does not have the exact layout supported by this fast path. +fn try_sort_primary_key_indices(batch: &RecordBatch) -> Result> { + let total_columns = batch.num_columns(); + let timestamp = batch.column(total_columns - 4); + let primary_key = batch.column(total_columns - 3); + let sequence = batch.column(total_columns - 2); + + let Some(primary_key) = primary_key + .as_any() + .downcast_ref::>() + else { + return Ok(None); + }; + let Some(sequence) = sequence.as_any().downcast_ref::() else { + return Ok(None); + }; + + if batch.num_rows() > u32::MAX as usize + || primary_key.null_count() != 0 + || primary_key.values().data_type() != &ArrowDataType::Binary + || primary_key.values().null_count() != 0 + || primary_key.values().len() > batch.num_rows() + || timestamp.null_count() != 0 + || sequence.null_count() != 0 + { + return Ok(None); + } + + let indices = match timestamp.data_type() { + ArrowDataType::Timestamp(TimeUnit::Second, _) => timestamp + .as_any() + .downcast_ref::() + .map(|timestamp| sort_primary_key_indices(primary_key, timestamp.values(), sequence)), + ArrowDataType::Timestamp(TimeUnit::Millisecond, _) => timestamp + .as_any() + .downcast_ref::() + .map(|timestamp| sort_primary_key_indices(primary_key, timestamp.values(), sequence)), + ArrowDataType::Timestamp(TimeUnit::Microsecond, _) => timestamp + .as_any() + .downcast_ref::() + .map(|timestamp| sort_primary_key_indices(primary_key, timestamp.values(), sequence)), + ArrowDataType::Timestamp(TimeUnit::Nanosecond, _) => timestamp + .as_any() + .downcast_ref::() + .map(|timestamp| sort_primary_key_indices(primary_key, timestamp.values(), sequence)), + _ => None, + }; + + indices.transpose() +} + +/// Computes sorted row indices by ranking dictionary values, scattering rows into primary-key +/// buckets, and sorting only buckets that are not already ordered by time and sequence. +fn sort_primary_key_indices( + primary_key: &PrimaryKeyArray, + timestamps: &[i64], + sequences: &UInt64Array, +) -> Result { + debug_assert_eq!(primary_key.len(), timestamps.len()); + debug_assert_eq!(primary_key.len(), sequences.len()); + debug_assert_eq!(primary_key.null_count(), 0); + debug_assert_eq!(primary_key.values().null_count(), 0); + debug_assert_eq!(sequences.null_count(), 0); + + // Ranks compare in the same order as the dictionary values. Equal dictionary values receive + // the same rank, which is required because Arrow dictionary concatenation only deduplicates + // values on a best-effort basis. + let value_ranks = datatypes::arrow::compute::rank( + primary_key.values().as_ref(), + Some(SortOptions { + descending: false, + nulls_first: true, + }), + ) + .context(ComputeArrowSnafu)?; + + // A rank is at most the number of dictionary values. It is not necessarily dense when the + // dictionary contains duplicate values, so keep one extra slot and allow empty buckets. + let mut counts = vec![0usize; value_ranks.len() + 1]; + for &key in primary_key.keys().values() { + counts[value_ranks[key as usize] as usize] += 1; + } + + let mut bucket_offsets = vec![0usize; counts.len()]; + let mut offset = 0; + for (bucket_offset, &count) in bucket_offsets.iter_mut().zip(&counts) { + *bucket_offset = offset; + offset += count; + } + + // Scatter in input order. This preserves already sorted runs within each primary key. + let mut next_offsets = bucket_offsets.clone(); + let mut indices = vec![0u32; primary_key.len()]; + for (row, &key) in primary_key.keys().values().iter().enumerate() { + let rank = value_ranks[key as usize] as usize; + indices[next_offsets[rank]] = row as u32; + next_offsets[rank] += 1; + } + + let sequences = sequences.values(); + for (&start, count) in bucket_offsets.iter().zip(counts) { + let end = start + count; + let bucket = &mut indices[start..end]; + if !bucket.is_sorted_by(|left, right| { + compare_time_sequence(timestamps, sequences, *left, *right).is_le() + }) { + bucket.sort_unstable_by(|left, right| { + compare_time_sequence(timestamps, sequences, *left, *right) + }); + } + } + + Ok(UInt32Array::from(indices)) +} + +#[inline] +fn compare_time_sequence( + timestamps: &[i64], + sequences: &[u64], + left: u32, + right: u32, +) -> std::cmp::Ordering { + let left = left as usize; + let right = right as usize; + timestamps[left] + .cmp(×tamps[right]) + .then_with(|| sequences[right].cmp(&sequences[left])) +} + +fn lexsort_primary_key_indices(batch: &RecordBatch) -> Result { let total_columns = batch.num_columns(); let sort_columns = vec![ // Primary key column (ascending) @@ -799,10 +942,7 @@ pub fn sort_primary_key_record_batch(batch: &RecordBatch) -> Result }, ]; - let indices = datatypes::arrow::compute::lexsort_to_indices(&sort_columns, None) - .context(ComputeArrowSnafu)?; - - datatypes::arrow::compute::take_record_batch(batch, &indices).context(ComputeArrowSnafu) + datatypes::arrow::compute::lexsort_to_indices(&sort_columns, None).context(ComputeArrowSnafu) } /// Converts a `BulkPart` that is unordered and without encoded primary keys into a `BulkPart` @@ -1674,6 +1814,8 @@ mod tests { use datatypes::prelude::{ConcreteDataType, Value}; use datatypes::schema::ColumnSchema; use mito_codec::row_converter::build_primary_key_codec; + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use store_api::storage::RegionId; use store_api::storage::consts::ReservedColumnId; @@ -1692,6 +1834,137 @@ mod tests { sequence: u64, } + #[test] + fn test_sort_primary_key_record_batch_with_duplicate_dictionary_values() { + // Dictionary keys 1 and 2 both represent "series_a". This can happen after Arrow + // concatenates dictionaries because dictionary deduplication is best-effort. + let primary_key = DictionaryArray::try_new( + UInt32Array::from(vec![0, 1, 0, 2, 2]), + Arc::new(BinaryArray::from_vec(vec![ + b"series_b".as_slice(), + b"series_a".as_slice(), + b"series_a".as_slice(), + ])), + ) + .unwrap(); + let timestamps = TimestampMillisecondArray::from(vec![20, 30, 10, 30, 10]); + let sequences = UInt64Array::from(vec![5, 6, 4, 8, 9]); + let row_ids = UInt32Array::from_iter_values(0..5); + let op_types = UInt8Array::from_value(OpType::Put as u8, 5); + let schema = Arc::new(Schema::new(vec![ + Field::new("row_id", ArrowDataType::UInt32, false), + Field::new( + "ts", + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new( + PRIMARY_KEY_COLUMN_NAME, + ArrowDataType::Dictionary( + Box::new(ArrowDataType::UInt32), + Box::new(ArrowDataType::Binary), + ), + false, + ), + Field::new("__sequence", ArrowDataType::UInt64, false), + Field::new("__op_type", ArrowDataType::UInt8, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(row_ids), + Arc::new(timestamps), + Arc::new(primary_key), + Arc::new(sequences), + Arc::new(op_types), + ], + ) + .unwrap(); + + let expected_indices = lexsort_primary_key_indices(&batch).unwrap(); + let expected = + datatypes::arrow::compute::take_record_batch(&batch, &expected_indices).unwrap(); + let actual = sort_primary_key_record_batch(&batch).unwrap(); + + let expected_row_ids = expected + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let actual_row_ids = actual + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(expected_row_ids, actual_row_ids); + assert_eq!(actual_row_ids.values(), &[4, 3, 1, 2, 0]); + } + + #[test] + fn test_sort_primary_key_record_batch_against_lexsort() { + let mut rng = StdRng::seed_from_u64(0x5eed); + for _ in 0..100 { + let num_rows = rng.random_range(0..128); + let keys = UInt32Array::from_iter_values((0..num_rows).map(|_| rng.random_range(0..8))); + let primary_key = DictionaryArray::try_new( + keys, + Arc::new(BinaryArray::from_vec(vec![ + b"d".as_slice(), + b"a".as_slice(), + b"c".as_slice(), + b"b".as_slice(), + b"a".as_slice(), + b"d".as_slice(), + b"b".as_slice(), + b"c".as_slice(), + ])), + ) + .unwrap(); + let timestamps = TimestampMillisecondArray::from_iter_values( + (0..num_rows).map(|_| rng.random_range(-1000..1000)), + ); + // Unique sequences ensure that the oracle and fast path have no fully equal sort keys. + let sequences = UInt64Array::from_iter_values(0..num_rows as u64); + let row_ids = UInt32Array::from_iter_values(0..num_rows as u32); + let op_types = UInt8Array::from_value(OpType::Put as u8, num_rows); + let schema = Arc::new(Schema::new(vec![ + Field::new("row_id", ArrowDataType::UInt32, false), + Field::new( + "ts", + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new( + PRIMARY_KEY_COLUMN_NAME, + ArrowDataType::Dictionary( + Box::new(ArrowDataType::UInt32), + Box::new(ArrowDataType::Binary), + ), + false, + ), + Field::new("__sequence", ArrowDataType::UInt64, false), + Field::new("__op_type", ArrowDataType::UInt8, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(row_ids), + Arc::new(timestamps), + Arc::new(primary_key), + Arc::new(sequences), + Arc::new(op_types), + ], + ) + .unwrap(); + + let expected_indices = lexsort_primary_key_indices(&batch).unwrap(); + let expected = + datatypes::arrow::compute::take_record_batch(&batch, &expected_indices).unwrap(); + let actual = sort_primary_key_record_batch(&batch).unwrap(); + assert_eq!(expected.column(0), actual.column(0)); + } + } + #[test] fn test_unordered_part_tracks_estimated_bytes() { let mut part = UnorderedPart::new(); From 764c93bf43bd82da6eae8a5fa671fc56c39cc156 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Thu, 13 Aug 2026 10:51:53 +0000 Subject: [PATCH 40/92] perf(query): choose bounded CTE as hash join build side (#8807) * fix(query): choose bounded CTE as hash join build side Signed-off-by: jeremyhi * perf(query): remove join estimate cap * test(compat): accept repartition in analyze plan Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- Cargo.lock | 68 +++--- Cargo.toml | 30 +-- src/query/src/dist_plan/merge_scan.rs | 219 +++++++++++++++++- src/query/src/optimizer/pass_distribution.rs | 4 +- src/standalone/src/options.rs | 10 + .../distributed/explain/step_aggr.result | 8 + .../explain/step_aggr_advance.result | 111 ++++++--- .../explain/step_aggr_basic.result | 10 + .../distributed/explain/subqueries.result | 13 +- .../distributed/flow-tql/tsid_on_phy.result | 4 +- .../cases/distributed/optimizer/count.result | 5 + .../optimizer/first_value_advance.result | 8 + .../optimizer/last_value_advance.result | 8 + .../optimizer/range_select_projection.result | 2 +- .../common/aggregate/distinct.result | 10 +- .../common/aggregate/multi_regions.result | 2 + .../create/metric_engine_partition.result | 11 +- .../common/cte/cte_join_build_side.result | 116 ++++++++++ .../common/cte/cte_join_build_side.sql | 46 ++++ .../filter/hash_join_topk_dyn_filter.result | 12 +- .../common/order/order_by_exceptions.result | 7 +- .../promql/histogram_multi_partition.result | 1 + .../promql/tsid_binary_join_regression.result | 20 +- .../tsid_histogram_quantile_regression.result | 2 +- .../cases/standalone/common/range/nest.result | 4 +- .../standalone/common/tql/partition.result | 5 +- .../standalone/common/tql/tql-cte.result | 24 +- tests/cases/standalone/optimizer/count.result | 5 + .../optimizer/first_value_advance.result | 8 + .../optimizer/last_value_advance.result | 8 + .../optimizer/rewrite_set_comparison.result | 22 +- .../tql-explain-analyze/analyze.result | 1 + .../tql-explain-analyze/tsid_column.result | 8 +- .../verify.result | 1 + tests/conf/datanode-test.toml.template | 3 + tests/conf/frontend-test.toml.template | 3 + tests/conf/standalone-test.toml.template | 3 + 37 files changed, 674 insertions(+), 148 deletions(-) create mode 100644 tests/cases/standalone/common/cte/cte_join_build_side.result create mode 100644 tests/cases/standalone/common/cte/cte_join_build_side.sql diff --git a/Cargo.lock b/Cargo.lock index 97b2fc76a79..a7a7ff9b877 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3733,7 +3733,7 @@ checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" [[package]] name = "datafusion" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "arrow-schema 58.3.0", @@ -3787,7 +3787,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3811,7 +3811,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3833,7 +3833,7 @@ dependencies = [ [[package]] name = "datafusion-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -3857,7 +3857,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "futures", "log", @@ -3867,7 +3867,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "async-compression", @@ -3901,7 +3901,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "arrow-ipc 58.3.0", @@ -3924,7 +3924,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3946,7 +3946,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3969,7 +3969,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3998,12 +3998,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" [[package]] name = "datafusion-execution" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "arrow-buffer 58.3.0", @@ -4025,7 +4025,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "async-trait", @@ -4047,7 +4047,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4059,7 +4059,7 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "arrow-buffer 58.3.0", @@ -4090,7 +4090,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4111,7 +4111,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4123,7 +4123,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "arrow-ord 58.3.0", @@ -4147,7 +4147,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "async-trait", @@ -4162,7 +4162,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4179,7 +4179,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -4188,7 +4188,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "datafusion-doc", "quote", @@ -4198,7 +4198,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "chrono", @@ -4247,7 +4247,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4270,7 +4270,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4284,7 +4284,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4300,7 +4300,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4318,7 +4318,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4349,7 +4349,7 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "chrono", @@ -4376,7 +4376,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4386,7 +4386,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4402,7 +4402,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "async-trait", "datafusion-common", @@ -4415,7 +4415,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "arrow 58.3.0", "bigdecimal 0.4.8", @@ -4433,7 +4433,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=452cb4b786ccadab24a2afa5a49034e85411135f#452cb4b786ccadab24a2afa5a49034e85411135f" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" dependencies = [ "async-recursion", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 63dbc76f383..bfd6153bf5b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -346,21 +346,21 @@ git = "https://github.com/GreptimeTeam/greptime-meter.git" rev = "5618e779cf2bb4755b499c630fba4c35e91898cb" [patch.crates-io] -datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } -datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "452cb4b786ccadab24a2afa5a49034e85411135f" } +datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "2aefa08a8d69c96eec2d6d6703598a009bba6e4c" } # on branch v0.61.x [profile.release] diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index aafc17c469f..8d3580ea452 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -42,8 +42,9 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, SendableRecordBatchStream, }; -use datafusion_common::{Column as ColumnExpr, DataFusionError, Result}; -use datafusion_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion_common::stats::Precision; +use datafusion_common::{Column as ColumnExpr, DataFusionError, Result, Statistics}; +use datafusion_expr::{Expr, Extension, FetchType, LogicalPlan, UserDefinedLogicalNodeCore}; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr::{Distribution, EquivalenceProperties, PhysicalSortExpr}; use datatypes::extension::json::{ @@ -84,6 +85,51 @@ fn query_engine_state_from_task_context(context: &TaskContext) -> Option Option { + match plan { + LogicalPlan::Limit(limit) => { + let input_bound = remote_plan_row_bound(&limit.input); + match limit.get_fetch_type() { + Ok(FetchType::Literal(Some(fetch))) => { + Some(input_bound.map_or(fetch, |bound| bound.min(fetch))) + } + _ => input_bound, + } + } + LogicalPlan::Sort(sort) => { + let input_bound = remote_plan_row_bound(&sort.input); + sort.fetch + .map(|fetch| input_bound.map_or(fetch, |bound| bound.min(fetch))) + .or(input_bound) + } + LogicalPlan::Projection(projection) => remote_plan_row_bound(&projection.input), + LogicalPlan::Filter(filter) => remote_plan_row_bound(&filter.input), + LogicalPlan::SubqueryAlias(alias) => remote_plan_row_bound(&alias.input), + LogicalPlan::Window(window) => remote_plan_row_bound(&window.input), + LogicalPlan::Repartition(repartition) => remote_plan_row_bound(&repartition.input), + LogicalPlan::Distinct(distinct) => remote_plan_row_bound(distinct.input()), + LogicalPlan::Aggregate(aggregate) => { + if aggregate + .group_expr + .iter() + .any(|expr| matches!(expr, Expr::GroupingSet(_))) + { + None + } else if aggregate.group_expr.is_empty() { + Some(1) + } else { + remote_plan_row_bound(&aggregate.input) + } + } + _ => None, + } +} + fn remote_dyn_filter_enabled(query_ctx: &QueryContextRef) -> Result { remote_dyn_filter_pushdown_enabled_from_extensions(&query_ctx.extensions()) .map_err(|err| DataFusionError::External(Box::new(err))) @@ -429,6 +475,7 @@ impl MergeScanExec { enable_per_region_metrics: bool, ) -> Result { let arrow_schema = maybe_amend_json2_field(arrow_schema); + let output_partition_count = Self::output_partition_count(regions.len(), target_partition); // States the output ordering of the plan. // @@ -438,7 +485,7 @@ impl MergeScanExec { // // Otherwise, we need to use the default ordering. let eq_properties = if let LogicalPlan::Sort(sort) = &plan - && target_partition >= regions.len() + && output_partition_count >= regions.len() { let lex_ordering = sort .expr @@ -477,7 +524,7 @@ impl MergeScanExec { } }) .collect(); - let partitioning = Partitioning::Hash(partition_exprs, target_partition); + let partitioning = Partitioning::Hash(partition_exprs, output_partition_count); let properties = Arc::new(PlanProperties::new( eq_properties, @@ -504,6 +551,28 @@ impl MergeScanExec { }) } + /// Conservative row-count upper bound for all selected regions. + /// + /// This is not an expected cardinality: DataFusion receives it as an + /// inexact estimate only because `Statistics` has no upper-bound precision. + fn estimated_num_rows(&self) -> Precision { + if self.regions.is_empty() { + return Precision::Inexact(0); + } + + let Some(rows_per_region) = remote_plan_row_bound(&self.plan) else { + return Precision::Absent; + }; + rows_per_region + .checked_mul(self.regions.len()) + .map_or(Precision::Absent, Precision::Inexact) + } + + /// Number of partitions populated by the region striping in [`Self::to_stream`]. + fn output_partition_count(num_regions: usize, target_partition: usize) -> usize { + num_regions.max(1).min(target_partition.max(1)) + } + pub fn to_stream( &self, context: Arc, @@ -518,7 +587,8 @@ impl MergeScanExec { let sub_stage_metrics_moved = self.sub_stage_metrics.clone(); let partition_metrics_moved = self.partition_metrics.clone(); let plan = self.plan.clone(); - let target_partition = self.target_partition; + let target_partition = + Self::output_partition_count(self.regions.len(), self.target_partition); let remote_dyn_filter_enabled = remote_dyn_filter_enabled(&self.query_ctx)?; let captured_remote_dyn_filters = if remote_dyn_filter_enabled { self.captured_remote_dyn_filters() @@ -855,7 +925,7 @@ impl MergeScanExec { metric: self.metric.clone(), properties: Arc::new(PlanProperties::new( self.properties.eq_properties.clone(), - Partitioning::Hash(overlaps, self.target_partition), + Partitioning::Hash(overlaps, self.partition_count()), self.properties.emission_type, self.properties.boundedness, )), @@ -906,7 +976,7 @@ impl MergeScanExec { } pub fn partition_count(&self) -> usize { - self.target_partition + Self::output_partition_count(self.regions.len(), self.target_partition) } pub fn region_count(&self) -> usize { @@ -1174,6 +1244,16 @@ impl ExecutionPlan for MergeScanExec { Some(self.metric.clone_inner()) } + fn partition_statistics(&self, partition: Option) -> Result { + if partition.is_some() { + return Ok(Statistics::new_unknown(&self.arrow_schema)); + } + + let mut statistics = Statistics::new_unknown(&self.arrow_schema); + statistics.num_rows = self.estimated_num_rows(); + Ok(statistics) + } + fn name(&self) -> &str { "MergeScanExec" } @@ -1404,7 +1484,6 @@ mod tests { region_count: u64, target_partition: usize, ) -> MergeScanExec { - let session_state = SessionStateBuilder::new().build(); let plan = LogicalPlanBuilder::empty(true) .project(vec![lit(1i64).alias("ts")]) .unwrap() @@ -1412,11 +1491,21 @@ mod tests { .unwrap() .build() .unwrap(); - let schema = plan.schema().as_arrow().clone(); let regions = (0..region_count) .map(|region_number| RegionId::new(1024, region_number as u32)) .collect(); + merge_scan_exec_with_plan(regions, plan, target_partition) + } + + fn merge_scan_exec_with_plan( + regions: Vec, + plan: LogicalPlan, + target_partition: usize, + ) -> MergeScanExec { + let session_state = SessionStateBuilder::new().build(); + let schema = plan.schema().as_arrow().clone(); + MergeScanExec::new( &session_state, // The table name is not relevant to these ordering metadata tests; @@ -1555,6 +1644,118 @@ mod tests { assert!(exec.properties().output_ordering().is_some()); } + #[test] + fn merge_scan_reports_populated_partition_count() { + let cases = [(0, 10, 1), (1, 10, 1), (3, 2, 2), (5, 10, 5), (3, 0, 1)]; + + for (region_count, target, expected) in cases { + let exec = merge_scan_exec_with_sorted_input(region_count, target); + assert_eq!(exec.partition_count(), expected); + assert_eq!( + exec.properties().output_partitioning().partition_count(), + expected + ); + } + } + + #[test] + fn merge_scan_reports_only_deterministic_plan_bounds() { + use datafusion::functions_aggregate::expr_fn::count; + use datafusion_expr::GroupingSet; + + let regions = vec![RegionId::new(1024, 1), RegionId::new(1024, 2)]; + let limited = LogicalPlanBuilder::empty(true) + .project(vec![lit(1i64).alias("col")]) + .unwrap() + .limit(0, Some(50)) + .unwrap() + .build() + .unwrap(); + assert_eq!( + merge_scan_exec_with_plan(regions.clone(), limited, 10) + .partition_statistics(None) + .unwrap() + .num_rows, + Precision::Inexact(100) + ); + + let large_bound = i32::MAX as usize + 1; + let large_limit = LogicalPlanBuilder::empty(true) + .project(vec![lit(1i64).alias("col")]) + .unwrap() + .limit(0, Some(large_bound)) + .unwrap() + .build() + .unwrap(); + assert_eq!( + merge_scan_exec_with_plan(vec![RegionId::new(1024, 1)], large_limit, 10) + .partition_statistics(None) + .unwrap() + .num_rows, + Precision::Inexact(large_bound) + ); + + let uncapped = LogicalPlanBuilder::empty(true) + .project(vec![lit(1i64).alias("col")]) + .unwrap() + .build() + .unwrap(); + assert_eq!( + merge_scan_exec_with_plan(regions.clone(), uncapped.clone(), 10) + .partition_statistics(None) + .unwrap() + .num_rows, + Precision::Absent + ); + assert_eq!( + merge_scan_exec_with_plan(Vec::new(), uncapped, 10) + .partition_statistics(None) + .unwrap() + .num_rows, + Precision::Inexact(0) + ); + + let global_aggregate = LogicalPlanBuilder::empty(true) + .project(vec![lit(1i64).alias("col")]) + .unwrap() + .limit(0, Some(0)) + .unwrap() + .aggregate(Vec::::new(), vec![count(lit(1))]) + .unwrap() + .build() + .unwrap(); + assert_eq!( + merge_scan_exec_with_plan(regions.clone(), global_aggregate, 10) + .partition_statistics(None) + .unwrap() + .num_rows, + Precision::Inexact(2) + ); + + let grouping_sets = LogicalPlanBuilder::empty(true) + .project(vec![lit(1i64).alias("col")]) + .unwrap() + .limit(0, Some(50)) + .unwrap() + .aggregate( + vec![Expr::GroupingSet(GroupingSet::GroupingSets(vec![ + vec![], + vec![col("col")], + ]))], + Vec::::new(), + ) + .unwrap() + .build() + .unwrap(); + assert_eq!( + merge_scan_exec_with_plan(regions, grouping_sets, 10) + .partition_statistics(None) + .unwrap() + .num_rows, + Precision::Absent + ); + } + #[test] fn sub_stage_metrics_are_sorted_by_region_id() { let exec = merge_scan_exec_with_sorted_input(0, 1); diff --git a/src/query/src/optimizer/pass_distribution.rs b/src/query/src/optimizer/pass_distribution.rs index 37274286482..c45d4695d93 100644 --- a/src/query/src/optimizer/pass_distribution.rs +++ b/src/query/src/optimizer/pass_distribution.rs @@ -277,8 +277,8 @@ mod tests { panic!("expected right merge scan hash partitioning"); }; - assert_eq!(*left_count, 32); - assert_eq!(*right_count, 32); + assert_eq!(*left_count, 2); + assert_eq!(*right_count, 2); assert_eq!( column_names(left_exprs), vec![DATA_SCHEMA_TSID_COLUMN_NAME, "greptime_timestamp"] diff --git a/src/standalone/src/options.rs b/src/standalone/src/options.rs index 637715013af..bedf9eebc70 100644 --- a/src/standalone/src/options.rs +++ b/src/standalone/src/options.rs @@ -157,6 +157,7 @@ impl StandaloneOptions { meta_client: None, logging: cloned_opts.logging, user_provider: cloned_opts.user_provider, + query: cloned_opts.query, slow_query: cloned_opts.slow_query, event_recorder: cloned_opts.event_recorder, heartbeat_env_vars: cloned_opts.heartbeat_env_vars.clone(), @@ -227,4 +228,13 @@ mod tests { &frontend_options.event_recorder.event_types, )); } + + #[test] + fn test_query_options_propagated_to_components() { + let mut options = StandaloneOptions::default(); + options.query.parallelism = 4; + + assert_eq!(options.frontend_options().query.parallelism, 4); + assert_eq!(options.datanode_options().query.parallelism, 4); + } } diff --git a/tests/cases/distributed/explain/step_aggr.result b/tests/cases/distributed/explain/step_aggr.result index 78218d514da..89634a3cc6d 100644 --- a/tests/cases/distributed/explain/step_aggr.result +++ b/tests/cases/distributed/explain/step_aggr.result @@ -67,6 +67,7 @@ FROM |_|_AggregateExec: mode=Final, gby=[], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -96,6 +97,7 @@ FROM |_|_|_AggregateExec: mode=Final, gby=[], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED @@ -148,6 +150,7 @@ FROM | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[avg(integers.i)]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[avg(integers.i)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -172,6 +175,7 @@ FROM | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[avg(integers.i)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[avg(integers.i)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__avg_state(integers.i)] REDACTED @@ -249,6 +253,7 @@ ORDER BY |_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)]_| |_|_RepartitionExec: partitioning=REDACTED |_|_AggregateExec: mode=Partial, gby=[ts@0 as ts], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -285,6 +290,7 @@ ORDER BY |_|_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[ts@0 as ts], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED @@ -362,6 +368,7 @@ ORDER BY |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("2 seconds"),integers.ts)@0 as date_bin(Utf8("2 seconds"),integers.ts)], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)]_| |_|_RepartitionExec: partitioning=REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("2 seconds"),integers.ts)@0 as date_bin(Utf8("2 seconds"),integers.ts)], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -398,6 +405,7 @@ ORDER BY |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("2 seconds"),integers.ts)@0 as date_bin(Utf8("2 seconds"),integers.ts)], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("2 seconds"),integers.ts)@0 as date_bin(Utf8("2 seconds"),integers.ts)], aggr=[count(integers.i), sum(integers.i), uddsketch_state(Int64(128),Float64(0.01),integers.i), hll(integers.i)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("2 seconds"),integers.ts)@0 as date_bin(Utf8("2 seconds"),integers.ts)], aggr=[__count_state(integers.i), __sum_state(integers.i), __uddsketch_state_state(Int64(128),Float64(0.01),integers.i), __hll_state(integers.i)] REDACTED diff --git a/tests/cases/distributed/explain/step_aggr_advance.result b/tests/cases/distributed/explain/step_aggr_advance.result index a78a177bc3d..9fc8dfe5a1c 100644 --- a/tests/cases/distributed/explain/step_aggr_advance.result +++ b/tests/cases/distributed/explain/step_aggr_advance.result @@ -108,8 +108,11 @@ tql explain (1752591864, 1752592164, '30s') sum by (a, b) (max_over_time(aggr_op | | ]] | | physical_plan | SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, greptime_timestamp@2 ASC NULLS LAST] | | | SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, greptime_timestamp@2 ASC NULLS LAST], preserve_partitioning=[true] | -| | AggregateExec: mode=SinglePartitioned, gby=[a@0 as a, b@1 as b, greptime_timestamp@2 as greptime_timestamp], aggr=[sum(prom_max_over_time(greptime_timestamp_range,greptime_value))] | -| | MergeScanExec: REDACTED +| | AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, greptime_timestamp@2 as greptime_timestamp], aggr=[sum(prom_max_over_time(greptime_timestamp_range,greptime_value))] | +| | RepartitionExec: partitioning=REDACTED +| | AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b, greptime_timestamp@2 as greptime_timestamp], aggr=[sum(prom_max_over_time(greptime_timestamp_range,greptime_value))] | +| | RepartitionExec: partitioning=REDACTED +| | MergeScanExec: REDACTED | | | +---------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -127,7 +130,10 @@ tql analyze (1752591864, 1752592164, '30s') sum by (a, b) (max_over_time(aggr_op +-+-+-+ | 0_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, greptime_timestamp@2 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, greptime_timestamp@2 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_AggregateExec: mode=SinglePartitioned, gby=[a@0 as a, b@1 as b, greptime_timestamp@2 as greptime_timestamp], aggr=[sum(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, greptime_timestamp@2 as greptime_timestamp], aggr=[sum(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b, greptime_timestamp@2 as greptime_timestamp], aggr=[sum(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, greptime_timestamp@2 as greptime_timestamp], aggr=[__sum_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED @@ -177,8 +183,11 @@ tql explain (1752591864, 1752592164, '30s') avg by (a) (max_over_time(aggr_optim | | ]] | | physical_plan | SortPreservingMergeExec: [a@0 ASC NULLS LAST, greptime_timestamp@1 ASC NULLS LAST] | | | SortExec: expr=[a@0 ASC NULLS LAST, greptime_timestamp@1 ASC NULLS LAST], preserve_partitioning=[true] | -| | AggregateExec: mode=SinglePartitioned, gby=[a@0 as a, greptime_timestamp@1 as greptime_timestamp], aggr=[avg(prom_max_over_time(greptime_timestamp_range,greptime_value))] | -| | MergeScanExec: REDACTED +| | AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, greptime_timestamp@1 as greptime_timestamp], aggr=[avg(prom_max_over_time(greptime_timestamp_range,greptime_value))] | +| | RepartitionExec: partitioning=REDACTED +| | AggregateExec: mode=Partial, gby=[a@0 as a, greptime_timestamp@1 as greptime_timestamp], aggr=[avg(prom_max_over_time(greptime_timestamp_range,greptime_value))] | +| | RepartitionExec: partitioning=REDACTED +| | MergeScanExec: REDACTED | | | +---------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -196,7 +205,10 @@ tql analyze (1752591864, 1752592164, '30s') avg by (a) (max_over_time(aggr_optim +-+-+-+ | 0_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, greptime_timestamp@1 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[a@0 ASC NULLS LAST, greptime_timestamp@1 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_AggregateExec: mode=SinglePartitioned, gby=[a@0 as a, greptime_timestamp@1 as greptime_timestamp], aggr=[avg(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, greptime_timestamp@1 as greptime_timestamp], aggr=[avg(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, greptime_timestamp@1 as greptime_timestamp], aggr=[avg(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, greptime_timestamp@1 as greptime_timestamp], aggr=[__avg_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED @@ -311,8 +323,11 @@ tql explain (1752591864, 1752592164, '30s') min by (b, c, d) (max_over_time(aggr | | ]] | | physical_plan | SortPreservingMergeExec: [b@0 ASC NULLS LAST, c@1 ASC NULLS LAST, d@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] | | | SortExec: expr=[b@0 ASC NULLS LAST, c@1 ASC NULLS LAST, d@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] | -| | AggregateExec: mode=SinglePartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[min(prom_max_over_time(greptime_timestamp_range,greptime_value))] | -| | MergeScanExec: REDACTED +| | AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[min(prom_max_over_time(greptime_timestamp_range,greptime_value))] | +| | RepartitionExec: partitioning=REDACTED +| | AggregateExec: mode=Partial, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[min(prom_max_over_time(greptime_timestamp_range,greptime_value))] | +| | RepartitionExec: partitioning=REDACTED +| | MergeScanExec: REDACTED | | | +---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -330,7 +345,10 @@ tql analyze (1752591864, 1752592164, '30s') min by (b, c, d) (max_over_time(aggr +-+-+-+ | 0_| 0_|_SortPreservingMergeExec: [b@0 ASC NULLS LAST, c@1 ASC NULLS LAST, d@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[b@0 ASC NULLS LAST, c@1 ASC NULLS LAST, d@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_AggregateExec: mode=SinglePartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[min(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[min(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[min(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[__min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED @@ -380,7 +398,8 @@ tql explain sum(aggr_optimize_not); | | AggregateExec: mode=FinalPartitioned, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[sum(aggr_optimize_not.greptime_value)] | | | RepartitionExec: partitioning=REDACTED | | AggregateExec: mode=Partial, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[sum(aggr_optimize_not.greptime_value)] | -| | MergeScanExec: REDACTED +| | RepartitionExec: partitioning=REDACTED +| | MergeScanExec: REDACTED | | | +---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -401,6 +420,7 @@ tql analyze sum(aggr_optimize_not); |_|_|_AggregateExec: mode=FinalPartitioned, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[sum(aggr_optimize_not.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[sum(aggr_optimize_not.greptime_value)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[__sum_state(aggr_optimize_not.greptime_value)] REDACTED @@ -477,14 +497,17 @@ tql explain (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize | physical_plan | ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp, sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))@5 / sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))@4 as aggr_optimize_not.sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000))) / aggr_optimize_not_count.sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] | | | REDACTED | | CoalescePartitionsExec | -| | AggregateExec: mode=SinglePartitioned, gby=[a@2 as a, b@3 as b, c@4 as c, greptime_timestamp@0 as greptime_timestamp], aggr=[sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] | -| | FilterExec: prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000))@1 IS NOT NULL | -| | ProjectionExec: expr=[greptime_timestamp@4 as greptime_timestamp, prom_rate(greptime_timestamp_range@6, greptime_value@5, greptime_timestamp@4, 120000) as prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)), a@0 as a, b@1 as b, c@2 as c] | -| | PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] | -| | PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] | -| | PromSeriesDivideExec: tags=["a", "b", "c", "d"] | -| | SortExec: expr=[a@0 ASC, b@1 ASC, c@2 ASC, d@3 ASC, greptime_timestamp@4 ASC], preserve_partitioning=[true] | -| | MergeScanExec: REDACTED +| | AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp], aggr=[sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] | +| | RepartitionExec: partitioning=REDACTED +| | AggregateExec: mode=Partial, gby=[a@2 as a, b@3 as b, c@4 as c, greptime_timestamp@0 as greptime_timestamp], aggr=[sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] | +| | FilterExec: prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000))@1 IS NOT NULL | +| | ProjectionExec: expr=[greptime_timestamp@4 as greptime_timestamp, prom_rate(greptime_timestamp_range@6, greptime_value@5, greptime_timestamp@4, 120000) as prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)), a@0 as a, b@1 as b, c@2 as c] | +| | PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] | +| | PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] | +| | PromSeriesDivideExec: tags=["a", "b", "c", "d"] | +| | SortExec: expr=[a@0 ASC, b@1 ASC, c@2 ASC, d@3 ASC, greptime_timestamp@4 ASC], preserve_partitioning=[true] | +| | RepartitionExec: partitioning=REDACTED +| | MergeScanExec: REDACTED | | RepartitionExec: partitioning=REDACTED | | MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] | | | SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] | @@ -507,13 +530,16 @@ tql analyze (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize | 0_| 0_|_ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp, sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))@5 / sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))@4 as aggr_optimize_not.sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000))) / aggr_optimize_not_count.sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] REDACTED |_|_|_REDACTED |_|_|_CoalescePartitionsExec REDACTED -|_|_|_AggregateExec: mode=SinglePartitioned, gby=[a@2 as a, b@3 as b, c@4 as c, greptime_timestamp@0 as greptime_timestamp], aggr=[sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp], aggr=[sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@2 as a, b@3 as b, c@4 as c, greptime_timestamp@0 as greptime_timestamp], aggr=[sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] REDACTED |_|_|_FilterExec: prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000))@1 IS NOT NULL REDACTED |_|_|_ProjectionExec: expr=[greptime_timestamp@4 as greptime_timestamp, prom_rate(greptime_timestamp_range@6, greptime_value@5, greptime_timestamp@4, 120000) as prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)), a@0 as a, b@1 as b, c@2 as c] REDACTED |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED |_|_|_SortExec: expr=[a@0 ASC, b@1 ASC, c@2 ASC, d@3 ASC, greptime_timestamp@4 ASC], preserve_partitioning=[true] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeSortExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED @@ -710,9 +736,11 @@ GROUP BY | | TableScan: aggr_optimize_not | | | ]] | | physical_plan | ProjectionExec: expr=[min(aggr_optimize_not.greptime_value)@2 as min(aggr_optimize_not.greptime_value)] | -| | AggregateExec: mode=SinglePartitioned, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value)] | -| | CooperativeExec | -| | MergeScanExec: REDACTED +| | AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value)] | +| | RepartitionExec: partitioning=REDACTED +| | AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value)] | +| | RepartitionExec: partitioning=REDACTED +| | MergeScanExec: REDACTED | | | +---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -736,8 +764,10 @@ GROUP BY | stage | node | plan_| +-+-+-+ | 0_| 0_|_ProjectionExec: expr=[min(aggr_optimize_not.greptime_value)@2 as min(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_AggregateExec: mode=SinglePartitioned, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_CooperativeExec REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[__min_state(aggr_optimize_not.greptime_value)] REDACTED @@ -776,9 +806,11 @@ GROUP BY | | TableScan: aggr_optimize_not | | | ]] | | physical_plan | ProjectionExec: expr=[min(aggr_optimize_not.greptime_value)@2 + max(aggr_optimize_not.greptime_value)@3 as min(aggr_optimize_not.greptime_value) + max(aggr_optimize_not.greptime_value)] | -| | AggregateExec: mode=SinglePartitioned, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value), max(aggr_optimize_not.greptime_value)] | -| | CooperativeExec | -| | MergeScanExec: REDACTED +| | AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value), max(aggr_optimize_not.greptime_value)] | +| | RepartitionExec: partitioning=REDACTED +| | AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value), max(aggr_optimize_not.greptime_value)] | +| | RepartitionExec: partitioning=REDACTED +| | MergeScanExec: REDACTED | | | +---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -802,8 +834,10 @@ GROUP BY | stage | node | plan_| +-+-+-+ | 0_| 0_|_ProjectionExec: expr=[min(aggr_optimize_not.greptime_value)@2 + max(aggr_optimize_not.greptime_value)@3 as min(aggr_optimize_not.greptime_value) + max(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_AggregateExec: mode=SinglePartitioned, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value), max(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_CooperativeExec REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value), max(aggr_optimize_not.greptime_value)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[min(aggr_optimize_not.greptime_value), max(aggr_optimize_not.greptime_value)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[__min_state(aggr_optimize_not.greptime_value), __max_state(aggr_optimize_not.greptime_value)] REDACTED @@ -851,9 +885,11 @@ GROUP BY | | Projection: aggr_optimize_not.a, aggr_optimize_not.b, aggr_optimize_not.greptime_value | | | TableScan: aggr_optimize_not | | | ]] | -| physical_plan | AggregateExec: mode=SinglePartitioned, gby=[a@0 as a], aggr=[min(aggr_optimize_not.greptime_value)] | -| | CooperativeExec | -| | MergeScanExec: REDACTED +| physical_plan | AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[min(aggr_optimize_not.greptime_value)] | +| | RepartitionExec: partitioning=REDACTED +| | AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[min(aggr_optimize_not.greptime_value)] | +| | RepartitionExec: partitioning=REDACTED +| | MergeScanExec: REDACTED | | | +---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -886,8 +922,10 @@ GROUP BY +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_AggregateExec: mode=SinglePartitioned, gby=[a@0 as a], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED -|_|_|_CooperativeExec REDACTED +| 0_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[min(aggr_optimize_not.greptime_value)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[__min_state(aggr_optimize_not.greptime_value)] REDACTED @@ -1017,6 +1055,7 @@ EXPLAIN SELECT COUNT(DISTINCT val_col_1) FROM step_aggr_extended; |_|_AggregateExec: mode=FinalPartitioned, gby=[alias1@0 as alias1], aggr=[]_| |_|_RepartitionExec: partitioning=REDACTED |_|_AggregateExec: mode=Partial, gby=[val_col_1@0 as alias1], aggr=[]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_ProjectionExec: expr=[val_col_1@2 as val_col_1]_| |_|_MergeScanExec: REDACTED |_|_| @@ -1060,6 +1099,7 @@ EXPLAIN SELECT pk_col_2, sum(val_col_1) FROM step_aggr_extended GROUP BY pk_col_ |_|_AggregateExec: mode=FinalPartitioned, gby=[pk_col_2@0 as pk_col_2], aggr=[sum(step_aggr_extended.val_col_1)]_| |_|_RepartitionExec: partitioning=REDACTED |_|_AggregateExec: mode=Partial, gby=[pk_col_2@0 as pk_col_2], aggr=[sum(step_aggr_extended.val_col_1)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -1096,6 +1136,7 @@ EXPLAIN SELECT SUM(val_col_3), COUNT(val_col_2), COUNT(val_col_3), COUNT(*) FROM |_|_AggregateExec: mode=Final, gby=[], aggr=[sum(step_aggr_extended.val_col_3), count(step_aggr_extended.val_col_2), count(step_aggr_extended.val_col_3), count(Int64(1))]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[sum(step_aggr_extended.val_col_3), count(step_aggr_extended.val_col_2), count(step_aggr_extended.val_col_3), count(Int64(1))]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -1130,6 +1171,7 @@ EXPLAIN SELECT MIN(pk_col_1), MAX(val_col_2) FROM step_aggr_extended; | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[min(step_aggr_extended.pk_col_1), max(step_aggr_extended.val_col_2)]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[min(step_aggr_extended.pk_col_1), max(step_aggr_extended.val_col_2)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -1167,6 +1209,7 @@ EXPLAIN SELECT SUM(val_col_1), COUNT(*) FROM step_aggr_extended WHERE pk_col_1 = |_|_AggregateExec: mode=Final, gby=[], aggr=[sum(step_aggr_extended.val_col_1), count(Int64(1))]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[sum(step_aggr_extended.val_col_1), count(Int64(1))]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ diff --git a/tests/cases/distributed/explain/step_aggr_basic.result b/tests/cases/distributed/explain/step_aggr_basic.result index 6a002115438..cf11b7f9ef0 100644 --- a/tests/cases/distributed/explain/step_aggr_basic.result +++ b/tests/cases/distributed/explain/step_aggr_basic.result @@ -60,6 +60,7 @@ FROM | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[count(integers.i)]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[count(integers.i)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -85,6 +86,7 @@ FROM | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(integers.i)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(integers.i)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(integers.i)] REDACTED @@ -156,6 +158,7 @@ ORDER BY |_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[count(integers.i)]_| |_|_RepartitionExec: partitioning=REDACTED |_|_AggregateExec: mode=Partial, gby=[ts@0 as ts], aggr=[count(integers.i)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -189,6 +192,7 @@ ORDER BY |_|_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[count(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[ts@0 as ts], aggr=[count(integers.i)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[__count_state(integers.i)] REDACTED @@ -262,6 +266,7 @@ ORDER BY |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("1 hour"),integers.ts)@0 as date_bin(Utf8("1 hour"),integers.ts)], aggr=[count(integers.i)]_| |_|_RepartitionExec: partitioning=REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("1 hour"),integers.ts)@0 as date_bin(Utf8("1 hour"),integers.ts)], aggr=[count(integers.i)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -296,6 +301,7 @@ ORDER BY |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("1 hour"),integers.ts)@0 as date_bin(Utf8("1 hour"),integers.ts)], aggr=[count(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("1 hour"),integers.ts)@0 as date_bin(Utf8("1 hour"),integers.ts)], aggr=[count(integers.i)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("1 hour"),integers.ts)@0 as date_bin(Utf8("1 hour"),integers.ts)], aggr=[__count_state(integers.i)] REDACTED @@ -377,6 +383,7 @@ ORDER BY |_|_AggregateExec: mode=FinalPartitioned, gby=[integers.ts + Int64(1)@0 as integers.ts + Int64(1), integers.i / Int64(2)@1 as integers.i / Int64(2)], aggr=[count(integers.i)] | |_|_RepartitionExec: partitioning=REDACTED |_|_AggregateExec: mode=Partial, gby=[integers.ts + Int64(1)@0 as integers.ts + Int64(1), integers.i / Int64(2)@1 as integers.i / Int64(2)], aggr=[count(integers.i)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -413,6 +420,7 @@ ORDER BY |_|_|_AggregateExec: mode=FinalPartitioned, gby=[integers.ts + Int64(1)@0 as integers.ts + Int64(1), integers.i / Int64(2)@1 as integers.i / Int64(2)], aggr=[count(integers.i)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[integers.ts + Int64(1)@0 as integers.ts + Int64(1), integers.i / Int64(2)@1 as integers.i / Int64(2)], aggr=[count(integers.i)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[integers.ts + Int64(1)@0 as integers.ts + Int64(1), integers.i / Int64(2)@1 as integers.i / Int64(2)], aggr=[__count_state(integers.i)] REDACTED @@ -500,6 +508,7 @@ FROM |_|_AggregateExec: mode=Final, gby=[], aggr=[uddsketch_merge(Int64(128),Float64(0.01),sink_table.udd_state), hll_merge(sink_table.hll_state)]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[uddsketch_merge(Int64(128),Float64(0.01),sink_table.udd_state), hll_merge(sink_table.hll_state)]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -527,6 +536,7 @@ FROM |_|_|_AggregateExec: mode=Final, gby=[], aggr=[uddsketch_merge(Int64(128),Float64(0.01),sink_table.udd_state), hll_merge(sink_table.hll_state)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[uddsketch_merge(Int64(128),Float64(0.01),sink_table.udd_state), hll_merge(sink_table.hll_state)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__uddsketch_merge_state(Int64(128),Float64(0.01),sink_table.udd_state), __hll_merge_state(sink_table.hll_state)] REDACTED diff --git a/tests/cases/distributed/explain/subqueries.result b/tests/cases/distributed/explain/subqueries.result index e6f7286535e..d7ad6ed920c 100644 --- a/tests/cases/distributed/explain/subqueries.result +++ b/tests/cases/distributed/explain/subqueries.result @@ -101,9 +101,9 @@ order by t.i desc; | physical_plan | SortPreservingMergeExec: [i@0 DESC]_| |_|_SortExec: expr=[i@0 DESC], preserve_partitioning=[true]_| |_|_CrossJoinExec_| -|_|_CoalescePartitionsExec_| |_|_ProjectionExec: expr=[i@0 as i]_| |_|_MergeScanExec: REDACTED +|_|_RepartitionExec: partitioning=REDACTED |_|_ProjectionExec: expr=[]_| |_|_MergeScanExec: REDACTED |_|_| @@ -179,7 +179,9 @@ EXPLAIN SELECT * FROM integers i1 WHERE EXISTS(SELECT i FROM integers WHERE i=i1 | physical_plan | SortPreservingMergeExec: [i@0 ASC NULLS LAST]_| |_|_SortExec: expr=[i@0 ASC NULLS LAST], preserve_partitioning=[true]_| |_|_REDACTED +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -212,7 +214,9 @@ EXPLAIN SELECT * FROM integers i1 WHERE EXISTS(SELECT count(i) FROM integers WHE | physical_plan | SortPreservingMergeExec: [i@0 ASC NULLS LAST]_| |_|_SortExec: expr=[i@0 ASC NULLS LAST], preserve_partitioning=[true]_| |_|_REDACTED +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED +|_|_RepartitionExec: partitioning=REDACTED |_|_ProjectionExec: expr=[i@1 as i]_| |_|_MergeScanExec: REDACTED |_|_| @@ -363,7 +367,10 @@ EXPLAIN SELECT DISTINCT x FROM (SELECT a AS x FROM t) sq ORDER BY x; |_| ]]_| | physical_plan | SortPreservingMergeExec: [x@0 ASC NULLS LAST]_| |_|_SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[true]_| -|_|_AggregateExec: mode=SinglePartitioned, gby=[x@0 as x], aggr=[]_| +|_|_AggregateExec: mode=FinalPartitioned, gby=[x@0 as x], aggr=[]_| +|_|_RepartitionExec: partitioning=REDACTED +|_|_AggregateExec: mode=Partial, gby=[x@0 as x], aggr=[]_| +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -560,7 +567,9 @@ ORDER BY u1.a; | physical_plan | SortPreservingMergeExec: [a@0 ASC NULLS LAST]_| |_|_SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true]_| |_|_REDACTED +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED +|_|_RepartitionExec: partitioning=REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ diff --git a/tests/cases/distributed/flow-tql/tsid_on_phy.result b/tests/cases/distributed/flow-tql/tsid_on_phy.result index 857e45d16ad..68dd01374c1 100644 --- a/tests/cases/distributed/flow-tql/tsid_on_phy.result +++ b/tests/cases/distributed/flow-tql/tsid_on_phy.result @@ -119,8 +119,8 @@ TQL EXPLAIN ( | | TableScan: phy projection=[ts, v, tag1, tag2, le, tag4, tag5, tag6, tag7, tag8, __table_id, __tsid], partial_filters=[phy.ts >= TimestampMillisecond(1769137200001, None), phy.ts <= TimestampMillisecond(1769139900000, None), phy.__table_id=UInt32(REDACTED)] | | | ]] | | physical_plan | HistogramFoldExec: le=@0, field=@4, quantile=0.5 | -| | SortExec: expr=[tag4@1 ASC NULLS LAST, tag5@2 ASC NULLS LAST, ts@3 ASC NULLS LAST, CAST(le@0 AS Float64) ASC NULLS LAST], preserve_partitioning=[true] | -| | RepartitionExec: REDACTED +| | RepartitionExec: REDACTED +| | SortExec: expr=[tag4@1 ASC NULLS LAST, tag5@2 ASC NULLS LAST, ts@3 ASC NULLS LAST, CAST(le@0 AS Float64) ASC NULLS LAST], preserve_partitioning=[false] | | | MergeScanExec: REDACTED | | | +---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ diff --git a/tests/cases/distributed/optimizer/count.result b/tests/cases/distributed/optimizer/count.result index a339bc3bf42..7368ad6bbd7 100644 --- a/tests/cases/distributed/optimizer/count.result +++ b/tests/cases/distributed/optimizer/count.result @@ -281,6 +281,7 @@ select count(1) from count_where_bug; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[{count[count]:REDACTED} as __count_state(count_where_bug.ts)] REDACTED @@ -316,6 +317,7 @@ select count(1) from count_where_bug where `tag` = 'b'; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED @@ -342,6 +344,7 @@ select count(1) from count_where_bug where `tag` = 'b'; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED @@ -378,6 +381,7 @@ select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED @@ -417,6 +421,7 @@ select count(1) from count_where_bug where num != 3; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED diff --git a/tests/cases/distributed/optimizer/first_value_advance.result b/tests/cases/distributed/optimizer/first_value_advance.result index 08e1b8a4112..519ec9d9c15 100644 --- a/tests/cases/distributed/optimizer/first_value_advance.result +++ b/tests/cases/distributed/optimizer/first_value_advance.result @@ -325,6 +325,7 @@ explain select first_value(ts order by ts) from t; | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -348,6 +349,7 @@ explain analyze | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -414,6 +416,7 @@ order by time_window, ordered_host; |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST], first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST], first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -447,6 +450,7 @@ order by time_window, ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST], first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST], first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__first_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -525,6 +529,7 @@ explain | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -548,6 +553,7 @@ explain analyze | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED @@ -711,6 +717,7 @@ order by time_window, ordered_host; |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -744,6 +751,7 @@ order by time_window, ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__first_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED diff --git a/tests/cases/distributed/optimizer/last_value_advance.result b/tests/cases/distributed/optimizer/last_value_advance.result index 7b131d44f37..6269c5d397f 100644 --- a/tests/cases/distributed/optimizer/last_value_advance.result +++ b/tests/cases/distributed/optimizer/last_value_advance.result @@ -325,6 +325,7 @@ explain select last_value(ts order by ts) from t; | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -348,6 +349,7 @@ explain analyze | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -414,6 +416,7 @@ order by time_window, ordered_host; |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -447,6 +450,7 @@ order by time_window, ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__last_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -525,6 +529,7 @@ explain | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -548,6 +553,7 @@ explain analyze | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED @@ -711,6 +717,7 @@ order by time_window, ordered_host; |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -744,6 +751,7 @@ order by time_window, ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__last_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED diff --git a/tests/cases/distributed/optimizer/range_select_projection.result b/tests/cases/distributed/optimizer/range_select_projection.result index c32878abbf0..451e743da9c 100644 --- a/tests/cases/distributed/optimizer/range_select_projection.result +++ b/tests/cases/distributed/optimizer/range_select_projection.result @@ -48,7 +48,7 @@ ORDER BY station, "channel", ts; | 0_| 0_|_SortExec: expr=[station@1 ASC NULLS LAST, channel@2 ASC NULLS LAST, ts@0 ASC NULLS LAST], preserve_REDACTED |_|_|_ProjectionExec: expr=[ts@1 as ts, station@2 as station, channel@3 as channel, avg(range_select_projection.value_a + range_select_projection.value_b) RANGE 10s@0 as avg_value] REDACTED |_|_|_RangeSelectExec: range_expr=[avg(range_select_projection.value_a + range_select_projection.value_b) RANGE 10s], align=5000ms, align_to=0ms, align_by=[station@1, channel@2], time_index=ts REDACTED -|_|_|_CoalescePartitionsExec REDACTED +|_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_CooperativeExec REDACTED diff --git a/tests/cases/standalone/common/aggregate/distinct.result b/tests/cases/standalone/common/aggregate/distinct.result index 5a721b7a481..5c10ee22239 100644 --- a/tests/cases/standalone/common/aggregate/distinct.result +++ b/tests/cases/standalone/common/aggregate/distinct.result @@ -169,7 +169,10 @@ EXPLAIN ANALYZE SELECT DISTINCT a FROM test ORDER BY a; +-+-+-+ | 0_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_AggregateExec: mode=SinglePartitioned, gby=[a@0 as a], aggr=[] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[] REDACTED @@ -199,7 +202,10 @@ EXPLAIN ANALYZE SELECT DISTINCT a, b FROM test ORDER BY a; +-+-+-+ | 0_| 0_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_AggregateExec: mode=SinglePartitioned, gby=[a@0 as a, b@1 as b], aggr=[] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[] REDACTED diff --git a/tests/cases/standalone/common/aggregate/multi_regions.result b/tests/cases/standalone/common/aggregate/multi_regions.result index 4b104df849d..8627b83c153 100644 --- a/tests/cases/standalone/common/aggregate/multi_regions.result +++ b/tests/cases/standalone/common/aggregate/multi_regions.result @@ -62,6 +62,7 @@ select sum(val) from t; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[sum(t.val)] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[sum(t.val)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__sum_state(t.val)] REDACTED @@ -96,6 +97,7 @@ select sum(val) from t group by idc; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[idc@0 as idc], aggr=[sum(t.val)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[idc@0 as idc], aggr=[sum(t.val)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[idc@0 as idc], aggr=[__sum_state(t.val)] REDACTED diff --git a/tests/cases/standalone/common/create/metric_engine_partition.result b/tests/cases/standalone/common/create/metric_engine_partition.result index 0a5899bd9b3..ff3fd40e940 100644 --- a/tests/cases/standalone/common/create/metric_engine_partition.result +++ b/tests/cases/standalone/common/create/metric_engine_partition.result @@ -148,7 +148,10 @@ select host, count(*) from logical_table_2 GROUP BY host ORDER BY host; | physical_plan | SortPreservingMergeExec: [host@0 ASC NULLS LAST]_| |_|_SortExec: expr=[host@0 ASC NULLS LAST], preserve_partitioning=[true]_| |_|_ProjectionExec: expr=[host@0 as host, count(Int64(1))@1 as count(*)]_| -|_|_AggregateExec: mode=SinglePartitioned, gby=[host@0 as host], aggr=[count(Int64(1))]_| +|_|_AggregateExec: mode=FinalPartitioned, gby=[host@0 as host], aggr=[count(Int64(1))]_| +|_|_RepartitionExec: REDACTED +|_|_AggregateExec: mode=Partial, gby=[host@0 as host], aggr=[count(Int64(1))]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -188,6 +191,7 @@ select ts, count(*) from logical_table_2 GROUP BY ts ORDER BY ts; |_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[count(Int64(1))]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[ts@0 as ts], aggr=[count(Int64(1))]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -283,6 +287,7 @@ select a, count(*) from logical_table_3 GROUP BY a ORDER BY a; |_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(Int64(1))]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(Int64(1))]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -351,7 +356,7 @@ EXPLAIN select count(*) from logical_table_4; | plan_type_| plan_| +-+-+ | logical_plan_| Projection: count(Int64(1)) AS count(*)_| -|_|_Aggregate: groupBy=[[]], aggr=[[__count_merge(__count_state(logical_table_4.ts)) AS count(Int64(1))]] | +|_|_Aggregate: groupBy=[[]], aggr=[[__count_merge(__count_state(logical_table_4.ts)) AS count(Int64(1))]]_| |_|_MergeScan [is_placeholder=false, remote_input=[_| |_| Aggregate: groupBy=[[]], aggr=[[__count_state(logical_table_4.ts)]]_| |_|_TableScan: logical_table_4_| @@ -360,6 +365,7 @@ EXPLAIN select count(*) from logical_table_4; |_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -399,6 +405,7 @@ select ts, count(*) from logical_table_4 GROUP BY ts ORDER BY ts; |_|_AggregateExec: mode=FinalPartitioned, gby=[ts@0 as ts], aggr=[count(Int64(1))]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[ts@0 as ts], aggr=[count(Int64(1))]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ diff --git a/tests/cases/standalone/common/cte/cte_join_build_side.result b/tests/cases/standalone/common/cte/cte_join_build_side.result new file mode 100644 index 00000000000..51ee4d98a4d --- /dev/null +++ b/tests/cases/standalone/common/cte/cte_join_build_side.result @@ -0,0 +1,116 @@ +-- Regression test for issue 6874: a bounded CTE derived from the same table +-- must become the hash-join build side in both standalone and distributed mode. +CREATE TABLE cte_join_logs ( + ts TIMESTAMP TIME INDEX, + "db" STRING, + val DOUBLE, +) WITH (append_mode = 'true'); + +Affected Rows: 0 + +INSERT INTO cte_join_logs VALUES + ('2024-01-01 00:00:00', 'a', 1.0), + ('2024-01-01 00:00:01', 'a', 2.0), + ('2024-01-01 00:00:02', 'b', 3.0), + ('2024-01-01 00:00:03', 'b', 4.0), + ('2024-01-01 00:00:04', 'c', 5.0), + ('2024-01-01 00:00:05', 'c', 6.0); + +Affected Rows: 6 + +-- Control: the plain Top-N aggregation should remain below MergeScan. +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +EXPLAIN SELECT "db", count(*) AS cnt +FROM cte_join_logs +GROUP BY "db" +ORDER BY cnt DESC +LIMIT 5; + ++---------------+----------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+----------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | Limit: skip=0, fetch=5 | +| | Sort: cnt DESC NULLS FIRST | +| | Projection: cte_join_logs.db, count(Int64(1)) AS count(*) AS cnt | +| | Aggregate: groupBy=[[cte_join_logs.db]], aggr=[[count(cte_join_logs.ts) AS count(Int64(1))]] | +| | TableScan: cte_join_logs | +| | ]] | +| physical_plan | CooperativeExec | +| | MergeScanExec: REDACTED +| | | ++---------------+----------------------------------------------------------------------------------------------------+ + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +EXPLAIN WITH top_dbs AS ( + SELECT "db", count(*) AS cnt FROM cte_join_logs + GROUP BY "db" ORDER BY cnt DESC LIMIT 5 +) +SELECT t."db", count(*) AS c +FROM cte_join_logs AS t +JOIN top_dbs td ON t."db" = td."db" +GROUP BY t."db" ORDER BY c DESC, t."db"; + ++---------------+----------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+----------------------------------------------------------------------------------------------------------+ +| logical_plan | Sort: c DESC NULLS FIRST, t.db ASC NULLS LAST | +| | Projection: t.db, count(Int64(1)) AS count(*) AS c | +| | Aggregate: groupBy=[[t.db]], aggr=[[count(Int64(1))]] | +| | Projection: t.db | +| | Inner Join: t.db = td.db | +| | Projection: t.db | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | SubqueryAlias: t | +| | Filter: cte_join_logs.db IS NOT NULL | +| | TableScan: cte_join_logs, partial_filters=[cte_join_logs.db IS NOT NULL] | +| | ]] | +| | Projection: td.db | +| | MergeScan [is_placeholder=false, remote_input=[ | +| | SubqueryAlias: td | +| | SubqueryAlias: top_dbs | +| | Filter: cte_join_logs.db IS NOT NULL | +| | Limit: skip=0, fetch=5 | +| | Sort: cnt DESC NULLS FIRST | +| | Projection: cte_join_logs.db, count(Int64(1)) AS count(*) AS cnt | +| | Aggregate: groupBy=[[cte_join_logs.db]], aggr=[[count(cte_join_logs.ts) AS count(Int64(1))]] | +| | TableScan: cte_join_logs | +| | ]] | +| physical_plan | SortPreservingMergeExec: [c@1 DESC, db@0 ASC NULLS LAST] | +| | SortExec: expr=[c@1 DESC, db@0 ASC NULLS LAST], preserve_REDACTED +| | ProjectionExec: expr=[db@0 as db, count(Int64(1))@1 as c] | +| | AggregateExec: mode=FinalPartitioned, gby=[db@0 as db], aggr=[count(Int64(1))] | +| | RepartitionExec: REDACTED +| | AggregateExec: mode=Partial, gby=[db@0 as db], aggr=[count(Int64(1))] | +| | HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(db@0, db@0)], projection=[db@1] | +| | ProjectionExec: expr=[db@0 as db] | +| | MergeScanExec: REDACTED +| | RepartitionExec: REDACTED +| | ProjectionExec: expr=[db@1 as db] | +| | MergeScanExec: REDACTED +| | | ++---------------+----------------------------------------------------------------------------------------------------------+ + +WITH top_dbs AS ( + SELECT "db", count(*) AS cnt FROM cte_join_logs + GROUP BY "db" ORDER BY cnt DESC LIMIT 5 +) +SELECT t."db", count(*) AS c +FROM cte_join_logs AS t +JOIN top_dbs td ON t."db" = td."db" +GROUP BY t."db" ORDER BY c DESC, t."db"; + ++----+---+ +| db | c | ++----+---+ +| a | 2 | +| b | 2 | +| c | 2 | ++----+---+ + +DROP TABLE cte_join_logs; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/cte/cte_join_build_side.sql b/tests/cases/standalone/common/cte/cte_join_build_side.sql new file mode 100644 index 00000000000..f609b826042 --- /dev/null +++ b/tests/cases/standalone/common/cte/cte_join_build_side.sql @@ -0,0 +1,46 @@ +-- Regression test for issue 6874: a bounded CTE derived from the same table +-- must become the hash-join build side in both standalone and distributed mode. +CREATE TABLE cte_join_logs ( + ts TIMESTAMP TIME INDEX, + "db" STRING, + val DOUBLE, +) WITH (append_mode = 'true'); + +INSERT INTO cte_join_logs VALUES + ('2024-01-01 00:00:00', 'a', 1.0), + ('2024-01-01 00:00:01', 'a', 2.0), + ('2024-01-01 00:00:02', 'b', 3.0), + ('2024-01-01 00:00:03', 'b', 4.0), + ('2024-01-01 00:00:04', 'c', 5.0), + ('2024-01-01 00:00:05', 'c', 6.0); + +-- Control: the plain Top-N aggregation should remain below MergeScan. +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +EXPLAIN SELECT "db", count(*) AS cnt +FROM cte_join_logs +GROUP BY "db" +ORDER BY cnt DESC +LIMIT 5; + +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (partitioning.*) REDACTED +EXPLAIN WITH top_dbs AS ( + SELECT "db", count(*) AS cnt FROM cte_join_logs + GROUP BY "db" ORDER BY cnt DESC LIMIT 5 +) +SELECT t."db", count(*) AS c +FROM cte_join_logs AS t +JOIN top_dbs td ON t."db" = td."db" +GROUP BY t."db" ORDER BY c DESC, t."db"; + +WITH top_dbs AS ( + SELECT "db", count(*) AS cnt FROM cte_join_logs + GROUP BY "db" ORDER BY cnt DESC LIMIT 5 +) +SELECT t."db", count(*) AS c +FROM cte_join_logs AS t +JOIN top_dbs td ON t."db" = td."db" +GROUP BY t."db" ORDER BY c DESC, t."db"; + +DROP TABLE cte_join_logs; diff --git a/tests/cases/standalone/common/filter/hash_join_topk_dyn_filter.result b/tests/cases/standalone/common/filter/hash_join_topk_dyn_filter.result index a81fc4918bb..5798ddd07c1 100644 --- a/tests/cases/standalone/common/filter/hash_join_topk_dyn_filter.result +++ b/tests/cases/standalone/common/filter/hash_join_topk_dyn_filter.result @@ -67,11 +67,11 @@ WHERE c.tier IN ('gold', 'bronze'); | | Filter: customers.customer_id IS NOT NULL AND (customers.tier = Utf8("gold") OR customers.tier = Utf8("bronze")) | | | TableScan: customers, partial_filters=[customers.tier = Utf8("gold") OR customers.tier = Utf8("bronze"), customers.customer_id IS NOT NULL] | | | ]] | -| physical_plan | HashJoinExec: mode=Partitioned, join_type=Inner, on=[(customer_id@1, customer_id@0)], projection=[id@0, amount@2, name@4, tier@5] | -| | RepartitionExec: partitioning=REDACTED -| | ProjectionExec: expr=[id@0 as id, customer_id@1 as customer_id, amount@2 as amount] | +| physical_plan | HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(customer_id@1, customer_id@0)], projection=[id@0, amount@2, name@4, tier@5] | +| | ProjectionExec: expr=[id@0 as id, customer_id@1 as customer_id, amount@2 as amount] | +| | CooperativeExec | | | MergeScanExec: REDACTED -| | RepartitionExec: partitioning=REDACTED +| | RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 | | | ProjectionExec: expr=[customer_id@0 as customer_id, name@1 as name, tier@2 as tier] | | | MergeScanExec: REDACTED | | | @@ -107,9 +107,9 @@ WHERE c.tier IN ('gold', 'bronze'); +-+-+-+ | stage | node | plan_| +-+-+-+ -| 0_| 0_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(customer_id@1, customer_id@0)], projection=[id@0, amount@2, name@4, tier@5] metrics=REDACTED_| -|_|_|_RepartitionExec: partitioning=REDACTED +| 0_| 0_|_HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(customer_id@1, customer_id@0)], projection=[id@0, amount@2, name@4, tier@5] metrics=REDACTED_| |_|_|_ProjectionExec: expr=[id@0 as id, customer_id@1 as customer_id, amount@2 as amount] metrics=REDACTED_| +|_|_|_CooperativeExec metrics=REDACTED_| |_|_|_MergeScanExec: REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_ProjectionExec: expr=[customer_id@0 as customer_id, name@1 as name, tier@2 as tier] metrics=REDACTED_| diff --git a/tests/cases/standalone/common/order/order_by_exceptions.result b/tests/cases/standalone/common/order/order_by_exceptions.result index bde25ea927c..bdf215f8f26 100644 --- a/tests/cases/standalone/common/order/order_by_exceptions.result +++ b/tests/cases/standalone/common/order/order_by_exceptions.result @@ -82,9 +82,10 @@ EXPLAIN SELECT a % 2, b FROM test UNION SELECT a % 2 AS k, b FROM test ORDER BY | | AggregateExec: mode=FinalPartitioned, gby=[test.a % Int64(2)@0 as test.a % Int64(2), b@1 as b], aggr=[] | | | RepartitionExec: REDACTED | | AggregateExec: mode=Partial, gby=[test.a % Int64(2)@0 as test.a % Int64(2), b@1 as b], aggr=[] | -| | InterleaveExec | -| | MergeScanExec: REDACTED -| | MergeScanExec: REDACTED +| | RepartitionExec: REDACTED +| | InterleaveExec | +| | MergeScanExec: REDACTED +| | MergeScanExec: REDACTED | | | +---------------+-----------------------------------------------------------------------------------------------------------+ diff --git a/tests/cases/standalone/common/promql/histogram_multi_partition.result b/tests/cases/standalone/common/promql/histogram_multi_partition.result index 091ed7bfcc3..3d5490f38a1 100644 --- a/tests/cases/standalone/common/promql/histogram_multi_partition.result +++ b/tests/cases/standalone/common/promql/histogram_multi_partition.result @@ -47,6 +47,7 @@ tql analyze (0, 10, '10s') histogram_quantile(0.5, sum by (le) (histogram_gap_bu |_|_|_AggregateExec: mode=FinalPartitioned, gby=[le@0 as le, ts@1 as ts], aggr=[sum(histogram_gap_bucket.val)] REDACTED |_|_|_RepartitionExec: partitioning=Hash([le@0, ts@1],REDACTED |_|_|_AggregateExec: mode=Partial, gby=[le@0 as le, ts@1 as ts], aggr=[sum(histogram_gap_bucket.val)] REDACTED +|_|_|_RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=2 REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[le@0 as le, ts@1 as ts], aggr=[__sum_state(histogram_gap_bucket.val)] REDACTED diff --git a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result index 75aaf8176c9..72a6dc9c5b4 100644 --- a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result +++ b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result @@ -117,11 +117,11 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / tsid_binary_join_right; +-+-+-+ | 0_| 0_|_ProjectionExec: expr=[host@2 as host, job@3 as job, ts@5 as ts, __tsid@4 as __tsid, greptime_value@0 / greptime_value@1 as tsid_binary_join_left.greptime_value / tsid_binary_join_right.greptime_value] REDACTED |_|_|_HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(__tsid@1, __tsid@3), (ts@2, ts@4)], projection=[greptime_value@0, greptime_value@3, host@4, job@5, __tsid@6, ts@7], NullsEqual: true REDACTED -|_|_|_CoalescePartitionsExec REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, __tsid@3 as __tsid, ts@4 as ts] REDACTED -|_|_|_MergeScanExec: REDACTED |_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED @@ -363,11 +363,11 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > bool tsid_binary_join_right; +-+-+-+ | 0_| 0_|_ProjectionExec: expr=[host@2 as host, job@3 as job, ts@5 as ts, __tsid@4 as __tsid, CAST(greptime_value@1 < greptime_value@0 AS Float64) as tsid_binary_join_left.greptime_value > tsid_binary_join_right.greptime_value] REDACTED |_|_|_HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(__tsid@1, __tsid@3), (ts@2, ts@4)], projection=[greptime_value@0, greptime_value@3, host@4, job@5, __tsid@6, ts@7], NullsEqual: true REDACTED -|_|_|_CoalescePartitionsExec REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, __tsid@3 as __tsid, ts@4 as ts] REDACTED -|_|_|_MergeScanExec: REDACTED |_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED @@ -411,7 +411,7 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > tsid_binary_join_right) / tsi |_|_|_RepartitionExec: partitioning=Hash([REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, __tsid@3 as __tsid, ts@4 as ts] REDACTED |_|_|_MergeScanExec: REDACTED -|_|_|_CooperativeExec REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED @@ -466,7 +466,7 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > bool tsid_binary_join_right) |_|_|_RepartitionExec: partitioning=Hash([REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, __tsid@3 as __tsid, ts@4 as ts] REDACTED |_|_|_MergeScanExec: REDACTED -|_|_|_CooperativeExec REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED @@ -516,12 +516,12 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi |_|_|_HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(__tsid@1, __tsid@3), (ts@0, ts@4)], projection=[greptime_value@2, greptime_value@3, host@4, job@5, __tsid@6, ts@7], NullsEqual: true REDACTED |_|_|_ProjectionExec: expr=[ts@0 as ts, __tsid@1 as __tsid, greptime_value@2 as greptime_value] REDACTED |_|_|_UnionDistinctOnExec: on col=[5, 6], ts_col=0 REDACTED -|_|_|_CoalescePartitionsExec REDACTED -|_|_|_MergeScanExec: REDACTED -|_|_|_CoalescePartitionsExec REDACTED +|_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[ts@2 as ts, __tsid@3 as __tsid, greptime_value@4 as greptime_value, host@5 as host, job@6 as job, CASE WHEN __common_expr_1@0 IS NOT NULL THEN __common_expr_1@0 ELSE_END as __promql_or_match_0, CASE WHEN __common_expr_2@1 IS NOT NULL THEN __common_expr_2@1 ELSE_END as __promql_or_match_1] REDACTED |_|_|_ProjectionExec: expr=[CAST(host@1 AS Utf8) as __common_expr_1, CAST(job@2 AS Utf8) as __common_expr_2, ts@4 as ts, __tsid@3 as __tsid, greptime_value@0 as greptime_value, host@1 as host, job@2 as job] REDACTED @@ -575,7 +575,7 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left / ignoring(host) group_left tsid |_|_|_RepartitionExec: partitioning=Hash([REDACTED |_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@2 as job, __tsid@3 as __tsid, ts@4 as ts] REDACTED |_|_|_MergeScanExec: REDACTED -|_|_|_CooperativeExec REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED diff --git a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result index fc883d96e71..ba2e97bcc2d 100644 --- a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result +++ b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result @@ -109,8 +109,8 @@ TQL ANALYZE (0, 10, '5s') histogram_quantile(0.5, tsid_no_aggr_histogram_bucket) | stage | node | plan_| +-+-+-+ | 0_| 0_|_HistogramFoldExec: le=@2, field=@0, quantile=0.5 REDACTED -|_|_|_SortExec: expr=[job@1 ASC NULLS LAST, ts@3 ASC NULLS LAST, CAST(le@2 AS Float64) ASC NULLS LAST], preserve_partitioning=[true] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_SortExec: expr=[job@1 ASC NULLS LAST, ts@3 ASC NULLS LAST, CAST(le@2 AS Float64) ASC NULLS LAST], preserve_partitioning=[false] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[val@0 as val, job@1 as job, le@2 as le, ts@4 as ts] REDACTED diff --git a/tests/cases/standalone/common/range/nest.result b/tests/cases/standalone/common/range/nest.result index 74742d5d862..184b27545d0 100644 --- a/tests/cases/standalone/common/range/nest.result +++ b/tests/cases/standalone/common/range/nest.result @@ -138,7 +138,7 @@ EXPLAIN SELECT ts, host, min(val) RANGE '5s' FROM host ALIGN '5s'; |_|_TableScan: host_| |_| ]]_| | physical_plan | RangeSelectExec: range_expr=[min(host.val) RANGE 5s], align=5000ms, align_to=0ms, align_by=[host@1], time_index=ts | -|_|_CoalescePartitionsExec_| +|_|_CooperativeExec_| |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -154,7 +154,7 @@ EXPLAIN ANALYZE SELECT ts, host, min(val) RANGE '5s' FROM host ALIGN '5s'; | stage | node | plan_| +-+-+-+ | 0_| 0_|_RangeSelectExec: range_expr=[min(host.val) RANGE 5s], align=5000ms, align_to=0ms, align_by=[host@1], time_index=ts REDACTED -|_|_|_CoalescePartitionsExec REDACTED +|_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_CooperativeExec REDACTED diff --git a/tests/cases/standalone/common/tql/partition.result b/tests/cases/standalone/common/tql/partition.result index 45927f8abbd..f1072edbb32 100644 --- a/tests/cases/standalone/common/tql/partition.result +++ b/tests/cases/standalone/common/tql/partition.result @@ -68,7 +68,10 @@ tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_SortPreservingMergeExec: [k@0 ASC NULLS LAST, j@1 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[k@0 ASC NULLS LAST, j@1 ASC NULLS LAST], preserve_partitioning=[true] REDACTED -|_|_|_AggregateExec: mode=SinglePartitioned, gby=[k@0 as k, j@1 as j], aggr=[avg(prom_irate(j_range,i))] REDACTED +|_|_|_AggregateExec: mode=FinalPartitioned, gby=[k@0 as k, j@1 as j], aggr=[avg(prom_irate(j_range,i))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_AggregateExec: mode=Partial, gby=[k@0 as k, j@1 as j], aggr=[avg(prom_irate(j_range,i))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[k@0 as k, j@1 as j], aggr=[__avg_state(prom_irate(j_range,i))] REDACTED diff --git a/tests/cases/standalone/common/tql/tql-cte.result b/tests/cases/standalone/common/tql/tql-cte.result index 0547754b483..8efe983763b 100644 --- a/tests/cases/standalone/common/tql/tql-cte.result +++ b/tests/cases/standalone/common/tql/tql-cte.result @@ -487,8 +487,8 @@ ORDER BY ts; | | ]] | | physical_plan | ProjectionExec: expr=[ts@0 as ts, val@1 as val, lag(tql_base.val,Int64(1)) ORDER BY [tql_base.ts ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as prev_value] | | | BoundedWindowAggExec: wdw=[lag(tql_base.val,Int64(1)) ORDER BY [tql_base.ts ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "lag(tql_base.val,Int64(1)) ORDER BY [tql_base.ts ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Float64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] | -| | SortPreservingMergeExec: [ts@0 ASC NULLS LAST] | -| | SortExec: expr=[ts@0 ASC NULLS LAST], preserve_REDACTED +| | SortExec: expr=[ts@0 ASC NULLS LAST], preserve_REDACTED +| | CooperativeExec | | | MergeScanExec: REDACTED | | | +---------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -765,11 +765,15 @@ LIMIT 5; | | ProjectionExec: expr=[ts@0 as ts, cpu@1 as avg_value, host@2 as host] | | | HashJoinExec: mode=Partitioned, join_type=Inner, on=[(date_trunc(Utf8("second"),t.ts)@2, date_trunc(Utf8("second"),l.ts)@2)], projection=[ts@0, cpu@1, host@4] | | | RepartitionExec: REDACTED -| | ProjectionExec: expr=[ts@0 as ts, cpu@2 as cpu, date_trunc(second, ts@0) as date_trunc(Utf8("second"),t.ts)] | -| | MergeScanExec: REDACTED +| | ProjectionExec: expr=[ts@0 as ts, cpu@1 as cpu, date_trunc(second, ts@0) as date_trunc(Utf8("second"),t.ts)] | +| | RepartitionExec: REDACTED +| | ProjectionExec: expr=[ts@0 as ts, cpu@2 as cpu] | +| | MergeScanExec: REDACTED | | RepartitionExec: REDACTED | | ProjectionExec: expr=[ts@0 as ts, host@1 as host, date_trunc(second, ts@0) as date_trunc(Utf8("second"),l.ts)] | -| | MergeScanExec: REDACTED +| | RepartitionExec: REDACTED +| | ProjectionExec: expr=[ts@0 as ts, host@1 as host] | +| | MergeScanExec: REDACTED | | | +---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -843,11 +847,15 @@ LIMIT 5; | | ProjectionExec: expr=[ts@1 as ts, cpu@0 as avg_value, host@2 as host] | | | HashJoinExec: mode=Partitioned, join_type=Inner, on=[(date_trunc(Utf8("second"),t.ts)@2, date_trunc(Utf8("second"),l.ts)@2)], projection=[cpu@0, ts@1, host@4] | | | RepartitionExec: REDACTED -| | ProjectionExec: expr=[cpu@0 as cpu, ts@2 as ts, date_trunc(second, ts@2) as date_trunc(Utf8("second"),t.ts)] | -| | MergeScanExec: REDACTED +| | ProjectionExec: expr=[cpu@0 as cpu, ts@1 as ts, date_trunc(second, ts@1) as date_trunc(Utf8("second"),t.ts)] | +| | RepartitionExec: REDACTED +| | ProjectionExec: expr=[cpu@0 as cpu, ts@2 as ts] | +| | MergeScanExec: REDACTED | | RepartitionExec: REDACTED | | ProjectionExec: expr=[ts@0 as ts, host@1 as host, date_trunc(second, ts@0) as date_trunc(Utf8("second"),l.ts)] | -| | MergeScanExec: REDACTED +| | RepartitionExec: REDACTED +| | ProjectionExec: expr=[ts@0 as ts, host@1 as host] | +| | MergeScanExec: REDACTED | | | +---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------+ diff --git a/tests/cases/standalone/optimizer/count.result b/tests/cases/standalone/optimizer/count.result index 99f910136fd..8a128a92000 100644 --- a/tests/cases/standalone/optimizer/count.result +++ b/tests/cases/standalone/optimizer/count.result @@ -278,6 +278,7 @@ select count(1) from count_where_bug; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[{count[count]:REDACTED} as __count_state(count_where_bug.ts)] REDACTED @@ -313,6 +314,7 @@ select count(1) from count_where_bug where `tag` = 'b'; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED @@ -339,6 +341,7 @@ select count(1) from count_where_bug where `tag` = 'b'; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED @@ -375,6 +378,7 @@ select count(1) from count_where_bug where ts > '2024-09-06T06:00:04Z'; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED @@ -414,6 +418,7 @@ select count(1) from count_where_bug where num != 3; | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(count_where_bug.ts)] REDACTED diff --git a/tests/cases/standalone/optimizer/first_value_advance.result b/tests/cases/standalone/optimizer/first_value_advance.result index afcefae4cff..c543a8df2a2 100644 --- a/tests/cases/standalone/optimizer/first_value_advance.result +++ b/tests/cases/standalone/optimizer/first_value_advance.result @@ -322,6 +322,7 @@ explain select first_value(ts order by ts) from t; | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -345,6 +346,7 @@ explain analyze | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -411,6 +413,7 @@ order by time_window, ordered_host; |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST], first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST], first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -444,6 +447,7 @@ order by time_window, ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST], first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[first_value(t.host) ORDER BY [t.ts ASC NULLS LAST], first_value(t.val) ORDER BY [t.ts ASC NULLS LAST], first_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__first_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __first_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -522,6 +526,7 @@ explain | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -545,6 +550,7 @@ explain analyze | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED @@ -705,6 +711,7 @@ order by time_window, ordered_host; |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -738,6 +745,7 @@ order by time_window, ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[first_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], first_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__first_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __first_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED diff --git a/tests/cases/standalone/optimizer/last_value_advance.result b/tests/cases/standalone/optimizer/last_value_advance.result index a1882a75887..3199692a73d 100644 --- a/tests/cases/standalone/optimizer/last_value_advance.result +++ b/tests/cases/standalone/optimizer/last_value_advance.result @@ -322,6 +322,7 @@ explain select last_value(ts order by ts) from t; | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -345,6 +346,7 @@ explain analyze | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -411,6 +413,7 @@ order by time_window, ordered_host; |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -444,6 +447,7 @@ order by time_window, ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[last_value(t.host) ORDER BY [t.ts ASC NULLS LAST], last_value(t.val) ORDER BY [t.ts ASC NULLS LAST], last_value(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t.ts)@0 as date_bin(Utf8("5 milliseconds"),t.ts)], aggr=[__last_value_state(t.host) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.val) ORDER BY [t.ts ASC NULLS LAST], __last_value_state(t.ts) ORDER BY [t.ts ASC NULLS LAST]] REDACTED @@ -522,6 +526,7 @@ explain | physical_plan | AggregateExec: mode=Final, gby=[], aggr=[last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| |_|_CoalescePartitionsExec_| |_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -545,6 +550,7 @@ explain analyze | 0_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED @@ -705,6 +711,7 @@ order by time_window, ordered_host; |_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| |_|_RepartitionExec: REDACTED |_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]]_| +|_|_RepartitionExec: REDACTED |_|_MergeScanExec: REDACTED |_|_| +-+-+ @@ -738,6 +745,7 @@ order by time_window, ordered_host; |_|_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED |_|_|_RepartitionExec: REDACTED |_|_|_AggregateExec: mode=Partial, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[last_value(t1.host) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.val) ORDER BY [t1.ts ASC NULLS LAST], last_value(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED +|_|_|_RepartitionExec: REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[date_bin(Utf8("5 milliseconds"),t1.ts)@0 as date_bin(Utf8("5 milliseconds"),t1.ts)], aggr=[__last_value_state(t1.host) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.val) ORDER BY [t1.ts ASC NULLS LAST], __last_value_state(t1.ts) ORDER BY [t1.ts ASC NULLS LAST]] REDACTED diff --git a/tests/cases/standalone/optimizer/rewrite_set_comparison.result b/tests/cases/standalone/optimizer/rewrite_set_comparison.result index eb41d5f78e8..6d9e04458c7 100644 --- a/tests/cases/standalone/optimizer/rewrite_set_comparison.result +++ b/tests/cases/standalone/optimizer/rewrite_set_comparison.result @@ -81,11 +81,12 @@ EXPLAIN SELECT v FROM sc_t WHERE v > ANY(SELECT v FROM sc_s) ORDER BY v; | | CoalescePartitionsExec | | | FilterExec: mark@1 OR NOT mark@1 AND NULL | | | NestedLoopJoinExec: join_type=LeftMark, filter=(v@0 > v@1) IS NOT DISTINCT FROM true | -| | CoalescePartitionsExec | -| | ProjectionExec: expr=[v@1 as v] | -| | MergeScanExec: REDACTED -| | MergeScanExec: REDACTED -| | MergeScanExec: REDACTED +| | ProjectionExec: expr=[v@1 as v] | +| | MergeScanExec: REDACTED +| | RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 | +| | MergeScanExec: REDACTED +| | RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 | +| | MergeScanExec: REDACTED | | | +---------------+---------------------------------------------------------------------------------------------------------------+ @@ -135,11 +136,12 @@ EXPLAIN SELECT v FROM sc_t WHERE v != ALL(SELECT v FROM sc_s) ORDER BY v; | | CoalescePartitionsExec | | | FilterExec: NOT mark@1, projection=[v@0] | | | NestedLoopJoinExec: join_type=LeftMark, filter=(v@0 != v@1) IS NOT DISTINCT FROM false | -| | CoalescePartitionsExec | -| | ProjectionExec: expr=[v@1 as v] | -| | MergeScanExec: REDACTED -| | MergeScanExec: REDACTED -| | MergeScanExec: REDACTED +| | ProjectionExec: expr=[v@1 as v] | +| | MergeScanExec: REDACTED +| | RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 | +| | MergeScanExec: REDACTED +| | RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 | +| | MergeScanExec: REDACTED | | | +---------------+----------------------------------------------------------------------------------------------------+ diff --git a/tests/cases/standalone/tql-explain-analyze/analyze.result b/tests/cases/standalone/tql-explain-analyze/analyze.result index 1753d59da91..203debf90bf 100644 --- a/tests/cases/standalone/tql-explain-analyze/analyze.result +++ b/tests/cases/standalone/tql-explain-analyze/analyze.result @@ -288,6 +288,7 @@ TQL ANALYZE sum(test2); |_|_|_AggregateExec: mode=FinalPartitioned, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[sum(test2.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED |_|_|_AggregateExec: mode=Partial, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[sum(test2.greptime_value)] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=FinalPartitioned, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[__sum_state(test2.greptime_value)] REDACTED diff --git a/tests/cases/standalone/tql-explain-analyze/tsid_column.result b/tests/cases/standalone/tql-explain-analyze/tsid_column.result index 4a7a8750605..8d27a55ab38 100644 --- a/tests/cases/standalone/tql-explain-analyze/tsid_column.result +++ b/tests/cases/standalone/tql-explain-analyze/tsid_column.result @@ -102,10 +102,10 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(count(tsid | 0_| 0_|_ProjectionExec: expr=[ts@1 as ts, sum(prom_irate(ts_range,val))@2 / scalar(count(count(tsid_metric.val)))@0 as lhs.sum(prom_irate(ts_range,val)) / rhs.scalar(count(count(tsid_metric.val)))] REDACTED |_|_|_REDACTED |_|_|_ScalarCalculateExec: tags=[] REDACTED -|_|_|_CoalescePartitionsExec REDACTED -|_|_|_MergeScanExec: REDACTED |_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ts@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED @@ -153,10 +153,10 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(sum(tsid_m | 0_| 0_|_ProjectionExec: expr=[ts@1 as ts, sum(prom_irate(ts_range,val))@2 / scalar(count(sum(tsid_metric.val)))@0 as lhs.sum(prom_irate(ts_range,val)) / rhs.scalar(count(sum(tsid_metric.val)))] REDACTED |_|_|_REDACTED |_|_|_ScalarCalculateExec: tags=[] REDACTED -|_|_|_CoalescePartitionsExec REDACTED -|_|_|_MergeScanExec: REDACTED |_|_|_CooperativeExec REDACTED |_|_|_MergeScanExec: REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED +|_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ts@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED diff --git a/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result b/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result index 07d0b8c764c..45aac8d9c39 100644 --- a/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result +++ b/tests/compatibility/cases/analyze_verbose_remote_metrics_extension/verify.result @@ -17,6 +17,7 @@ EXPLAIN ANALYZE VERBOSE SELECT count(*) FROM t_analyze_verbose_remote_metrics_ex |_|_|_AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] REDACTED |_|_|_CoalescePartitionsExec REDACTED |_|_|_AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] REDACTED +|_|_|_RepartitionExec: partitioning=REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[__count_state(t_analyze_verbose_remote_REDACTED diff --git a/tests/conf/datanode-test.toml.template b/tests/conf/datanode-test.toml.template index 63d3be2f8c4..907fcee4ddf 100644 --- a/tests/conf/datanode-test.toml.template +++ b/tests/conf/datanode-test.toml.template @@ -3,6 +3,9 @@ mode = 'distributed' require_lease_before_startup = true rpc_runtime_size = 8 +[query] +parallelism = 4 + [[region_engine]] [region_engine.mito] {{ if enable_flat_format }} diff --git a/tests/conf/frontend-test.toml.template b/tests/conf/frontend-test.toml.template index 0042716bd80..632947a3cde 100644 --- a/tests/conf/frontend-test.toml.template +++ b/tests/conf/frontend-test.toml.template @@ -2,5 +2,8 @@ bind_addr = "{addrs.grpc_addr}" server_addr = "{addrs.grpc_addr}" +[query] +parallelism = 4 + [event_recorder] event_types = [] diff --git a/tests/conf/standalone-test.toml.template b/tests/conf/standalone-test.toml.template index 88e3b1018cd..ebe25f82487 100644 --- a/tests/conf/standalone-test.toml.template +++ b/tests/conf/standalone-test.toml.template @@ -2,6 +2,9 @@ mode = 'standalone' enable_memory_catalog = false require_lease_before_startup = true +[query] +parallelism = 4 + [[region_engine]] [region_engine.mito] {{ if enable_flat_format }} From 354921c80eaf13c975b38e8cdd6f6767001f5a8e Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Thu, 13 Aug 2026 10:58:29 +0000 Subject: [PATCH 41/92] chore: update mysql test drivers and lru (#8868) * chore: update mysql test drivers and lru * fix: test --- Cargo.lock | 192 +++++++++++-------- src/servers/Cargo.toml | 5 +- src/servers/tests/mysql/mysql_server_test.rs | 2 +- tests-integration/Cargo.toml | 2 +- tests/runner/Cargo.toml | 2 +- 5 files changed, 123 insertions(+), 80 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a7a7ff9b877..ece381b7b84 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -424,7 +424,7 @@ checksum = "0c6cd424c2693bcdbc150d843dc9d4d137dd2de4782ce6df491ad11a3a0416c0" dependencies = [ "bytes", "half", - "num-bigint", + "num-bigint 0.4.6", "num-traits", ] @@ -1254,7 +1254,7 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6773ddc0eafc0e509fb60e48dff7f450f8e674a0686ae8605e8d9901bd5eefa" dependencies = [ - "num-bigint", + "num-bigint 0.4.6", "num-integer", "num-traits", "serde", @@ -1268,7 +1268,7 @@ checksum = "1a22f228ab7a1b23027ccc6c350b72868017af7ea8356fbdf19f8d991c690013" dependencies = [ "autocfg", "libm", - "num-bigint", + "num-bigint 0.4.6", "num-integer", "num-traits", ] @@ -1517,6 +1517,15 @@ dependencies = [ "num-traits", ] +[[package]] +name = "btoi" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b5ab9db53bcda568284df0fd39f6eac24ad6f7ba7ff1168b9e76eba6576b976" +dependencies = [ + "num-traits", +] + [[package]] name = "bufstream" version = "0.1.4" @@ -1620,7 +1629,7 @@ dependencies = [ "maybe-owned", "rustix 1.0.7", "rustix-linux-procfs", - "windows-sys 0.60.2", + "windows-sys 0.61.2", "winx", ] @@ -2232,7 +2241,7 @@ checksum = "fe6d2e5af09e8c8ad56c969f2157a3d4238cebc7c55f0a517728c38f7b200f81" dependencies = [ "serde", "termcolor", - "unicode-width 0.1.14", + "unicode-width 0.2.1", ] [[package]] @@ -2243,7 +2252,7 @@ checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" dependencies = [ "serde", "termcolor", - "unicode-width 0.1.14", + "unicode-width 0.2.1", ] [[package]] @@ -2259,7 +2268,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" dependencies = [ "lazy_static", - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] @@ -5126,7 +5135,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] @@ -6197,6 +6206,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" dependencies = [ "allocator-api2", + "equivalent", + "foldhash 0.2.0", ] [[package]] @@ -6539,7 +6550,7 @@ dependencies = [ "tokio", "tokio-rustls", "tower-service", - "webpki-roots 1.0.1", + "webpki-roots 1.0.9", ] [[package]] @@ -6585,7 +6596,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.5.10", + "socket2 0.6.4", "tokio", "tower-service", "tracing", @@ -7084,7 +7095,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20fd6de4ccfcc187e38bc21cfa543cb5a302cb86a8b114eb7f0bf0dc9f8ac00f" dependencies = [ "io-lifetimes 3.0.1", - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] @@ -7825,7 +7836,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667" dependencies = [ "cfg-if", - "windows-targets 0.48.5", + "windows-targets 0.53.5", ] [[package]] @@ -8099,15 +8110,6 @@ dependencies = [ "vob", ] -[[package]] -name = "lru" -version = "0.12.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "234cf4f4a04dc1f57e24b96cc0cd600cf2af460d4161ac5ecdd0af8e1f3b2a38" -dependencies = [ - "hashbrown 0.15.4", -] - [[package]] name = "lru" version = "0.16.4" @@ -8117,6 +8119,15 @@ dependencies = [ "hashbrown 0.16.1", ] +[[package]] +name = "lru" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d2f2f9b4ba7e6b24d95e7e899329d35be83bcded72c8540cdd5368932d1d90a" +dependencies = [ + "hashbrown 0.17.1", +] + [[package]] name = "lru-slab" version = "0.1.2" @@ -8734,28 +8745,29 @@ checksum = "2195bf6aa996a481483b29d62a7663eed3fe39600c460e323f8ff41e90bdd89b" [[package]] name = "mysql" -version = "26.0.0" +version = "28.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "64453aedc258ac8c720b46c8264302fad39cef6c02483f68adbad4bcd22d6fab" +checksum = "4a732193888328fc060ab901c0ed1355521267a51ffbfd9a0b3786434c6b8e7f" dependencies = [ "bufstream", "bytes", "crossbeam-queue", + "crossbeam-utils", "flate2", "io-enum", "libc", - "lru 0.12.5", - "mysql_common 0.34.1", + "lru 0.16.4", + "mysql_common 0.37.3", "named_pipe", "pem", "percent-encoding", "rustls", "rustls-pemfile", - "socket2 0.5.10", + "socket2 0.6.4", "twox-hash", "url", "webpki", - "webpki-roots 0.26.11", + "webpki-roots 1.0.9", ] [[package]] @@ -8766,7 +8778,7 @@ checksum = "63c3512cf11487168e0e9db7157801bf5273be13055a9cc95356dc9e0035e49c" dependencies = [ "darling 0.20.11", "heck 0.5.0", - "num-bigint", + "num-bigint 0.4.6", "proc-macro-crate 3.3.0", "proc-macro-error2", "proc-macro2", @@ -8784,7 +8796,7 @@ checksum = "66f62cad7623a9cb6f8f64037f0c4f69c8db8e82914334a83c9788201c2c1bfa" dependencies = [ "darling 0.20.11", "heck 0.5.0", - "num-bigint", + "num-bigint 0.4.6", "proc-macro-crate 3.3.0", "proc-macro-error2", "proc-macro2", @@ -8796,36 +8808,32 @@ dependencies = [ [[package]] name = "mysql_async" -version = "0.35.1" +version = "0.37.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d14cf024116ba8fef4a7fec5abf0bd5de89b9fb29a7e55818a119ac5ec745077" +checksum = "3519e91b0d254ac1ffa495bc42053286cb2172ad7241d5b3b1b9f8a891f21ee2" dependencies = [ "bytes", - "crossbeam", + "crossbeam-queue", + "crossbeam-utils", "flate2", "futures-core", "futures-sink", "futures-util", "keyed_priority_queue", - "lru 0.12.5", - "mysql_common 0.34.1", - "pem", + "lru 0.18.2", + "mysql_common 0.37.3", "percent-encoding", - "pin-project", - "rand 0.8.5", + "rand 0.10.1", "rustls", - "rustls-pemfile", "serde", - "serde_json", - "socket2 0.5.10", + "socket2 0.6.4", "thiserror 2.0.17", "tokio", "tokio-rustls", "tokio-util", "twox-hash", "url", - "webpki", - "webpki-roots 0.26.11", + "webpki-roots 1.0.9", ] [[package]] @@ -8839,7 +8847,7 @@ dependencies = [ "bindgen", "bitflags 2.12.1", "bitvec", - "btoi", + "btoi 0.4.3", "byteorder", "bytes", "cc", @@ -8850,7 +8858,7 @@ dependencies = [ "frunk", "lazy_static", "mysql-common-derive 0.31.2", - "num-bigint", + "num-bigint 0.4.6", "num-traits", "rand 0.8.5", "regex", @@ -8870,35 +8878,58 @@ dependencies = [ [[package]] name = "mysql_common" -version = "0.34.1" +version = "0.37.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34a9141e735d5bb02414a7ac03add09522466d4db65bdd827069f76ae0850e58" +checksum = "0f27695f286b461da077b8c2f72f47feaa04ce3c3f9c0976257410e90e21208a" dependencies = [ "base64 0.22.1", "bitflags 2.12.1", - "btoi", + "btoi 0.5.0", "byteorder", "bytes", - "cc", - "cmake", "crc32fast", "flate2", - "lazy_static", + "getrandom 0.3.4", "mysql-common-derive 0.32.1", - "num-bigint", + "num-bigint 0.4.6", "num-traits", - "rand 0.8.5", "regex", "saturating", "serde", "serde_json", "sha1 0.10.6", "sha2 0.10.9", - "subprocess", - "thiserror 1.0.69", + "thiserror 2.0.17", "time", "uuid", - "zstd", +] + +[[package]] +name = "mysql_common" +version = "0.38.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d11189bb0e285f58b71744ee23aa8b943bd5120ef0877555813d2107374d3970" +dependencies = [ + "base64 0.22.1", + "bitflags 2.12.1", + "btoi 0.5.0", + "byteorder", + "bytes", + "crc32fast", + "flate2", + "getrandom 0.4.1", + "hex", + "mysql-common-derive 0.32.1", + "num-bigint 0.5.1", + "num-traits", + "regex", + "saturating", + "serde", + "serde_json", + "sha1 0.11.0", + "sha2 0.11.0", + "thiserror 2.0.17", + "uuid", ] [[package]] @@ -9094,7 +9125,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -9103,7 +9134,7 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23" dependencies = [ - "num-bigint", + "num-bigint 0.4.6", "num-complex", "num-integer", "num-iter", @@ -9121,6 +9152,16 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-bigint" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-bigint-dig" version = "0.8.6" @@ -9205,7 +9246,7 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" dependencies = [ - "num-bigint", + "num-bigint 0.4.6", "num-integer", "num-traits", ] @@ -9288,7 +9329,7 @@ version = "0.7.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77e878c846a8abae00dd069496dbe8751b16ac1c3d6bd2a7283a938e8228f90d" dependencies = [ - "proc-macro-crate 1.3.1", + "proc-macro-crate 3.3.0", "proc-macro2", "quote", "syn 2.0.117", @@ -10202,7 +10243,7 @@ dependencies = [ "half", "hashbrown 0.17.1", "lz4_flex 0.13.1", - "num-bigint", + "num-bigint 0.4.6", "num-integer", "num-traits", "object_store", @@ -11223,7 +11264,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4" dependencies = [ "bytes", - "heck 0.4.1", + "heck 0.5.0", "itertools 0.12.1", "log", "multimap", @@ -11243,8 +11284,8 @@ version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ - "heck 0.4.1", - "itertools 0.10.5", + "heck 0.5.0", + "itertools 0.14.0", "log", "multimap", "once_cell", @@ -11292,7 +11333,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.10.5", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -11305,7 +11346,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9120690fafc389a67ba3803df527d0ec9cbbc9cc45e4cc20b332996dfb672425" dependencies = [ "anyhow", - "itertools 0.10.5", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -12300,7 +12341,7 @@ dependencies = [ "wasm-bindgen-futures", "wasm-streams 0.4.2", "web-sys", - "webpki-roots 1.0.1", + "webpki-roots 1.0.9", ] [[package]] @@ -12824,7 +12865,7 @@ dependencies = [ "security-framework 3.7.0", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -13364,7 +13405,7 @@ dependencies = [ "mime_guess", "moka", "mysql_async", - "mysql_common 0.34.1", + "mysql_common 0.38.2", "notify", "object-pool", "once_cell", @@ -13397,6 +13438,7 @@ dependencies = [ "rustls-pki-types", "serde", "serde_json", + "servers", "session", "simd-json", "simdutf8", @@ -13708,7 +13750,7 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1961e2ef424c1424204d3a5d6975f934f56b6d50ff5732382d84ebf460e147f7" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "proc-macro2", "quote", "syn 2.0.117", @@ -13737,7 +13779,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -14733,7 +14775,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix 1.0.7", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -16414,14 +16456,14 @@ version = "0.26.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" dependencies = [ - "webpki-roots 1.0.1", + "webpki-roots 1.0.9", ] [[package]] name = "webpki-roots" -version = "1.0.1" +version = "1.0.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8782dd5a41a24eed3a4f40b606249b3e236ca61adf1f25ea4d45c73de122b502" +checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a" dependencies = [ "rustls-pki-types", ] @@ -16481,7 +16523,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index c5e98d789a6..90dd1890222 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -83,7 +83,7 @@ log-query.workspace = true loki-proto.workspace = true metric-engine.workspace = true mime_guess = "2.0" -mysql_common = "0.34" +mysql_common = "0.38" notify.workspace = true object-pool = "0.5" once_cell.workspace = true @@ -157,12 +157,13 @@ common-test-util.workspace = true criterion = "0.5" json5 = "0.4" moka.workspace = true -mysql_async = { version = "0.35", default-features = false, features = [ +mysql_async = { version = "0.37", default-features = false, features = [ "default-rustls", ] } permutation = "0.4" rand.workspace = true serde_json.workspace = true +servers = { workspace = true, features = ["testing"] } session = { workspace = true, features = ["testing"] } table.workspace = true tempfile = "3.0.0" diff --git a/src/servers/tests/mysql/mysql_server_test.rs b/src/servers/tests/mysql/mysql_server_test.rs index e71369df09b..9ed5ac041e8 100644 --- a/src/servers/tests/mysql/mysql_server_test.rs +++ b/src/servers/tests/mysql/mysql_server_test.rs @@ -961,7 +961,7 @@ async fn test_query_prepared() -> Result<()> { .await { Err(mysql_async::Error::Server(e)) => assert_eq!( - "ERROR HY000 (1210): (InvalidArguments): Invalid prepare statement: Invalid SQL syntax: sql parser error: INTERVAL requires a unit after the literal value", + "ERROR 1210 (HY000): (InvalidArguments): Invalid prepare statement: Invalid SQL syntax: sql parser error: INTERVAL requires a unit after the literal value", e.to_string() ), _ => unreachable!(), diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 1456693704a..5376d7520ab 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -105,7 +105,7 @@ hex.workspace = true http.workspace = true itertools.workspace = true jsonb.workspace = true -mysql_async = { version = "0.35", default-features = false, features = [ +mysql_async = { version = "0.37", default-features = false, features = [ "time", "default-rustls-ring", ] } diff --git a/tests/runner/Cargo.toml b/tests/runner/Cargo.toml index 9801f6bc0bb..8d07c2ad1b3 100644 --- a/tests/runner/Cargo.toml +++ b/tests/runner/Cargo.toml @@ -19,7 +19,7 @@ datatypes = { workspace = true } flate2 = "1.0" hex = "0.4" local-ip-address = "0.6" -mysql = { version = "26", default-features = false, features = ["minimal", "rustls-tls"] } +mysql = { version = "28", default-features = false, features = ["minimal", "rustls-tls"] } num_cpus = "1.16" regex.workspace = true reqwest = { version = "0.12", default-features = false, features = ["rustls-tls", "stream"] } From 7539e60139657203ab0ae48fa92cff1b9ebb171f Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Thu, 13 Aug 2026 12:10:49 +0000 Subject: [PATCH 42/92] refactor: remove trivial tests (#8877) * refactor: remove trivial tests Signed-off-by: Lei, HUANG * fix: remove unused trace test import Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/auth/src/user_info.rs | 20 ------------------- .../src/instance/otlp/trace_ingest/tests.rs | 9 --------- src/meta-srv/src/gc/dropped.rs | 16 --------------- src/mito2/src/manifest/action.rs | 10 ---------- src/puffin/src/partial_reader/async.rs | 7 ------- .../src/tests/reconcile_table.rs | 3 --- 6 files changed, 65 deletions(-) diff --git a/src/auth/src/user_info.rs b/src/auth/src/user_info.rs index ab7c7f32f46..908f7bbb71a 100644 --- a/src/auth/src/user_info.rs +++ b/src/auth/src/user_info.rs @@ -223,26 +223,6 @@ mod tests { } } - #[test] - fn test_default_user_info_with_name() { - let user_info = DefaultUserInfo::with_name("test_user"); - assert_eq!(user_info.username(), "test_user"); - } - - #[test] - fn test_default_user_info_with_name_and_permission() { - let user_info = - DefaultUserInfo::with_name_and_permission("test_user", PermissionMode::ReadOnly); - assert_eq!(user_info.username(), "test_user"); - - // Cast to DefaultUserInfo to access permission_mode - let default_user = user_info - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(default_user.permission_mode, PermissionMode::ReadOnly); - } - #[test] fn test_user_info_as_any() { let user_info = DefaultUserInfo::with_name("test_user"); diff --git a/src/frontend/src/instance/otlp/trace_ingest/tests.rs b/src/frontend/src/instance/otlp/trace_ingest/tests.rs index 0f18a7194be..b3a0570db51 100644 --- a/src/frontend/src/instance/otlp/trace_ingest/tests.rs +++ b/src/frontend/src/instance/otlp/trace_ingest/tests.rs @@ -32,7 +32,6 @@ use servers::otlp::trace::SERVICE_NAME_COLUMN; use servers::otlp::trace::attributes::Attributes; use servers::otlp::trace::span::{SpanEvents, SpanLinks, TraceSpan}; use servers::otlp::trace::v1::{TraceBinaryType, TraceRetryColumn}; -use servers::query_handler::TraceIngestOutcome; use super::{ ChunkFailureReaction, Instance, TraceChunkRetry, TraceChunkSchemaState, TraceRequestSchema, @@ -101,14 +100,6 @@ fn test_classify_trace_prewrite_failure() { ); } -#[test] -fn test_add_trace_write_cost() { - let mut outcome = TraceIngestOutcome::default(); - Instance::add_trace_write_cost(&mut outcome, 3); - Instance::add_trace_write_cost(&mut outcome, 5); - assert_eq!(outcome.write_cost, 8); -} - #[test] fn test_finish_trace_failure_message() { let message = Instance::finish_trace_failure_message( diff --git a/src/meta-srv/src/gc/dropped.rs b/src/meta-srv/src/gc/dropped.rs index b4793ee5561..d0c7b1d2d64 100644 --- a/src/meta-srv/src/gc/dropped.rs +++ b/src/meta-srv/src/gc/dropped.rs @@ -284,19 +284,3 @@ impl<'a> DroppedRegionCollector<'a> { Ok(assignment) } } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_dropped_region_info() { - let info = DroppedRegionInfo { - region_id: RegionId::new(1, 1), - table_id: 1, - dst_regions: HashSet::new(), - }; - assert_eq!(info.region_id, RegionId::new(1, 1)); - assert_eq!(info.table_id, 1); - } -} diff --git a/src/mito2/src/manifest/action.rs b/src/mito2/src/manifest/action.rs index 2cde150782b..2d0c17137a8 100644 --- a/src/mito2/src/manifest/action.rs +++ b/src/mito2/src/manifest/action.rs @@ -679,16 +679,6 @@ mod tests { .unwrap(); } - #[test] - fn test_region_manifest_builder() { - // TODO(ruihang): port this test case - } - - #[test] - fn test_encode_decode_region_checkpoint() { - // TODO(ruihang): port this test case - } - #[test] fn test_region_manifest_compatibility() { // Test deserializing RegionManifest from old schema where FileId is a UUID string diff --git a/src/puffin/src/partial_reader/async.rs b/src/puffin/src/partial_reader/async.rs index 0b9a5c54c71..0525628ea34 100644 --- a/src/puffin/src/partial_reader/async.rs +++ b/src/puffin/src/partial_reader/async.rs @@ -110,11 +110,4 @@ mod tests { // seeking past the portion returns an error assert!(reader.read(31..32).await.is_err()); } - - #[tokio::test] - async fn is_eof_returns_true_at_end_of_portion() { - let data: Vec = (0..100).collect(); - let reader = PartialReader::new(data, 10, 30); - let _ = reader.read(0..20).await.unwrap(); - } } diff --git a/tests-integration/src/tests/reconcile_table.rs b/tests-integration/src/tests/reconcile_table.rs index c2feecf30ed..d31845be461 100644 --- a/tests-integration/src/tests/reconcile_table.rs +++ b/tests-integration/src/tests/reconcile_table.rs @@ -541,9 +541,6 @@ async fn test_renamed_table() { check_output_stream(output.data, expected).await; } -#[tokio::test] -async fn test_rename_table() {} - async fn unwrap_err(output: OutputData) -> String { let error = match output { OutputData::Stream(stream) => collect_batches(stream).await.unwrap_err(), From 3c80df043a46f0511bd81a64a072ca18436efac3 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Thu, 13 Aug 2026 14:17:09 +0000 Subject: [PATCH 43/92] fix(mito2): keep deletion markers when compacting part of a window (#8872) `TwcsPicker::find_inputs` decides `filter_deleted` from the shape of the whole time window, but the compaction inputs are only a subset of it: `reduce_runs` and `merge_seq_files` narrow the selection down and the max input file num limit narrows it further. When a deletion marker lands in the compacted set while the file holding the row it masks stays behind, the marker is dropped from the output and the old row becomes visible again. Re-check the final selection against the rest of the window and stop filtering deleted rows whenever something left behind still overlaps the inputs. The check compares ranges inclusively, so it also covers files that share only a boundary timestamp: run detection treats those as non-overlapping, which is how a single timestamp delete file ends up in the same run as the file it deletes rows from. Signed-off-by: jeremyhi --- src/mito2/src/compaction/twcs.rs | 117 ++++++++++++++++++- src/mito2/src/engine/compaction_test.rs | 149 ++++++++++++++++++++++++ 2 files changed, 265 insertions(+), 1 deletion(-) diff --git a/src/mito2/src/compaction/twcs.rs b/src/mito2/src/compaction/twcs.rs index bbfc1171987..342426bbefe 100644 --- a/src/mito2/src/compaction/twcs.rs +++ b/src/mito2/src/compaction/twcs.rs @@ -169,7 +169,7 @@ impl TwcsPicker { let found_runs = sorted_runs.len(); // We only remove deletion markers if we found less than 2 runs and not in append mode. // because after compaction there will be no overlapping files. - let filter_deleted = + let mut filter_deleted = found_runs <= 2 && !self.append_mode && !window_has_overlap(files, windows); if found_runs == 0 { return (vec![], filter_deleted); @@ -210,6 +210,14 @@ impl TwcsPicker { ); } + // The inputs are only a subset of the window: `reduce_runs` and `merge_seq_files` both + // narrow the selection down and the file num limit above narrows it further. A deletion + // marker may only be dropped when every file that can hold a row it masks is compacted + // along with it, so re-check the final selection against what stays in the window. + if filter_deleted && overlaps_files_left_behind(&inputs, &files_to_merge) { + filter_deleted = false; + } + if inputs.len() > 1 { // If we have more than one group to compact. log_pick_result( @@ -227,6 +235,20 @@ impl TwcsPicker { } } +/// Returns whether any file group of `window_files` that is not part of `inputs` overlaps +/// `inputs`. Such a group keeps rows a deletion marker in `inputs` masks, so the compaction +/// must not filter deleted rows out. +/// +/// Overlapping is inclusive on both boundaries here: two files that only share a boundary +/// timestamp can still hold the same row. +fn overlaps_files_left_behind(inputs: &[FileGroup], window_files: &[FileGroup]) -> bool { + let picked: HashSet<_> = inputs.iter().flat_map(|fg| fg.file_ids()).collect(); + window_files + .iter() + .filter(|fg| !fg.file_ids().iter().any(|id| picked.contains(id))) + .any(|fg| inputs.iter().any(|input| input.overlap_inclusive(fg))) +} + #[allow(clippy::too_many_arguments)] fn log_pick_result( region_id: RegionId, @@ -1456,6 +1478,99 @@ mod tests { assert_eq!(output[0].inputs.len(), 32); } + #[tokio::test] + async fn test_limit_max_input_files_keeps_deletion_markers() { + common_telemetry::init_default_ut_logging(); + + // One large file group spanning the whole window plus 32 small ones nested inside + // it. That is exactly 2 runs, so the window on its own allows filtering deletions. + let mut files = vec![new_file_handle_with_size_and_sequence( + FileId::random(), + 0, + 3_000_000, + 0, + 1, + 1024 * 1024 * 1024, + )]; + files.extend((0..32).map(|idx: i64| { + new_file_handle_with_size_and_sequence( + FileId::random(), + (idx + 1) * 10_000, + (idx + 1) * 10_000 + 1_000, + 0, + (idx + 2) as u64, + 1024, + ) + })); + + let windows = assign_to_windows(files.iter(), 3600); + + let picker = TwcsPicker { + trigger_file_num: 4, + time_window_seconds: Some(3600), + max_output_file_size: None, + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let active_window = find_latest_window_in_seconds(files.iter(), 3600); + let output = picker + .build_output_with_time_range(RegionId::from_u64(123), windows, active_window, None) + .await + .unwrap(); + + assert_eq!(1, output.len()); + // The input file num limit picks the smallest groups first, so the large group is + // left behind while the small groups that overlap it are compacted. + assert_eq!(32, output[0].inputs.len()); + assert!( + !output[0].filter_deleted, + "deletion markers must be kept once the file num limit drops files they may mask" + ); + } + + #[tokio::test] + async fn test_limit_max_input_files_still_filters_without_overlap() { + common_telemetry::init_default_ut_logging(); + + // 40 file groups with disjoint time ranges, i.e. a single run. Nothing the file num + // limit leaves behind can hold a row masked by a deletion marker we compact. + let files: Vec<_> = (0..40i64) + .map(|idx| { + new_file_handle_with_size_and_sequence( + FileId::random(), + (idx + 1) * 10_000, + (idx + 1) * 10_000 + 1_000, + 0, + (idx + 1) as u64, + 1024, + ) + }) + .collect(); + + let windows = assign_to_windows(files.iter(), 3600); + + let picker = TwcsPicker { + trigger_file_num: 4, + time_window_seconds: Some(3600), + max_output_file_size: Some(1024 * 1024 * 1024), + append_mode: false, + max_background_tasks: None, + time_range: None, + }; + + let active_window = find_latest_window_in_seconds(files.iter(), 3600); + let output = picker + .build_output_with_time_range(RegionId::from_u64(123), windows, active_window, None) + .await + .unwrap(); + + assert_eq!(1, output.len()); + assert_eq!(32, output[0].inputs.len()); + assert!(output[0].filter_deleted); + } + #[tokio::test] async fn test_newer_windows_have_priority() { let older_file_ids = [FileId::random(), FileId::random()]; diff --git a/src/mito2/src/engine/compaction_test.rs b/src/mito2/src/engine/compaction_test.rs index 4404b8d46a1..1fc3a8f785d 100644 --- a/src/mito2/src/engine/compaction_test.rs +++ b/src/mito2/src/engine/compaction_test.rs @@ -1342,6 +1342,155 @@ async fn test_compaction_region_with_overlapping_delete_all_with_format(flat_for assert!(vec.is_empty()); } +#[tokio::test] +async fn test_compaction_input_limit_keeps_rows_deleted() { + test_compaction_input_limit_keeps_rows_deleted_with_format(false).await; + test_compaction_input_limit_keeps_rows_deleted_with_format(true).await; +} + +/// Creates a region that only compacts when asked, so that a test can build an exact file +/// layout with `put_and_flush` and `delete_and_flush`. +async fn env_for_manual_compaction( + env: &mut TestEnv, + region_id: RegionId, + flat_format: bool, +) -> (MitoEngine, Vec) { + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + min_compaction_interval: Duration::from_secs(3600), + ..Default::default() + }) + .await; + + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + let request = CreateRequestBuilder::new() + .insert_option("compaction.type", "twcs") + .insert_option("compaction.twcs.time_window", "1h") + .build(); + let column_schemas = request + .column_metadatas + .iter() + .map(column_metadata_to_column_schema) + .collect::>(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + (engine, column_schemas) +} + +/// The picker caps a compaction at 32 input files and drops the largest file groups to get +/// there. A deletion marker among the picked files must not be filtered out while the file +/// holding the rows it masks stays behind, otherwise those rows become visible again. +async fn test_compaction_input_limit_keeps_rows_deleted_with_format(flat_format: bool) { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::new().await; + let region_id = RegionId::new(1, 1); + let (engine, column_schemas) = + env_for_manual_compaction(&mut env, region_id, flat_format).await; + + // One large file spanning the whole time window. + put_and_flush(&engine, region_id, &column_schemas, 0..3000).await; + // Deletes 6 rows of that file. The markers land in a tiny file that overlaps it. + delete_and_flush(&engine, region_id, &column_schemas, 10..16).await; + // 31 more tiny files that overlap the large one but not each other, so the window holds + // 33 file groups forming 2 runs. + for i in 2..33 { + put_and_flush(&engine, region_id, &column_schemas, i * 10..i * 10 + 6).await; + } + + let scanner = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap(); + assert_eq!( + 33, + scanner.num_files(), + "unexpected files: {:?}", + scanner.file_ids() + ); + + compact(&engine, region_id).await; + + let scanner = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap(); + // The 32 tiny files are merged into one; the large file exceeds the input file num limit + // and is left behind. + assert_eq!( + 2, + scanner.num_files(), + "unexpected files: {:?}", + scanner.file_ids() + ); + let stream = scanner.scan().await.unwrap(); + let vec = collect_stream_ts(stream).await; + assert!( + !(10..16).any(|ts| vec.contains(&(ts * 1000))), + "deleted rows are visible again after compaction" + ); + assert_eq!(2994, vec.len()); +} + +#[tokio::test] +async fn test_compaction_of_part_of_a_run_keeps_rows_deleted() { + test_compaction_of_part_of_a_run_keeps_rows_deleted_with_format(false).await; + test_compaction_of_part_of_a_run_keeps_rows_deleted_with_format(true).await; +} + +/// A run is supposed to hold no overlapping files, which is what lets `merge_seq_files` +/// compact part of a run and still filter deleted rows. Run detection compares time ranges +/// exclusively though, so a file covering a single timestamp lands in the same run as the +/// file it deletes rows from. The rows must stay deleted when only one of the two is picked. +async fn test_compaction_of_part_of_a_run_keeps_rows_deleted_with_format(flat_format: bool) { + common_telemetry::init_default_ut_logging(); + let mut env = TestEnv::new().await; + let region_id = RegionId::new(1, 1); + let (engine, column_schemas) = + env_for_manual_compaction(&mut env, region_id, flat_format).await; + + // One large file spanning the whole time window. + put_and_flush(&engine, region_id, &column_schemas, 0..3000).await; + // Deletes a single row of that file, so the marker lands in a file covering one timestamp. + delete_and_flush(&engine, region_id, &column_schemas, 1..2).await; + // 31 more single row files, each holding another key at its own timestamp. + for ts in 1..32 { + let rows = Rows { + schema: column_schemas.clone(), + rows: build_rows_for_key("b", ts * 10, ts * 10 + 1, 0), + }; + put_rows(&engine, region_id, rows).await; + flush(&engine, region_id).await; + } + + compact(&engine, region_id).await; + + let scanner = engine + .scanner(region_id, ScanRequest::default()) + .await + .unwrap(); + let stream = scanner.scan().await.unwrap(); + let vec = collect_stream_ts(stream).await; + assert!( + !vec.contains(&1000), + "deleted row is visible again after compaction" + ); + assert_eq!(3030, vec.len()); +} + // For issue https://github.com/GreptimeTeam/greptimedb/issues/3633 #[tokio::test] async fn test_readonly_during_compaction() { From f692ac32f3f92584ee81abe04cd7fc0e669f3a3d Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 14 Aug 2026 04:21:46 +0000 Subject: [PATCH 44/92] test(mito2): isolate sequence publication barrier (#8876) * test(mito2): isolate sequence publication barrier Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(mito2): trim sequence barrier test noise Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/mito2/src/region_write_ctx.rs | 84 +++++++++++----------------- src/mito2/src/worker/handle_write.rs | 47 ++-------------- 2 files changed, 39 insertions(+), 92 deletions(-) diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index f64f9810f03..400f7db1fc4 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -199,7 +199,6 @@ impl RegionWriteCtx { &self.version } - /// Returns the version control of the region. #[cfg(test)] pub(crate) fn version_control(&self) -> &VersionControlRef { &self.version_control @@ -327,7 +326,7 @@ impl RegionWriteCtx { return; } #[cfg(test)] - test_hooks::pause_before_bulk_install(self.region_id).await; + test_hooks::pause_before_bulk_install(self.region_id, &self.version_control).await; let _timer = metrics::REGION_WORKER_HANDLE_WRITE_ELAPSED .with_label_values(&["write_bulk"]) .start_timer(); @@ -375,20 +374,10 @@ impl RegionWriteCtx { } } - /// Publishes the sequences and entry id assigned by this context to the - /// region's committed watermark. - /// - /// Must be called after both [`write_memtable`](Self::write_memtable) and - /// [`write_bulk`](Self::write_bulk) have completed, so the committed - /// sequence never covers rows that are not yet physically installed in the - /// memtable (bulk part sequences are assigned in [`push_bulk`](Self::push_bulk) - /// before installation). Since we store the last sequence and entry id in - /// the region, we decrease `next_sequence` and `next_entry_id` by 1. - /// - /// If the write operation failed (e.g. the WAL entry could not be built), - /// no rows were installed and nothing must be published: advancing the - /// committed watermark here would make it cover rows that were never - /// written. + /// Publishes the assigned sequences and entry id to the region's committed + /// watermark. Only call after both [`write_memtable`](Self::write_memtable) + /// and [`write_bulk`](Self::write_bulk) have completed; a failed context + /// must not publish. pub(crate) fn publish_sequence_and_entry_id(&self) { if self.failed { return; @@ -407,25 +396,20 @@ pub(crate) mod test_hooks { use store_api::storage::RegionId; use tokio::sync::watch; - /// Channels of an armed bulk-install barrier. - /// - /// `reached` signals that a bulk write paused between ordinary-memtable - /// handling and bulk installation; `release` unblocks it. Dropping the - /// senders (by disarming the barrier) also unblocks paused writes because - /// their `wait_for` fails on a closed channel. + use crate::region::version::VersionControlRef; + + /// Channels of an armed bulk-install barrier; dropping the senders (by + /// disarming) unblocks writes paused on it. struct ActiveBarrier { id: u64, - /// Only bulk writes for this region pause at the barrier; writes for - /// other regions pass through untouched, so concurrently running tests - /// can never satisfy (or be blocked by) each other's barrier. + /// Only bulk writes for this region on this version control (Arc + /// identity) pause at the barrier. target_region_id: RegionId, + target_version_control: VersionControlRef, reached: watch::Sender, release: watch::Sender, } - /// The currently armed barrier, if any. Wrapped in a `Mutex` so a test can - /// arm a fresh barrier after a previous one was released (or dropped), and - /// so a barrier can never leak past the test that owns it. static ACTIVE_BARRIER: Mutex> = Mutex::new(None); static NEXT_BARRIER_ID: AtomicU64 = AtomicU64::new(1); @@ -437,12 +421,8 @@ pub(crate) mod test_hooks { .unwrap_or_else(|poisoned| poisoned.into_inner()) } - /// RAII guard for the bulk-install barrier. - /// - /// The guard owns the armed barrier. Releasing it — or dropping the guard, - /// e.g. when a test panics — unblocks any write paused at the barrier and - /// disarms it, so a paused write can never hang and later tests can arm a - /// fresh barrier. + /// RAII guard: releasing (or dropping) it unblocks a paused write and + /// disarms the barrier. pub(crate) struct BulkInstallBarrier { id: u64, reached_rx: watch::Receiver, @@ -451,20 +431,17 @@ pub(crate) mod test_hooks { } impl BulkInstallBarrier { - /// Waits until a bulk write paused at the barrier. pub(crate) async fn wait_until_reached(&mut self) { if !*self.reached_rx.borrow() { let _ = self.reached_rx.wait_for(|reached| *reached).await; } } - /// Unblocks the paused write and disarms the barrier. pub(crate) fn release(&mut self) { if self.released { return; } self.released = true; - // Flip the release value so writes paused on this barrier proceed. let _ = self.release_tx.send(true); disarm_barrier(self.id); } @@ -476,12 +453,12 @@ pub(crate) mod test_hooks { } } - /// Arms the bulk-install barrier for `target_region_id` and returns a - /// guard that owns it. - /// - /// Any previously armed barrier is replaced: its senders are dropped, which - /// unblocks writes that were paused on it instead of letting them hang. - pub(crate) fn arm_bulk_install_barrier(target_region_id: RegionId) -> BulkInstallBarrier { + /// Arms the bulk-install barrier and returns the owning guard; any + /// previously armed barrier is replaced. + pub(crate) fn arm_bulk_install_barrier( + target_region_id: RegionId, + target_version_control: VersionControlRef, + ) -> BulkInstallBarrier { let (reached_tx, reached_rx) = watch::channel(false); let (release_tx, _release_rx) = watch::channel(false); let id = NEXT_BARRIER_ID.fetch_add(1, Ordering::Relaxed); @@ -489,6 +466,7 @@ pub(crate) mod test_hooks { *active = Some(ActiveBarrier { id, target_region_id, + target_version_control, reached: reached_tx, release: release_tx.clone(), }); @@ -500,7 +478,6 @@ pub(crate) mod test_hooks { } } - /// Removes the barrier with `id` from the statics if it is still active. fn disarm_barrier(id: u64) { let mut active = lock_active_barrier(); if active.as_ref().is_some_and(|barrier| barrier.id == id) { @@ -508,20 +485,27 @@ pub(crate) mod test_hooks { } } - /// Pauses a bulk write for `region_id` before installing its parts until - /// the test releases the barrier (or it is disarmed). Bulk writes for - /// other regions return immediately. - pub(crate) async fn pause_before_bulk_install(region_id: RegionId) { + /// Pauses a bulk write before installing its parts until the barrier is + /// released or disarmed. + pub(crate) async fn pause_before_bulk_install( + region_id: RegionId, + version_control: &VersionControlRef, + ) { let (reached_tx, release_rx) = { let active = lock_active_barrier(); match active.as_ref() { - Some(barrier) if barrier.target_region_id == region_id => { + Some(barrier) + if barrier.target_region_id == region_id + && std::sync::Arc::ptr_eq( + &barrier.target_version_control, + version_control, + ) => + { (barrier.reached.clone(), barrier.release.subscribe()) } _ => return, } }; - // Signal that a bulk write reached the pause point. let _ = reached_tx.send(true); let mut release_rx = release_rx; if !*release_rx.borrow() { diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 0c2eecdea20..96940756248 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -128,9 +128,6 @@ impl RegionWorkerLoop { let mut region_ctx = region_ctxs.into_values().next().unwrap(); region_ctx.write_memtable().await; region_ctx.write_bulk().await; - // Publish only after all rows (including bulk parts) are - // physically installed, so a scan opening on the committed - // sequence can never bind H to invisible rows. region_ctx.publish_sequence_and_entry_id(); put_rows += region_ctx.put_num; delete_rows += region_ctx.delete_num; @@ -142,8 +139,6 @@ impl RegionWorkerLoop { common_runtime::spawn_global(async move { region_ctx.write_memtable().await; region_ctx.write_bulk().await; - // The spawned task owns the moved ctx, so publish - // inside the task after the memtable writes. region_ctx.publish_sequence_and_entry_id(); (region_ctx.put_num, region_ctx.delete_num) }) @@ -763,10 +758,6 @@ mod tests { use crate::test_util::ts_ms_value; use crate::test_util::version_util::VersionControlBuilder; - /// Creates a bulk part with `num_rows` rows. The schema carries the - /// builder metadata's columns (`tag_0` primary key + `ts` time index) so - /// the bulk-install conversion succeeds; `timestamp_index` points at the - /// `ts` column. fn new_bulk_part(num_rows: i64) -> BulkPart { let schema = Arc::new(Schema::new(vec![ Field::new("tag_0", DataType::Utf8, true), @@ -887,7 +878,6 @@ mod tests { } } - /// Creates a write context for `region_id` with one pending mutation of one row. fn new_region_ctx( region_id: RegionId, ) -> (RegionWriteCtx, oneshot::Receiver>) { @@ -940,16 +930,13 @@ mod tests { region_ctxs.insert(ok_region, ctx); let entry_id = region_ctxs[&ok_region].next_entry_id(); - // The failed region must not fail the batch or panic the worker. + // The failed region must not fail the batch. assert!(write_wal(&wal, &mut region_ctxs).await); assert!(region_ctxs[&failing_region].is_failed()); assert!(!region_ctxs[&ok_region].is_failed()); assert_eq!(entry_id + 1, region_ctxs[&ok_region].next_entry_id()); - // Run the memtable and publication phase the worker runs after `write_wal`. - // The failed region installed no rows, so its committed sequence must not - // advance; the successful sibling region advances normally. for region_ctx in region_ctxs.values_mut() { region_ctx.write_memtable().await; region_ctx.write_bulk().await; @@ -969,30 +956,15 @@ mod tests { .committed_sequence() ); - // Waiters of the failed region get the error while others get the result. drop(region_ctxs); assert!(failing_rx.await.unwrap().is_err()); assert_eq!(1, ok_rx.await.unwrap().unwrap()); } - // The committed sequence must not be published before the bulk part's rows - // are physically installed in the memtable: publication happens strictly - // after `write_bulk` returns. Armed with the bulk-install test barrier, - // this deterministically pauses the worker's memtable phase between the - // ordinary-memtable handling and the bulk installation and asserts the - // region's committed sequence is still its initial value (0) — read via - // the `version_control()` accessor, never through a scanner API. After the - // barrier is released, the committed sequence must advance to cover the - // bulk rows. #[tokio::test] async fn test_bulk_write_sequence_not_committed_before_install_worker_level() { let region_id = RegionId::new(1, 1); let version_control = Arc::new(VersionControlBuilder::new().build()); - assert_eq!( - 0, - version_control.committed_sequence(), - "the builder must start at sequence 0" - ); let mut region_ctxs = HashMap::new(); let mut ctx = RegionWriteCtx::new( @@ -1002,21 +974,18 @@ mod tests { None, ); let (tx, rx) = oneshot::channel(); - // Push a 3-row bulk part without an explicit sequence, exactly like the - // worker's `process_bulk_requests`. assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(3), None)); region_ctxs.insert(region_id, ctx); - // Run the WAL phase the worker runs before the memtable phase. let wal = Wal::new(Arc::new(MockLogStore::default())); assert!(write_wal(&wal, &mut region_ctxs).await); assert!(!region_ctxs[®ion_id].is_failed()); - // Arm the bulk-install barrier: `write_bulk` pauses right before the - // parts are physically installed into the memtable. - let mut barrier = crate::region_write_ctx::test_hooks::arm_bulk_install_barrier(region_id); + let mut barrier = crate::region_write_ctx::test_hooks::arm_bulk_install_barrier( + region_id, + version_control.clone(), + ); - // Run the worker's memtable and publication phases in the background. let write_handle = tokio::spawn(async move { let mut region_ctx = region_ctxs.remove(®ion_id).unwrap(); region_ctx.write_memtable().await; @@ -1024,7 +993,6 @@ mod tests { region_ctx.publish_sequence_and_entry_id(); }); - // Wait until the write paused at the barrier: deterministic, no sleeps. tokio::time::timeout( std::time::Duration::from_secs(10), barrier.wait_until_reached(), @@ -1032,16 +1000,12 @@ mod tests { .await .expect("bulk write never reached the install barrier"); - // The committed sequence must not have advanced yet: publication - // happens strictly after the bulk part is installed. assert_eq!( 0, version_control.committed_sequence(), "committed sequence leaked before the bulk part was installed" ); - // Release the barrier: the bulk part installs and the committed - // sequence advances to cover all 3 bulk rows. barrier.release(); write_handle.await.expect("bulk write should complete"); assert_eq!( @@ -1050,7 +1014,6 @@ mod tests { "committed sequence must cover the installed bulk rows" ); - // The bulk waiter is notified with the number of installed rows. assert_eq!(3, rx.await.unwrap().unwrap()); } From b1263fc65ffb95a80353177f9fef01f551879f68 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Fri, 14 Aug 2026 06:15:42 +0000 Subject: [PATCH 45/92] perf: optimize Prometheus remote write v2 decoding (#8873) * test: add bench for prom decode Signed-off-by: shuiyisong * chore: merge v2 decode manually and add to bench Signed-off-by: shuiyisong * refactor: update v2 path Signed-off-by: shuiyisong * refactor: update v2 path Signed-off-by: shuiyisong * fix: use constant Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/servers/Cargo.toml | 5 + src/servers/benches/prom_decode.rs | 247 ++++- src/servers/src/http/prom_store.rs | 33 +- src/servers/src/prom_remote_write/README.md | 33 +- src/servers/src/prom_remote_write/v2.rs | 998 +++++++++++++++--- src/servers/src/row_writer.rs | 45 +- .../tests/prom_remote_write_v2_test.rs | 2 +- 7 files changed, 1147 insertions(+), 216 deletions(-) diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 90dd1890222..d64a1bd5a5d 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -180,6 +180,11 @@ common-version.workspace = true name = "bench_prom" harness = false +[[bench]] +name = "prom_decode" +harness = false +required-features = ["testing"] + [[bench]] name = "to_http_output" harness = false diff --git a/src/servers/benches/prom_decode.rs b/src/servers/benches/prom_decode.rs index e9db0ad7716..06b0d02e940 100644 --- a/src/servers/benches/prom_decode.rs +++ b/src/servers/benches/prom_decode.rs @@ -12,21 +12,137 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; use std::time::Duration; -use api::prom_store::remote::WriteRequest; -use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main}; +use api::greptime_proto::io::prometheus::write::v2 as write_v2; +use api::prom_store::remote::{self as write_v1, WriteRequest}; +use criterion::{ + BatchSize, BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main, +}; use prost::Message; use servers::prom_remote_write::decode::{PromSeriesProcessor, PromWriteRequest}; +use servers::prom_remote_write::v2::test_util as remote_write_v2; use servers::prom_remote_write::validation::{PromValidationMode, validate_label_name}; use servers::prom_store::to_grpc_row_insert_requests; -fn bench_decode_prom_request(c: &mut Criterion) { +fn load_fixture_v1_bytes() -> Vec { let mut d = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")); d.push("benches"); d.push("write_request.pb.data"); + std::fs::read(d).expect("read write_request.pb.data fixture") +} - let data = std::fs::read(d).unwrap(); +/// Convert a PRW v1 `WriteRequest` into an equivalent PRW v2 `Request`. +/// +/// Labels are interned into the v2 symbol table so payloads with repeated label +/// names/values exercise the main on-wire advantage of v2. +fn write_request_v1_to_v2(v1: &WriteRequest) -> write_v2::Request { + let mut symbols = vec![String::new()]; + let mut symbol_index: HashMap = HashMap::new(); + symbol_index.insert(String::new(), 0); + + let mut intern = |value: &str| -> u32 { + if let Some(&idx) = symbol_index.get(value) { + return idx; + } + let idx = symbols.len() as u32; + symbols.push(value.to_string()); + symbol_index.insert(value.to_string(), idx); + idx + }; + + let timeseries = v1 + .timeseries + .iter() + .map(|series| { + let mut labels: Vec<&write_v1::Label> = series.labels.iter().collect(); + // PRW v2 requires lexicographically sorted label pairs. + labels.sort_by(|a, b| a.name.cmp(&b.name).then_with(|| a.value.cmp(&b.value))); + + let mut labels_refs = Vec::with_capacity(labels.len() * 2); + for label in labels { + labels_refs.push(intern(&label.name)); + labels_refs.push(intern(&label.value)); + } + + write_v2::TimeSeries { + labels_refs, + samples: series + .samples + .iter() + .map(|sample| write_v2::Sample { + value: sample.value, + timestamp: sample.timestamp, + start_timestamp: 0, + }) + .collect(), + histograms: Vec::new(), + exemplars: Vec::new(), + metadata: None, + } + }) + .collect(); + + write_v2::Request { + symbols, + timeseries, + } +} + +/// Synthetic sample workload with controllable label reuse. +/// +/// `shared_labels` keeps common label names/values across series (v2-friendly). +/// Each series also gets a unique `series_id` label so cardinality stays high. +fn generate_sample_workload( + series_count: usize, + samples_per_series: usize, + shared_label_count: usize, +) -> (WriteRequest, write_v2::Request) { + let shared_labels: Vec = (0..shared_label_count) + .map(|i| write_v1::Label { + name: format!("label_{i}"), + value: format!("value_{}", i % 16), + }) + .collect(); + + let timeseries = (0..series_count) + .map(|series_idx| { + let mut labels = Vec::with_capacity(shared_label_count + 2); + labels.push(write_v1::Label { + name: "__name__".to_string(), + value: format!("metric_{}", series_idx % 32), + }); + labels.extend(shared_labels.iter().cloned()); + labels.push(write_v1::Label { + name: "series_id".to_string(), + value: series_idx.to_string(), + }); + + write_v1::TimeSeries { + labels, + samples: (0..samples_per_series) + .map(|sample_idx| write_v1::Sample { + value: (series_idx * 1000 + sample_idx) as f64, + timestamp: 1_700_000_000_000 + sample_idx as i64 * 1_000, + }) + .collect(), + exemplars: Vec::new(), + histograms: Vec::new(), + } + }) + .collect(); + + let v1 = WriteRequest { + timeseries, + metadata: Vec::new(), + }; + let v2 = write_request_v1_to_v2(&v1); + (v1, v2) +} + +fn bench_decode_prom_request(c: &mut Criterion) { + let data = load_fixture_v1_bytes(); let mut group = c.benchmark_group("decode_prom_request"); group.measurement_time(Duration::from_secs(3)); @@ -64,6 +180,128 @@ fn bench_decode_prom_request(c: &mut Criterion) { group.finish(); } +/// Compare Prometheus remote-write v1 vs v2 decoding efficiency. +/// +/// Two layers are measured on equivalent logical payloads: +/// 1. protobuf decode only (`WriteRequest` / v2 `Request`) +/// 2. full Greptime path to row-insert requests +/// +/// Workloads: +/// - `fixture`: existing `write_request.pb.data` converted v1 -> v2 +/// - `synthetic_*`: generated series with shared labels (symbol-table friendly) +#[allow(clippy::print_stderr)] +fn bench_prom_v1_vs_v2_decode(c: &mut Criterion) { + let fixture_v1_bytes = load_fixture_v1_bytes(); + let fixture_v1 = WriteRequest::decode(fixture_v1_bytes.as_slice()).unwrap(); + let fixture_v2 = write_request_v1_to_v2(&fixture_v1); + let fixture_v2_bytes = fixture_v2.encode_to_vec(); + + let synthetic = [ + ( + "synthetic_1k_series_1_sample", + generate_sample_workload(1_000, 1, 8), + ), + ( + "synthetic_1k_series_10_samples", + generate_sample_workload(1_000, 10, 8), + ), + ( + "synthetic_5k_series_1_sample", + generate_sample_workload(5_000, 1, 8), + ), + ]; + + let mut workloads: Vec<(&str, Vec, Vec)> = + vec![("fixture", fixture_v1_bytes, fixture_v2_bytes)]; + for (name, (v1, v2)) in &synthetic { + workloads.push((*name, v1.encode_to_vec(), v2.encode_to_vec())); + } + + eprintln!("\nprom remote-write v1 vs v2 payload sizes:"); + for (name, v1_bytes, v2_bytes) in &workloads { + let ratio = v2_bytes.len() as f64 / v1_bytes.len() as f64; + eprintln!( + " {name}: v1={} B, v2={} B, v2/v1={ratio:.3}", + v1_bytes.len(), + v2_bytes.len() + ); + } + + // --- protobuf decode only --- + { + let mut group = c.benchmark_group("prom_rw_protobuf_decode"); + group.measurement_time(Duration::from_secs(3)); + + for (name, v1_bytes, v2_bytes) in &workloads { + group.throughput(Throughput::Bytes(v1_bytes.len() as u64)); + group.bench_with_input(BenchmarkId::new("v1", name), v1_bytes, |b, bytes| { + b.iter(|| { + black_box(WriteRequest::decode(black_box(bytes.as_slice())).unwrap()); + }); + }); + + group.throughput(Throughput::Bytes(v2_bytes.len() as u64)); + group.bench_with_input(BenchmarkId::new("v2", name), v2_bytes, |b, bytes| { + b.iter(|| { + black_box(write_v2::Request::decode(black_box(bytes.as_slice())).unwrap()); + }); + }); + } + + group.finish(); + } + + // --- full path: protobuf -> Greptime row inserts --- + { + let mut group = c.benchmark_group("prom_rw_decode_to_rows"); + group.sample_size(50); + group.measurement_time(Duration::from_secs(5)); + + for (name, v1_bytes, v2_bytes) in &workloads { + group.throughput(Throughput::Bytes(v1_bytes.len() as u64)); + group.bench_with_input(BenchmarkId::new("v1_custom", name), v1_bytes, |b, bytes| { + let mut prom_request = PromWriteRequest::default(); + let mut processor = PromSeriesProcessor::default_processor(); + b.iter_batched( + || bytes.clone(), + |bytes| { + prom_request + .decode(black_box(bytes), PromValidationMode::Strict, &mut processor) + .unwrap(); + let rows = prom_request.as_row_insert_requests(); + black_box(&rows); + }, + BatchSize::LargeInput, + ); + }); + + // Baseline: stock prost WriteRequest + existing converter. + group.throughput(Throughput::Bytes(v1_bytes.len() as u64)); + group.bench_with_input(BenchmarkId::new("v1_prost", name), v1_bytes, |b, bytes| { + b.iter(|| { + let request = WriteRequest::decode(black_box(bytes.as_slice())).unwrap(); + black_box(to_grpc_row_insert_requests(&request).unwrap()); + }); + }); + + group.throughput(Throughput::Bytes(v2_bytes.len() as u64)); + group.bench_with_input(BenchmarkId::new("v2_manual", name), v2_bytes, |b, bytes| { + b.iter(|| { + black_box( + remote_write_v2::decode_uncompressed_write_requests( + black_box(bytes.as_slice()), + true, + ) + .unwrap(), + ); + }); + }); + } + + group.finish(); + } +} + /// Benchmark comparing UTF-8 string validation (`decode_string`) vs /// direct byte-level Prometheus label name validation (`decode_label_name`). fn bench_label_name_validation(c: &mut Criterion) { @@ -155,6 +393,7 @@ fn bench_utf8_validation(c: &mut Criterion) { criterion_group!( benches, bench_decode_prom_request, + bench_prom_v1_vs_v2_decode, bench_label_name_validation, bench_utf8_validation ); diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index ac4f39628d1..ed960264b0d 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -48,7 +48,7 @@ use crate::http::header::{ use crate::pending_rows_batcher::PendingRowsBatcher; use crate::prom_remote_write::decode::PromSeriesProcessor; use crate::prom_remote_write::decode_remote_write_request; -use crate::prom_remote_write::v2::{decode_remote_write_v2_request, into_write_requests}; +use crate::prom_remote_write::v2::decode_remote_write_v2; use crate::prom_remote_write::validation::PromValidationMode; use crate::prom_store::snappy_decompress; use crate::query_handler::{PipelineHandlerRef, PromStoreProtocolHandlerRef, PromStoreResponse}; @@ -235,23 +235,11 @@ async fn remote_write_v2( let (db, mut query_ctx, _timer) = prepare_remote_write_context(¶ms, query_ctx, REMOTE_WRITE_V2_VERSION); - let request = match decode_remote_write_v2_request(is_zstd, body) { - Ok(request) => request, - Err(error) => return Ok(remote_write_v2_error_response(error, 0, 0, 0)), - }; - if !experimental_enable_prometheus_native_histogram && request_has_native_histograms(&request) { - return Ok(remote_write_v2_error_response( - error::InvalidPromRemoteRequestSnafu { - msg: "prometheus remote write v2 native histogram ingestion is experimental; set prom_store.experimental_enable_prometheus_native_histogram = true to enable it" - .to_string(), - } - .build(), - 0, - 0, - 0, - )); - } - let req = match into_write_requests(request) { + let req = match decode_remote_write_v2( + is_zstd, + body, + experimental_enable_prometheus_native_histogram, + ) { Ok(req) => req, Err(error) => return Ok(remote_write_v2_error_response(error, 0, 0, 0)), }; @@ -302,15 +290,6 @@ async fn remote_write_v2( Ok((StatusCode::NO_CONTENT, headers).into_response()) } -fn request_has_native_histograms( - request: &api::greptime_proto::io::prometheus::write::v2::Request, -) -> bool { - request - .timeseries - .iter() - .any(|series| !series.histograms.is_empty()) -} - fn vm_proto_version_response(params: &RemoteWriteQuery) -> Option { params .get_vm_proto_version diff --git a/src/servers/src/prom_remote_write/README.md b/src/servers/src/prom_remote_write/README.md index d21eead1901..9e86a8753dd 100644 --- a/src/servers/src/prom_remote_write/README.md +++ b/src/servers/src/prom_remote_write/README.md @@ -11,28 +11,33 @@ Remote write v2 enters through `remote_write_v2` in ```mermaid flowchart TD A["HTTP /v1/prometheus/write"] --> B["remote_write_v2"] - B --> C["decode_remote_write_v2_request"] - C --> D["into_write_requests"] + B --> C["decode_remote_write_v2"] + C --> D["borrow symbols and time-series bytes"] + D --> E["decode leaf messages and build rows"] - D --> E["samples ContextReq"] - D --> F["histograms ContextReq"] + E --> F["samples ContextReq"] + E --> G["histograms ContextReq"] - E --> G["write_prometheus_rows_with_progress"] - G --> H["metric engine / pending batcher when enabled"] + F --> H["write_prometheus_rows_with_progress"] + H --> I["metric engine / pending batcher when enabled"] - F --> I["histogram ContextOpt"] - I --> J["write_prometheus_rows_with_progress"] - J --> K["same metric-engine flag as samples, no batcher"] - K --> L["table: "] + G --> J["histogram ContextOpt"] + J --> K["write_prometheus_rows_with_progress"] + K --> L["same metric-engine flag as samples, no batcher"] + L --> M["table: "] - L --> M["field: configured native-histogram Struct"] - M --> N["struct children: counts, spans, buckets, sum, schema"] - H --> P["written headers and counters"] - K --> P + M --> N["field: configured native-histogram Struct"] + N --> O["struct children: counts, spans, buckets, sum, schema"] + I --> P["written headers and counters"] + L --> P ``` The conversion step splits one v2 request into two `ContextReq`s: +The `decode` codec metric covers decompression and the borrowed request +envelope. The `convert` metric covers per-series scans, leaf prost decoding, +validation, and row construction. + - samples keep the existing sample table name and can use the metric-engine physical table path and pending rows batcher; - native histograms keep the existing metric table name. diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 5a25feb566a..31da7ac1f35 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -17,19 +17,24 @@ use std::collections::hash_map::Entry; use ahash::{HashMap, HashMapExt, HashSet, HashSetExt}; use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; use api::greptime_proto::io::prometheus::write::v2::{ - BucketSpan, Histogram, Request, Sample, TimeSeries, metadata, + BucketSpan, Exemplar, Histogram, Metadata, Sample, metadata, }; #[cfg(test)] -use api::greptime_proto::io::prometheus::write::v2::{Exemplar, Metadata}; +use api::greptime_proto::io::prometheus::write::v2::{Request, TimeSeries}; use api::helper::ColumnDataTypeWrapper; use api::v1::value::ValueData; -use api::v1::{ColumnSchema, ListValue, RowInsertRequest, Rows, SemanticType, Value}; -use bytes::Bytes; +use api::v1::{ + ColumnDataType, ColumnSchema, ListValue, RowInsertRequest, Rows, SemanticType, Value, +}; +use bytes::{Buf, Bytes}; use common_grpc::precision::Precision; use common_query::native_histogram::*; use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; use pipeline::{ContextOpt, ContextReq}; -use prost::Message; +use prost::encoding::{ + DecodeContext, WireType, decode_key, decode_varint, message, skip_field, uint32, +}; +use prost::{DecodeError, Message}; use snafu::{OptionExt, ResultExt, ensure}; use table::requests::{ METADATA_QUALITY_DECLARED, SEMANTIC_METRIC_METADATA_QUALITY, SEMANTIC_METRIC_TYPE, @@ -52,23 +57,60 @@ use crate::semantic::{ openmetrics_unit_to_ucum, }; -type PromTags = Vec<(String, String)>; -type ResolvedSeriesLabels = (PromCtx, String, PromTags); +type PromTags<'a> = Vec<(&'a str, String)>; +type ResolvedSeriesLabels<'a> = (PromCtx, String, PromTags<'a>); const MAX_REMOTE_WRITE_V2_SCHEMA: i32 = 8; const MAX_REDUCIBLE_REMOTE_WRITE_V2_SCHEMA: i32 = 52; +const TIME_SERIES_LABELS_REFS_TAG: u32 = 1; +const TIME_SERIES_SAMPLES_TAG: u32 = 2; +const TIME_SERIES_HISTOGRAMS_TAG: u32 = 3; +const TIME_SERIES_EXEMPLARS_TAG: u32 = 4; +const TIME_SERIES_METADATA_TAG: u32 = 5; -pub(crate) fn decode_remote_write_v2_request(is_zstd: bool, body: Bytes) -> Result { - let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_DECODE_ELAPSED.start_timer(); +struct BorrowedRequest<'a> { + symbols: Vec<&'a str>, + timeseries: Vec<&'a [u8]>, +} - // Match the v1 decoder's VictoriaMetrics fallback: some clients may send a - // mismatched content-encoding header, so try the other compression on failure. - let buf = if let Ok(buf) = try_decompress(is_zstd, &body[..]) { - buf - } else { - try_decompress(!is_zstd, &body[..])? - }; +impl<'a> BorrowedRequest<'a> { + fn decode(mut buf: &'a [u8]) -> std::result::Result { + let mut symbols = Vec::new(); + let mut timeseries = Vec::new(); - Request::decode(&buf[..]).context(error::DecodePromRemoteRequestSnafu) + while buf.has_remaining() { + let (tag, wire_type) = decode_key(&mut buf)?; + match tag { + 4 => { + let value = + take_length_delimited(wire_type, &mut buf).map_err(|mut error| { + error.push("Request", "symbols"); + error + })?; + let symbol = std::str::from_utf8(value).map_err(|_| { + let mut error = + DecodeError::new("invalid string value: data is not UTF-8 encoded"); + error.push("Request", "symbols"); + error + })?; + symbols.push(symbol); + } + 5 => { + let series = + take_length_delimited(wire_type, &mut buf).map_err(|mut error| { + error.push("Request", "timeseries"); + error + })?; + timeseries.push(series); + } + _ => skip_field(wire_type, tag, &mut buf, DecodeContext::default())?, + } + } + + Ok(Self { + symbols, + timeseries, + }) + } } pub(crate) struct RemoteWriteV2WriteRequests { @@ -81,19 +123,33 @@ pub(crate) struct RemoteWriteV2WriteRequests { pub semantic_index: SemanticIndexes, } -/// Converts a PRW v2 request into normal sample writes and native histogram writes. -/// -/// A metric name may appear in only one payload kind per request because the -/// metric-engine logical table is either a float metric or a native histogram. -pub(crate) fn into_write_requests(request: Request) -> Result { - let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_CONVERT_ELAPSED.start_timer(); - let Request { - symbols, - timeseries, - } = request; +pub(crate) fn decode_remote_write_v2( + is_zstd: bool, + body: Bytes, + native_histograms_enabled: bool, +) -> Result { + let decode_timer = crate::metrics::METRIC_HTTP_PROM_STORE_DECODE_ELAPSED.start_timer(); + // Match the v1 decoder's VictoriaMetrics fallback: some clients may send a + // mismatched content-encoding header, so try the other compression on failure. + let buf = if let Ok(buf) = try_decompress(is_zstd, &body[..]) { + buf + } else { + try_decompress(!is_zstd, &body[..])? + }; + let request = BorrowedRequest::decode(&buf).context(error::DecodePromRemoteRequestSnafu)?; + drop(decode_timer); + + let _convert_timer = crate::metrics::METRIC_HTTP_PROM_STORE_CONVERT_ELAPSED.start_timer(); + convert_remote_write_v2(request, native_histograms_enabled) +} + +fn convert_remote_write_v2( + request: BorrowedRequest<'_>, + native_histograms_enabled: bool, +) -> Result { ensure!( - symbols.first().map(|s| s.as_str()) == Some(""), + request.symbols.first().copied() == Some(""), error::InvalidPromRemoteRequestSnafu { msg: "remote write v2 symbols must start with an empty string".to_string(), } @@ -104,28 +160,40 @@ pub(crate) fn into_write_requests(request: Request) -> Result 0 { - histogram_count > 0 + let has_other_value_type = if counts.samples > 0 { + counts.histograms > 0 || histogram_tables .get(&prom_ctx) .is_some_and(|tables| tables.contains_key(&table_name)) @@ -143,42 +211,40 @@ pub(crate) fn into_write_requests(request: Request) -> Result 0 && histogram_count == 0 { - // Fast path for regular sample-only series. Move the resolved labels into - // the sample writer instead of cloning them for a histogram path we won't use. - let table_data = get_or_create_table_data( - &mut sample_tables, - prom_ctx, - table_name, - tags.len() + 2, - sample_count, - ); + let column_count = + tags.len() + .checked_add(2) + .with_context(|| error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 series has too many labels".to_string(), + })?; + let (writer, row_count) = if counts.samples > 0 { + ( + SeriesWriter::Samples(get_or_create_table_data( + &mut sample_tables, + prom_ctx, + table_name, + column_count, + counts.samples, + )), + counts.samples, + ) + } else { + ( + SeriesWriter::Histograms(get_or_create_table_data( + &mut histogram_tables, + prom_ctx, + table_name, + column_count, + counts.histograms, + )), + counts.histograms, + ) + }; - write_samples(table_data, series.samples, tags)?; - sample_count_total += sample_count as u64; - // The owned labels were moved above, so skip the mixed-series path below. - continue; - } - - if histogram_count > 0 { - let table_data = get_or_create_table_data( - &mut histogram_tables, - prom_ctx, - table_name, - tags.len() + 2, - histogram_count, - ); - - let mut histograms = series.histograms; - let Some(last_histogram) = histograms.pop() else { - continue; - }; - for histogram in &histograms { - write_native_histogram(table_data, histogram, tags.iter().cloned())?; - } - write_native_histogram(table_data, &last_histogram, tags.into_iter())?; - histogram_count_total += histogram_count as u64; - } + decode_series_leaves(series, Some(writer), tags, row_count, &mut scratch)?; + sample_count_total = checked_total(sample_count_total, counts.samples, "sample")?; + histogram_count_total = + checked_total(histogram_count_total, counts.histograms, "histogram")?; } Ok(RemoteWriteV2WriteRequests { @@ -201,14 +267,11 @@ pub(crate) fn into_write_requests(request: Request) -> Result Result<()> { - let Some(series_metadata) = &series.metadata else { - return Ok(()); - }; // Symbol 0 is the mandatory empty string: no help / no unit. if series_metadata.help_ref != 0 { symbol_ref(symbols, series_metadata.help_ref, "metadata help")?; @@ -258,6 +321,256 @@ fn metric_type_value(wire_type: i32) -> Option<&'static str> { } } +#[derive(Default)] +struct SeriesCounts { + samples: usize, + histograms: usize, +} + +fn scan_series( + mut buf: &[u8], + labels_refs: &mut Vec, + metadata: &mut Metadata, +) -> std::result::Result { + labels_refs.clear(); + metadata.clear(); + let mut counts = SeriesCounts::default(); + + while buf.has_remaining() { + let (tag, wire_type) = decode_key(&mut buf)?; + match tag { + TIME_SERIES_LABELS_REFS_TAG => { + uint32::merge_repeated(wire_type, labels_refs, &mut buf, DecodeContext::default()) + .map_err(|mut error| { + error.push("TimeSeries", "labels_refs"); + error + })? + } + TIME_SERIES_SAMPLES_TAG => { + take_length_delimited(wire_type, &mut buf).map_err(|mut error| { + error.push("TimeSeries", "samples"); + error + })?; + counts.samples = counts.samples.checked_add(1).ok_or_else(|| { + DecodeError::new("remote write v2 sample count overflows usize") + })?; + } + TIME_SERIES_HISTOGRAMS_TAG => { + take_length_delimited(wire_type, &mut buf).map_err(|mut error| { + error.push("TimeSeries", "histograms"); + error + })?; + counts.histograms = counts.histograms.checked_add(1).ok_or_else(|| { + DecodeError::new("remote write v2 histogram count overflows usize") + })?; + } + TIME_SERIES_EXEMPLARS_TAG => { + take_length_delimited(wire_type, &mut buf)?; + } + TIME_SERIES_METADATA_TAG => { + message::merge(wire_type, metadata, &mut buf, DecodeContext::default()).map_err( + |mut error| { + error.push("TimeSeries", "metadata"); + error + }, + )? + } + _ => skip_field(wire_type, tag, &mut buf, DecodeContext::default())?, + } + } + + Ok(counts) +} + +#[derive(Default)] +struct LeafScratch { + sample: Sample, + histogram: Histogram, + exemplar: Exemplar, +} + +enum SeriesWriter<'a> { + Samples(&'a mut TableData), + Histograms(&'a mut TableData), +} + +fn decode_series_leaves( + mut buf: &[u8], + mut writer: Option>, + mut tags: PromTags<'_>, + mut rows_remaining: usize, + scratch: &mut LeafScratch, +) -> Result<()> { + let mut sample_row_template = None; + + while buf.has_remaining() { + let (tag, wire_type) = decode_key(&mut buf).context(error::DecodePromRemoteRequestSnafu)?; + match tag { + TIME_SERIES_LABELS_REFS_TAG => { + skip_field(wire_type, tag, &mut buf, DecodeContext::default()) + .context(error::DecodePromRemoteRequestSnafu)? + } + TIME_SERIES_SAMPLES_TAG => { + scratch.sample.clear(); + message::merge( + wire_type, + &mut scratch.sample, + &mut buf, + DecodeContext::default(), + ) + .map_err(|mut error| { + error.push("TimeSeries", "samples"); + error + }) + .context(error::DecodePromRemoteRequestSnafu)?; + rows_remaining = rows_remaining.checked_sub(1).with_context(|| { + error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 sample count changed between scans".to_string(), + } + })?; + if let Some(SeriesWriter::Samples(table_data)) = &mut writer { + if sample_row_template.is_none() { + let timestamp_index = table_data.ensure_column_by_name( + greptime_timestamp(), + ColumnDataType::TimestampMillisecond, + SemanticType::Timestamp, + )?; + let value_index = table_data.ensure_column_by_name( + greptime_value(), + ColumnDataType::Float64, + SemanticType::Field, + )?; + let mut row = table_data.alloc_one_row(); + row_writer::write_tags( + table_data, + std::mem::take(&mut tags).into_iter(), + &mut row, + )?; + sample_row_template = Some((row, timestamp_index, value_index)); + } + + if let Some((row, timestamp_index, value_index)) = sample_row_template.as_mut() + { + row[*timestamp_index].value_data = Some( + ValueData::TimestampMillisecondValue(scratch.sample.timestamp), + ); + row[*value_index].value_data = + Some(ValueData::F64Value(scratch.sample.value)); + let row = if rows_remaining == 0 { + std::mem::take(row) + } else { + row.clone() + }; + table_data.add_row(row); + } + } + } + TIME_SERIES_HISTOGRAMS_TAG => { + scratch.histogram.clear(); + message::merge( + wire_type, + &mut scratch.histogram, + &mut buf, + DecodeContext::default(), + ) + .map_err(|mut error| { + error.push("TimeSeries", "histograms"); + error + }) + .context(error::DecodePromRemoteRequestSnafu)?; + rows_remaining = rows_remaining.checked_sub(1).with_context(|| { + error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 histogram count changed between scans".to_string(), + } + })?; + if let Some(SeriesWriter::Histograms(table_data)) = &mut writer { + if rows_remaining == 0 { + write_native_histogram( + table_data, + &scratch.histogram, + std::mem::take(&mut tags).into_iter(), + )?; + } else { + write_native_histogram( + table_data, + &scratch.histogram, + tags.iter().cloned(), + )?; + } + } + } + TIME_SERIES_EXEMPLARS_TAG => { + scratch.exemplar.clear(); + message::merge( + wire_type, + &mut scratch.exemplar, + &mut buf, + DecodeContext::default(), + ) + .map_err(|mut error| { + error.push("TimeSeries", "exemplars"); + error + }) + .context(error::DecodePromRemoteRequestSnafu)?; + } + TIME_SERIES_METADATA_TAG => { + take_length_delimited(wire_type, &mut buf) + .map_err(|mut error| { + error.push("TimeSeries", "metadata"); + error + }) + .context(error::DecodePromRemoteRequestSnafu)?; + } + _ => skip_field(wire_type, tag, &mut buf, DecodeContext::default()) + .context(error::DecodePromRemoteRequestSnafu)?, + } + } + + ensure!( + rows_remaining == 0, + error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 row count changed between scans".to_string(), + } + ); + + Ok(()) +} + +fn take_length_delimited<'a>( + wire_type: WireType, + buf: &mut &'a [u8], +) -> std::result::Result<&'a [u8], DecodeError> { + if wire_type != WireType::LengthDelimited { + return Err(DecodeError::new(format!( + "invalid wire type: {wire_type:?} (expected LengthDelimited)" + ))); + } + + let len = decode_varint(buf)?; + let len = + usize::try_from(len).map_err(|_| DecodeError::new("length delimiter exceeds usize"))?; + if len > buf.len() { + return Err(DecodeError::new("buffer underflow")); + } + let (value, remaining) = buf.split_at(len); + *buf = remaining; + Ok(value) +} + +fn checked_total(total: u64, count: usize, name: &str) -> Result { + let count = + u64::try_from(count) + .ok() + .with_context(|| error::InvalidPromRemoteRequestSnafu { + msg: format!("remote write v2 {name} count exceeds u64"), + })?; + total + .checked_add(count) + .with_context(|| error::InvalidPromRemoteRequestSnafu { + msg: format!("remote write v2 {name} count overflows u64"), + }) +} + fn get_or_create_table_data( tables: &mut HashMap>, prom_ctx: PromCtx, @@ -275,46 +588,10 @@ fn get_or_create_table_data( } } -fn write_samples( - table_data: &mut TableData, - mut samples: Vec, - tags: PromTags, -) -> Result<()> { - let Some(last_sample) = samples.pop() else { - return Ok(()); - }; - - for sample in &samples { - write_sample(table_data, sample, tags.iter().cloned())?; - } - - write_sample(table_data, &last_sample, tags.into_iter()) -} - -fn write_sample( - table_data: &mut TableData, - sample: &Sample, - tags: impl Iterator, -) -> Result<()> { - let mut row = table_data.alloc_one_row(); - row_writer::write_ts_to_millis( - table_data, - greptime_timestamp(), - Some(sample.timestamp), - Precision::Millisecond, - &mut row, - )?; - row_writer::write_f64(table_data, greptime_value(), sample.value, &mut row)?; - row_writer::write_tags(table_data, tags, &mut row)?; - table_data.add_row(row); - - Ok(()) -} - -fn write_native_histogram( +fn write_native_histogram<'a>( table_data: &mut TableData, histogram: &Histogram, - tags: impl Iterator, + tags: impl Iterator, ) -> Result<()> { // Persist both int and float families into the logical table schema. Only one // family is populated per row; the other is written as NULL so PromQL can @@ -579,7 +856,7 @@ fn validate_native_histogram_spans( let span_len = spans .iter() .try_fold(0usize, |sum, span| sum.checked_add(span.length as usize)) - .context(error::InvalidPromRemoteRequestSnafu { + .with_context(|| error::InvalidPromRemoteRequestSnafu { msg: format!("remote write v2 native histogram {name} spans overflow"), })?; ensure!( @@ -606,13 +883,13 @@ fn validate_native_histogram_spans( current_index = if span_index == 0 { span.offset } else { - current_index.checked_add(span.offset).context( + current_index.checked_add(span.offset).with_context(|| { error::InvalidPromRemoteRequestSnafu { msg: format!( "remote write v2 native histogram {name} span index overflows i32" ), - }, - )? + } + })? }; for _ in 0..span.length { @@ -701,7 +978,7 @@ fn validate_integer_native_histogram_counts( .try_fold(zero_count, |total, bucket| { total.checked_add(*bucket as u64) }) - .context(error::InvalidPromRemoteRequestSnafu { + .with_context(|| error::InvalidPromRemoteRequestSnafu { msg: "remote write v2 native histogram bucket total overflows u64".to_string(), })?; ensure!( @@ -824,11 +1101,12 @@ fn bucket_counts_from_deltas(deltas: &[i64]) -> Result> { let mut buckets = Vec::with_capacity(deltas.len()); for delta in deltas { - count = count - .checked_add(*delta) - .context(error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 native histogram bucket count overflows i64".to_string(), - })?; + count = + count + .checked_add(*delta) + .with_context(|| error::InvalidPromRemoteRequestSnafu { + msg: "remote write v2 native histogram bucket count overflows i64".to_string(), + })?; ensure!( count >= 0, error::InvalidPromRemoteRequestSnafu { @@ -841,11 +1119,11 @@ fn bucket_counts_from_deltas(deltas: &[i64]) -> Result> { Ok(buckets) } -fn ensure_no_internal_histogram_labels(tags: &PromTags) -> Result<()> { +fn ensure_no_internal_histogram_labels(tags: &PromTags<'_>) -> Result<()> { // The histogram field column is generated from the protobuf payload. for (name, _) in tags { ensure!( - name != greptime_native_histogram() && name != NATIVE_HISTOGRAM_FIELD, + *name != greptime_native_histogram() && *name != NATIVE_HISTOGRAM_FIELD, error::InvalidPromRemoteRequestSnafu { msg: format!( "remote write v2 label `{name}` conflicts with an internal native histogram label" @@ -858,12 +1136,12 @@ fn ensure_no_internal_histogram_labels(tags: &PromTags) -> Result<()> { } fn resolve_series_labels<'a>( - symbols: &'a [String], - series: &TimeSeries, + symbols: &'a [&str], + labels_refs: &[u32], label_names: &mut HashSet<&'a str>, -) -> Result { +) -> Result> { ensure!( - series.labels_refs.len().is_multiple_of(2), + labels_refs.len().is_multiple_of(2), error::InvalidPromRemoteRequestSnafu { msg: "remote write v2 labels_refs must contain name/value pairs".to_string(), } @@ -871,10 +1149,10 @@ fn resolve_series_labels<'a>( let mut prom_ctx = PromCtx::default(); let mut table_name = None; - let mut tags = Vec::with_capacity(series.labels_refs.len() / 2); + let mut tags = Vec::with_capacity(labels_refs.len() / 2); label_names.clear(); - for pair in series.labels_refs.chunks_exact(2) { + for pair in labels_refs.chunks_exact(2) { let name = symbol_ref(symbols, pair[0], "label name")?; let value = symbol_ref(symbols, pair[1], "label value")?; validate_label(name)?; @@ -893,10 +1171,10 @@ fn resolve_series_labels<'a>( continue; } - tags.push((name.to_string(), value.to_string())); + tags.push((name, value.to_string())); } - let table_name = table_name.context(error::InvalidPromRemoteRequestSnafu { + let table_name = table_name.with_context(|| error::InvalidPromRemoteRequestSnafu { msg: "missing '__name__' label in time-series".to_string(), })?; ensure!( @@ -920,10 +1198,15 @@ fn validate_label(name: &str) -> Result<()> { Ok(()) } -fn symbol_ref<'a>(symbols: &'a [String], idx: u32, field: &str) -> Result<&'a str> { +fn symbol_ref<'a>(symbols: &'a [&str], idx: u32, field: &str) -> Result<&'a str> { + let idx = usize::try_from(idx) + .ok() + .with_context(|| error::InvalidPromRemoteRequestSnafu { + msg: format!("remote write v2 {field} symbol reference exceeds usize"), + })?; symbols - .get(idx as usize) - .map(String::as_str) + .get(idx) + .copied() .with_context(|| error::InvalidPromRemoteRequestSnafu { msg: format!( "remote write v2 {field} symbol reference {idx} is out of range, symbols len: {}", @@ -994,8 +1277,12 @@ pub mod test_util { use api::greptime_proto::io::prometheus::write::v2::{Histogram, Request, Sample, TimeSeries}; use api::v1::RowInsertRequest; use bytes::Bytes; + use prost::Message; + use snafu::ResultExt; - use crate::error::Result; + use crate::error::{self, Result}; + use crate::prom_remote_write::try_decompress; + use crate::prom_store::snappy_compress; pub fn request_with_labels_and_samples( labels: Vec<(&str, &str)>, @@ -1012,13 +1299,42 @@ pub mod test_util { } pub fn decode_request(is_zstd: bool, body: Bytes) -> Result { - super::decode_remote_write_v2_request(is_zstd, body) + let buf = if let Ok(buf) = try_decompress(is_zstd, &body[..]) { + buf + } else { + try_decompress(!is_zstd, &body[..])? + }; + Request::decode(&buf[..]).context(error::DecodePromRemoteRequestSnafu) } pub fn write_requests( request: Request, ) -> Result<(Vec, Vec, u64, u64)> { - let requests = super::into_write_requests(request)?; + let body = Bytes::from(snappy_compress(&request.encode_to_vec())?); + decode_write_requests(false, body, true) + } + + pub fn decode_write_requests( + is_zstd: bool, + body: Bytes, + native_histograms_enabled: bool, + ) -> Result<(Vec, Vec, u64, u64)> { + let requests = super::decode_remote_write_v2(is_zstd, body, native_histograms_enabled)?; + Ok(( + requests.samples.all_req().collect(), + requests.histograms.all_req().collect(), + requests.sample_count, + requests.histogram_count, + )) + } + + pub fn decode_uncompressed_write_requests( + body: &[u8], + native_histograms_enabled: bool, + ) -> Result<(Vec, Vec, u64, u64)> { + let request = + super::BorrowedRequest::decode(body).context(error::DecodePromRemoteRequestSnafu)?; + let requests = super::convert_remote_write_v2(request, native_histograms_enabled)?; Ok(( requests.samples.all_req().collect(), requests.histograms.all_req().collect(), @@ -1109,7 +1425,7 @@ mod tests { let body = Bytes::from(crate::prom_store::snappy_compress(&request.encode_to_vec()).unwrap()); - let decoded = decode_remote_write_v2_request(false, body).unwrap(); + let decoded = test_util::decode_request(false, body.clone()).unwrap(); assert_eq!(decoded.symbols, request.symbols); assert_eq!(decoded.timeseries.len(), 1); @@ -1117,11 +1433,291 @@ mod tests { assert_eq!(decoded.timeseries[0].samples.len(), 1); assert_eq!(decoded.timeseries[0].samples[0].value, 42.0); assert_eq!(decoded.timeseries[0].metadata.as_ref().unwrap().r#type, 1); + assert_eq!( + decode_remote_write_v2(true, body, true) + .unwrap() + .sample_count, + 1 + ); + } + + #[test] + fn test_fused_decoder_accepts_arbitrary_field_order_and_split_labels() { + let mut first_sample = Sample { + value: 42.0, + timestamp: 1000, + start_timestamp: 500, + } + .encode_to_vec(); + first_sample.extend(varint_field(90, 1)); + let second_sample = Sample { + value: 43.0, + timestamp: 2000, + start_timestamp: 1000, + } + .encode_to_vec(); + + let mut series = encoded_message_field(2, &first_sample); + series.extend(packed_u32_field(1, &[1])); + series.extend(varint_field(90, 1)); + series.extend(encoded_message_field(2, &second_sample)); + series.extend(varint_field(1, 2)); + + let mut wire = encoded_message_field(5, &series); + wire.extend(string_field(4, b"")); + wire.extend(varint_field(90, 1)); + wire.extend(string_field(4, METRIC_NAME_LABEL.as_bytes())); + wire.extend(encoded_message_field(5, &packed_u32_field(1, &[99]))); + wire.extend(string_field(4, b"http_requests_total")); + + let requests = decode_wire(&wire, true).unwrap(); + assert_eq!(requests.sample_count, 2); + assert_eq!(requests.histogram_count, 0); + let rows = requests.samples.all_req().next().unwrap().rows.unwrap(); + assert_eq!(rows.rows.len(), 2); + assert_eq!( + rows.schema + .iter() + .map(|column| column.column_name.as_str()) + .collect::>(), + vec![greptime_timestamp(), greptime_value()] + ); + assert_eq!( + rows.rows[0].values[1].value_data, + Some(ValueData::F64Value(42.0)) + ); + assert_eq!( + rows.rows[1].values[1].value_data, + Some(ValueData::F64Value(43.0)) + ); + } + + #[test] + fn test_fused_decoder_accepts_histogram_before_labels_and_unknown_fields() { + let mut histogram = Histogram { + count: Some(Count::CountInt(0)), + zero_count: Some(ZeroCount::ZeroCountInt(0)), + timestamp: 2000, + start_timestamp: 1000, + ..Default::default() + } + .encode_to_vec(); + histogram.extend(varint_field(90, 1)); + + let mut series = encoded_message_field(3, &histogram); + series.extend(packed_u32_field(1, &[1, 2])); + let wire = request_wire(&["", METRIC_NAME_LABEL, "metric"], &[series]); + + let requests = decode_wire(&wire, true).unwrap(); + assert_eq!(requests.histogram_count, 1); + let rows = requests.histograms.all_req().next().unwrap().rows.unwrap(); + assert_eq!( + histogram_field_value(&rows, 0, START_TIMESTAMP_FIELD), + Some(ValueData::TimestampMillisecondValue(1000)) + ); + } + + #[test] + fn test_fused_decoder_rejects_malformed_wire() { + let mut wrong_request_wire = varint_field(4, 0); + let invalid_utf8 = string_field(4, &[0xff]); + + let mut truncated_request = Vec::new(); + prost::encoding::encode_key(5, WireType::LengthDelimited, &mut truncated_request); + prost::encoding::encode_varint(2, &mut truncated_request); + truncated_request.push(0); + + let mut oversized_request = Vec::new(); + prost::encoding::encode_key(4, WireType::LengthDelimited, &mut oversized_request); + prost::encoding::encode_varint(u64::MAX, &mut oversized_request); + + let malformed_varint = vec![0x80; 10]; + + let mut wrong_labels_wire = Vec::new(); + prost::encoding::encode_key(1, WireType::SixtyFourBit, &mut wrong_labels_wire); + wrong_labels_wire.extend([0; 8]); + wrong_labels_wire = request_wire(&[""], &[std::mem::take(&mut wrong_labels_wire)]); + + let wrong_series_wire = request_wire(&[""], &[varint_field(2, 0)]); + + let mut truncated_series = Vec::new(); + prost::encoding::encode_key(2, WireType::LengthDelimited, &mut truncated_series); + prost::encoding::encode_varint(2, &mut truncated_series); + truncated_series.push(0x08); + let truncated_series = request_wire(&[""], &[truncated_series]); + + let mut oversized_series = Vec::new(); + prost::encoding::encode_key(3, WireType::LengthDelimited, &mut oversized_series); + prost::encoding::encode_varint(u64::MAX, &mut oversized_series); + let oversized_series = request_wire(&[""], &[oversized_series]); + + let invalid_sample = series_wire(&[1, 2], 2, &varint_field(1, 1)); + let invalid_sample = request_wire(&["", METRIC_NAME_LABEL, "metric"], &[invalid_sample]); + let invalid_histogram = series_wire(&[1, 2], 3, &varint_field(3, 1)); + let invalid_histogram = + request_wire(&["", METRIC_NAME_LABEL, "metric"], &[invalid_histogram]); + + for (name, wire) in [ + ( + "wrong request wire type", + std::mem::take(&mut wrong_request_wire), + ), + ("invalid utf8", invalid_utf8), + ("truncated request", truncated_request), + ("oversized request length", oversized_request), + ("malformed varint", malformed_varint), + ("wrong labels wire type", wrong_labels_wire), + ("wrong series wire type", wrong_series_wire), + ("truncated series", truncated_series), + ("oversized series length", oversized_series), + ("invalid sample", invalid_sample), + ("invalid histogram", invalid_histogram), + ] { + let error = decode_wire_error(&wire, true, name); + assert!( + matches!(error, error::Error::DecodePromRemoteRequest { .. }), + "{name}: {error}" + ); + } + } + + #[test] + fn test_fused_decoder_rejects_malformed_ignored_messages() { + for tag in [4, 5] { + let mut series = series_wire(&[1, 2], 2, &Sample::default().encode_to_vec()); + series.extend(encoded_message_field(tag, &[0x08])); + let wire = request_wire(&["", METRIC_NAME_LABEL, "metric"], &[series]); + + let error = decode_wire_error(&wire, true, "malformed ignored message"); + assert!(matches!( + error, + error::Error::DecodePromRemoteRequest { .. } + )); + } + } + + #[test] + fn test_fused_decoder_ignores_exemplar_symbol_refs() { + let request = Request { + symbols: vec![ + String::new(), + METRIC_NAME_LABEL.to_string(), + "metric".to_string(), + ], + timeseries: vec![TimeSeries { + labels_refs: vec![1, 2], + samples: vec![Sample::default()], + exemplars: vec![Exemplar { + labels_refs: vec![99], + ..Default::default() + }], + ..Default::default() + }], + }; + + assert_eq!(decode_test_request(request).unwrap().sample_count, 1); + } + + #[test] + fn test_fused_decoder_preserves_empty_request_and_series_behavior() { + let request = Request { + symbols: vec![ + String::new(), + METRIC_NAME_LABEL.to_string(), + "metric".to_string(), + "job".to_string(), + ], + timeseries: vec![ + TimeSeries { + labels_refs: vec![99, 99], + ..Default::default() + }, + TimeSeries { + labels_refs: vec![1], + ..Default::default() + }, + TimeSeries { + labels_refs: vec![3, 2, 3, 2], + ..Default::default() + }, + TimeSeries::default(), + ], + }; + + let requests = decode_test_request(request).unwrap(); + assert_eq!(requests.sample_count, 0); + assert_eq!(requests.histogram_count, 0); + + let requests = decode_test_request(Request { + symbols: vec![String::new()], + timeseries: Vec::new(), + }) + .unwrap(); + assert_eq!(requests.sample_count, 0); + assert!(decode_wire(&[], true).is_err()); + assert!(decode_wire(&[0x0a, 0x00], true).is_err()); + } + + #[test] + fn test_fused_decoder_rejects_duplicate_resolved_label_names() { + let request = Request { + symbols: vec![ + String::new(), + METRIC_NAME_LABEL.to_string(), + "metric".to_string(), + "job".to_string(), + "job".to_string(), + "api".to_string(), + "worker".to_string(), + ], + timeseries: vec![TimeSeries { + labels_refs: vec![1, 2, 3, 5, 4, 6], + samples: vec![Sample::default()], + ..Default::default() + }], + }; + + assert_invalid( + "duplicate resolved labels", + request, + "label name `job` is repeated", + ); + } + + #[test] + fn test_fused_decoder_pins_experimental_error_precedence() { + let histogram = series_wire(&[1, 2], 3, &Histogram::default().encode_to_vec()); + let malformed_sample = series_wire(&[1, 2], 2, &[0x08]); + + let wire = request_wire( + &["", METRIC_NAME_LABEL, "metric"], + &[histogram.clone(), malformed_sample.clone()], + ); + let error = decode_wire_error(&wire, false, "histogram before malformed series"); + assert!(error.to_string().contains("ingestion is experimental")); + + let wire = request_wire( + &["", METRIC_NAME_LABEL, "metric"], + &[malformed_sample, histogram.clone()], + ); + let error = decode_wire_error(&wire, false, "malformed series before histogram"); + assert!(matches!( + error, + error::Error::DecodePromRemoteRequest { .. } + )); + + let missing_name = series_wire(&[3, 4], 2, &Sample::default().encode_to_vec()); + let wire = request_wire( + &["", METRIC_NAME_LABEL, "metric", "job", "api"], + &[missing_name, histogram], + ); + let error = decode_wire_error(&wire, false, "conversion error before histogram"); + assert!(error.to_string().contains("missing '__name__'")); } #[test] fn test_into_context_req_samples() { - let ctx_req = into_write_requests(test_util::request_with_labels_and_samples( + let ctx_req = decode_test_request(test_util::request_with_labels_and_samples( vec![ (METRIC_NAME_LABEL, "http_requests_total"), ("job", "api"), @@ -1183,11 +1779,19 @@ mod tests { rows.rows[1].values[1].value_data, Some(ValueData::F64Value(43.0)) ); + assert_eq!( + rows.rows[1].values[2].value_data, + Some(ValueData::StringValue("api".to_string())) + ); + assert_eq!( + rows.rows[1].values[3].value_data, + Some(ValueData::StringValue("localhost:9090".to_string())) + ); } #[test] fn test_into_context_req_special_labels() { - let ctx_req = into_write_requests(test_util::request_with_labels_and_samples( + let ctx_req = decode_test_request(test_util::request_with_labels_and_samples( vec![ (METRIC_NAME_LABEL, "cpu_usage"), (DATABASE_LABEL, "tenant_a"), @@ -1518,7 +2122,7 @@ mod tests { #[test] fn test_into_context_req_allows_nan_observations_outside_buckets() { - into_write_requests(request_with_histogram(Histogram { + decode_test_request(request_with_histogram(Histogram { count: Some(Count::CountInt(2)), sum: f64::NAN, positive_spans: vec![BucketSpan { @@ -1533,7 +2137,7 @@ mod tests { #[test] fn test_into_context_req_allows_empty_label_values() { - let ctx_req = into_write_requests(test_util::request_with_labels_and_samples( + let ctx_req = decode_test_request(test_util::request_with_labels_and_samples( vec![(METRIC_NAME_LABEL, "metric"), ("job", "")], vec![Sample { value: 1.0, @@ -1649,7 +2253,7 @@ mod tests { }]; histogram.negative_deltas = vec![1]; } - into_write_requests(request_with_histogram(histogram.clone())).unwrap(); + decode_test_request(request_with_histogram(histogram.clone())).unwrap(); let beyond = max_index + 1; if positive { @@ -1694,7 +2298,7 @@ mod tests { ], }; - let ctx_req = into_write_requests(request).unwrap(); + let ctx_req = decode_test_request(request).unwrap(); assert_eq!(ctx_req.sample_count, 1); assert_eq!(ctx_req.histogram_count, 1); @@ -1708,7 +2312,7 @@ mod tests { test_util::request_with_labels_and_samples(vec![(METRIC_NAME_LABEL, "metric")], vec![]); request.timeseries[0].histograms.push(Histogram::default()); - let ctx_req = into_write_requests(request).unwrap(); + let ctx_req = decode_test_request(request).unwrap(); assert_eq!(ctx_req.sample_count, 0); assert_eq!(ctx_req.histogram_count, 1); @@ -1744,7 +2348,7 @@ mod tests { #[test] fn test_into_context_req_preserves_histogram_start_timestamp() { - let ctx_req = into_write_requests(test_util::request_with_labels_and_histograms( + let ctx_req = decode_test_request(test_util::request_with_labels_and_histograms( vec![(METRIC_NAME_LABEL, "metric")], vec![Histogram { timestamp: 2000, @@ -1774,7 +2378,7 @@ mod tests { ); request.timeseries[0].histograms.push(Histogram::default()); - let err = match into_write_requests(request) { + let err = match decode_test_request(request) { Ok(_) => panic!("expected invalid request error"), Err(err) => err, }; @@ -1790,7 +2394,7 @@ mod tests { assert_eq!(greptime_native_histogram(), "custom_native_histogram"); let err = ensure_no_internal_histogram_labels(&vec![( - NATIVE_HISTOGRAM_FIELD.to_string(), + NATIVE_HISTOGRAM_FIELD, "user_value".to_string(), )]) .unwrap_err(); @@ -1837,7 +2441,7 @@ mod tests { ], }; - let ctx_req = into_write_requests(request).unwrap(); + let ctx_req = decode_test_request(request).unwrap(); assert_eq!(ctx_req.histogram_count, 2); let mut inserts = ctx_req.histograms.all_req().collect::>(); @@ -1883,6 +2487,65 @@ mod tests { )); } + fn decode_wire( + wire: &[u8], + native_histograms_enabled: bool, + ) -> Result { + let body = Bytes::from(crate::prom_store::snappy_compress(wire).unwrap()); + decode_remote_write_v2(false, body, native_histograms_enabled) + } + + fn decode_wire_error(wire: &[u8], native_histograms_enabled: bool, name: &str) -> error::Error { + match decode_wire(wire, native_histograms_enabled) { + Ok(_) => panic!("{name}: expected decoder error"), + Err(error) => error, + } + } + + fn request_wire(symbols: &[&str], series: &[Vec]) -> Vec { + let mut wire = Vec::new(); + for symbol in symbols { + wire.extend(string_field(4, symbol.as_bytes())); + } + for series in series { + wire.extend(encoded_message_field(5, series)); + } + wire + } + + fn series_wire(labels_refs: &[u32], leaf_tag: u32, leaf: &[u8]) -> Vec { + let mut series = packed_u32_field(1, labels_refs); + series.extend(encoded_message_field(leaf_tag, leaf)); + series + } + + fn string_field(tag: u32, value: &[u8]) -> Vec { + encoded_message_field(tag, value) + } + + fn encoded_message_field(tag: u32, value: &[u8]) -> Vec { + let mut field = Vec::new(); + prost::encoding::encode_key(tag, WireType::LengthDelimited, &mut field); + prost::encoding::encode_varint(u64::try_from(value.len()).unwrap(), &mut field); + field.extend_from_slice(value); + field + } + + fn varint_field(tag: u32, value: u64) -> Vec { + let mut field = Vec::new(); + prost::encoding::encode_key(tag, WireType::Varint, &mut field); + prost::encoding::encode_varint(value, &mut field); + field + } + + fn packed_u32_field(tag: u32, values: &[u32]) -> Vec { + let mut packed = Vec::new(); + for value in values { + prost::encoding::encode_varint(u64::from(*value), &mut packed); + } + encoded_message_field(tag, &packed) + } + fn request_with_sample(labels: Vec<(&str, &str)>) -> Request { test_util::request_with_labels_and_samples( labels, @@ -1901,8 +2564,21 @@ mod tests { ) } + fn decode_test_request(request: Request) -> Result { + decode_test_request_with_histograms(request, true) + } + + fn decode_test_request_with_histograms( + request: Request, + native_histograms_enabled: bool, + ) -> Result { + let body = + Bytes::from(crate::prom_store::snappy_compress(&request.encode_to_vec()).unwrap()); + decode_remote_write_v2(false, body, native_histograms_enabled) + } + fn assert_invalid(name: &str, request: Request, expected: &str) { - let err = match into_write_requests(request) { + let err = match decode_test_request(request) { Ok(_) => panic!("{name}: expected invalid request error"), Err(err) => err, }; @@ -1993,7 +2669,7 @@ mod tests { >; fn decoded_index(request: Request) -> DecodedIndex { - let encoded = into_write_requests(request) + let encoded = decode_test_request(request) .unwrap() .semantic_index .encode("public") @@ -2045,7 +2721,7 @@ mod tests { assert!(!untyped.contains_key(SEMANTIC_METRIC_TYPE)); assert!(!untyped.contains_key(SEMANTIC_METRIC_METADATA_QUALITY)); - let requests = into_write_requests(metadata_request(vec![( + let requests = decode_test_request(metadata_request(vec![( vec![(METRIC_NAME_LABEL, "untyped_unitless_total")], metadata::MetricType::Unspecified as i32, None, @@ -2054,7 +2730,7 @@ mod tests { assert!(requests.semantic_index.is_empty()); // No metadata at all behaves the same. - let requests = into_write_requests(test_util::request_with_labels_and_samples( + let requests = decode_test_request(test_util::request_with_labels_and_samples( vec![(METRIC_NAME_LABEL, "bare_total")], vec![Sample { value: 1.0, @@ -2123,7 +2799,7 @@ mod tests { None, )]); request.timeseries[0].metadata.as_mut().unwrap().unit_ref = 999; - let err = into_write_requests(request).err().unwrap(); + let err = decode_test_request(request).err().unwrap(); assert!(err.to_string().contains("out of range"), "{err}"); // unit_ref must be validated even when the type is UNSPECIFIED @@ -2134,7 +2810,7 @@ mod tests { None, )]); request.timeseries[0].metadata.as_mut().unwrap().unit_ref = 999; - let err = into_write_requests(request).err().unwrap(); + let err = decode_test_request(request).err().unwrap(); assert!(err.to_string().contains("out of range"), "{err}"); // help_ref is validated although help is never persisted. @@ -2144,7 +2820,7 @@ mod tests { None, )]); request.timeseries[0].metadata.as_mut().unwrap().help_ref = 999; - let err = into_write_requests(request).err().unwrap(); + let err = decode_test_request(request).err().unwrap(); assert!(err.to_string().contains("out of range"), "{err}"); } } diff --git a/src/servers/src/row_writer.rs b/src/servers/src/row_writer.rs index 45afcbca6de..61fd9d77e5b 100644 --- a/src/servers/src/row_writer.rs +++ b/src/servers/src/row_writer.rs @@ -92,6 +92,26 @@ impl TableData { Ok(index) } + /// Ensures a default column schema without allocating its name when it already exists. + pub(crate) fn ensure_column_by_name( + &mut self, + name: &str, + datatype: ColumnDataType, + semantic_type: SemanticType, + ) -> Result { + if let Some(index) = self.column_indexes.get(name).copied() { + check_schema(datatype, semantic_type, &self.schema[index])?; + return Ok(index); + } + + self.ensure_column(ColumnSchema { + column_name: name.to_string(), + datatype: datatype as i32, + semantic_type: semantic_type as i32, + ..Default::default() + }) + } + #[allow(dead_code)] pub fn columns(&self) -> &Vec { &self.schema @@ -211,11 +231,14 @@ impl MultiTableData { } /// Write data as tags into the table data. -pub fn write_tags( +pub fn write_tags( table_data: &mut TableData, - tags: impl Iterator, + tags: impl Iterator, one_row: &mut Vec, -) -> Result<()> { +) -> Result<()> +where + K: AsRef + Into, +{ let ktv_iter = tags.map(|(k, v)| (k, ColumnDataType::String, Some(ValueData::StringValue(v)))); write_by_semantic_type(table_data, SemanticType::Tag, ktv_iter, one_row) } @@ -326,12 +349,15 @@ pub(crate) fn write_by_schema( Ok(()) } -fn write_by_semantic_type( +fn write_by_semantic_type( table_data: &mut TableData, semantic_type: SemanticType, - ktv_iter: impl Iterator)>, + ktv_iter: impl Iterator)>, one_row: &mut Vec, -) -> Result<()> { +) -> Result<()> +where + K: AsRef + Into, +{ let TableData { schema, column_indexes, @@ -339,12 +365,13 @@ fn write_by_semantic_type( } = table_data; for (name, datatype, value) in ktv_iter { - let index = column_indexes.get(&name); + let index = column_indexes.get(name.as_ref()).copied(); if let Some(index) = index { - check_schema(datatype, semantic_type, &schema[*index])?; - one_row[*index].value_data = value; + check_schema(datatype, semantic_type, &schema[index])?; + one_row[index].value_data = value; } else { let index = schema.len(); + let name = name.into(); schema.push(ColumnSchema { column_name: name.clone(), datatype: datatype as i32, diff --git a/src/servers/tests/prom_remote_write_v2_test.rs b/src/servers/tests/prom_remote_write_v2_test.rs index 94d260d3fac..75fee046f7c 100644 --- a/src/servers/tests/prom_remote_write_v2_test.rs +++ b/src/servers/tests/prom_remote_write_v2_test.rs @@ -77,7 +77,7 @@ fn test_decode_remote_write_v2_native_histogram_dump() { assert_eq!(histogram.timestamp, 1782358160412); let (sample_inserts, histogram_inserts, sample_count, histogram_count) = - remote_write_v2::write_requests(decoded).unwrap(); + remote_write_v2::decode_write_requests(false, Bytes::from_static(BODY), true).unwrap(); assert!(sample_inserts.is_empty()); assert_eq!(sample_count, 0); assert_eq!(histogram_count, 1); From 6538db6d618279f640958c40cdbb5a1796fb40cd Mon Sep 17 00:00:00 2001 From: fys <40801205+fengys1996@users.noreply.github.com> Date: Fri, 14 Aug 2026 06:21:17 +0000 Subject: [PATCH 46/92] refactor(json2): push down json2 type hints to parquet reads (#8833) * refactor(mito): push down json2 type hints to parquet reads Signed-off-by: fys * refactor(mito): share json2 target types with arc Signed-off-by: fys * refactor(mito): derive json2 output schema from target types Signed-off-by: fys * refactor(mito): simplify read columns construction Signed-off-by: fys * fix: cargo check Signed-off-by: fys * fix(mito): reject JSON hints for non-JSON2 read columns Signed-off-by: fys * fix: do not pushdown json type hint of non-json2-col Signed-off-by: fys * fix: unit test Signed-off-by: fys * refactor(query): simplify JSON type hint application Signed-off-by: fys * refactor: clean up JSON2 type hint handling Signed-off-by: fys * refactor(mito2): keep JSON2 hints with flat read format Signed-off-by: fys * fix(mito): use raw parquet projection for output schema Signed-off-by: fys * fix(query): note JSON2 hint scope limitation Signed-off-by: fys * refactor(mito): store JSON target types as native types Signed-off-by: fys * test(json2): cover join hint qualifier limitation Signed-off-by: fys * refactor(mito): remove JSON2 fallback from compat cast Signed-off-by: fys * docs(mito): document ReadColumns ordering contract Signed-off-by: fys * fix: cargo clippy Signed-off-by: fys --------- Signed-off-by: fys --- src/cmd/src/datanode/parquetbench.rs | 2 +- src/mito2/src/cache.rs | 20 +- src/mito2/src/compaction/reader.rs | 33 +- src/mito2/src/engine/scan_test.rs | 50 +-- src/mito2/src/memtable/bulk/context.rs | 4 +- src/mito2/src/read/compat.rs | 50 +-- src/mito2/src/read/flat_projection.rs | 83 ++--- src/mito2/src/read/last_row.rs | 15 +- src/mito2/src/read/projection.rs | 3 +- src/mito2/src/read/range_cache.rs | 4 +- src/mito2/src/read/read_columns.rs | 125 +++---- src/mito2/src/read/scan_region.rs | 327 ++++++++++++------ src/mito2/src/sst/parquet/file_range.rs | 9 +- src/mito2/src/sst/parquet/flat_format.rs | 79 ++++- src/mito2/src/sst/parquet/format.rs | 154 +++++---- src/mito2/src/sst/parquet/json_align/mod.rs | 2 - .../src/sst/parquet/json_align/schema.rs | 250 ------------- src/mito2/src/sst/parquet/prefilter.rs | 30 +- src/mito2/src/sst/parquet/reader.rs | 16 +- .../src/optimizer/json_type_concretize.rs | 91 +++-- src/table/src/table/adapter.rs | 5 + .../common/types/json/json2_limit.result | 69 ++++ .../common/types/json/json2_limit.sql | 39 +++ 23 files changed, 729 insertions(+), 731 deletions(-) delete mode 100644 src/mito2/src/sst/parquet/json_align/schema.rs diff --git a/src/cmd/src/datanode/parquetbench.rs b/src/cmd/src/datanode/parquetbench.rs index 301c6440ec1..2b46a775f76 100644 --- a/src/cmd/src/datanode/parquetbench.rs +++ b/src/cmd/src/datanode/parquetbench.rs @@ -554,7 +554,7 @@ async fn run_flat_prune_iteration( let reader_builder = ParquetReaderBuilder::new(table_dir, path_type, file_handle, object_store) .expected_metadata(Some(region_meta)) .cache(CacheStrategy::Disabled) - .projection(projection.map(ReadColumns::from_deduped_column_ids)); + .projection(projection.map(ReadColumns::new)); let mut reader_metrics = ReaderMetrics::default(); let start = Instant::now(); let mut stats = IterationStats::default(); diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 27db7daa5ab..ac1f5b66918 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -49,7 +49,7 @@ use puffin::puffin_manager::cache::{PuffinMetadataCache, PuffinMetadataCacheRef} use smallvec::SmallVec; use snafu::{OptionExt, ResultExt}; use store_api::metadata::{RegionMetadata, RegionMetadataRef}; -use store_api::storage::{ConcreteDataType, FileId, RegionId, TimeSeriesRowSelector}; +use store_api::storage::{ColumnId, ConcreteDataType, FileId, RegionId, TimeSeriesRowSelector}; use crate::cache::cache_size::parquet_meta_size; use crate::cache::file_cache::{FileType, IndexKey}; @@ -65,6 +65,7 @@ use crate::memtable::record_batch_estimated_size; use crate::metrics::{CACHE_BYTES, CACHE_EVICTION, CACHE_HIT, CACHE_MISS}; use crate::read::Batch; use crate::read::range_cache::{RangeScanCacheKey, RangeScanCacheValue}; +use crate::read::read_columns::JsonTargetTypes; use crate::sst::file::{RegionFileId, RegionIndexId}; use crate::sst::parquet::PARQUET_METADATA_KEY; use crate::sst::parquet::read_columns::ParquetReadColumns; @@ -2025,6 +2026,11 @@ pub struct SelectorResultValue { pub result: SelectorResult, /// The read columns of rows. pub read_cols: ParquetReadColumns, + /// JSON2 target types used by flat-format reads. + /// + /// JSON2 projection is query-driven; the same parquet columns can produce + /// different cached batches under different type hints. + pub json_target_types: JsonTargetTypes, } impl SelectorResultValue { @@ -2033,6 +2039,7 @@ impl SelectorResultValue { SelectorResultValue { result: SelectorResult::PrimaryKey(result), read_cols, + json_target_types: Arc::default(), } } @@ -2040,21 +2047,26 @@ impl SelectorResultValue { pub fn new_flat( result: Vec, read_cols: ParquetReadColumns, + json_target_types: JsonTargetTypes, ) -> SelectorResultValue { SelectorResultValue { result: SelectorResult::Flat(result), read_cols, + json_target_types, } } /// Returns memory used by the value (estimated). fn estimated_size(&self) -> usize { - match &self.result { + let result_size: usize = match &self.result { SelectorResult::PrimaryKey(batches) => { batches.iter().map(|batch| batch.memory_size()).sum() } SelectorResult::Flat(batches) => batches.iter().map(record_batch_estimated_size).sum(), - } + }; + result_size + + self.json_target_types.len() + * (mem::size_of::() + mem::size_of::()) } } @@ -2752,7 +2764,7 @@ mod tests { region_id: RegionId::new(1, 1), row_groups: vec![(FileId::random(), 0)], scan: ScanRequestFingerprintBuilder { - read_columns: ReadColumns::from_deduped_column_ids(std::iter::empty()), + read_columns: ReadColumns::new(std::iter::empty()), read_column_types: vec![], filters: vec!["tag_0 = 1".to_string()], time_filters: vec![], diff --git a/src/mito2/src/compaction/reader.rs b/src/mito2/src/compaction/reader.rs index 89b70f89c0e..f862d11e1a9 100644 --- a/src/mito2/src/compaction/reader.rs +++ b/src/mito2/src/compaction/reader.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; use common_time::Timestamp; @@ -118,15 +118,28 @@ impl CompactionSstReaderBuilder<'_> { }; let projection = (0..self.metadata.column_metadatas.len()).collect(); - let read_columns = ReadColumns::from_deduped_column_ids( - self.metadata.column_metadatas.iter().map(|x| x.column_id), - ); - let mapper = FlatProjectionMapper::new_with_read_columns( - &self.metadata, - projection, - read_columns, - json_type_hint.as_ref(), - )?; + let read_column_ids = self + .metadata + .column_metadatas + .iter() + .map(|x| x.column_id) + .collect::>(); + let json_target_types = json_type_hint + .as_ref() + .map(|hint| { + hint.iter() + .filter_map(|(col_name, json_type)| { + self.metadata + .column_by_name(col_name) + .map(|col| (col.column_id, json_type.clone())) + }) + .collect::>() + }) + .unwrap_or_default(); + let read_columns = + ReadColumns::new(read_column_ids).with_json_target_types(json_target_types); + let mapper = + FlatProjectionMapper::new_with_read_columns(&self.metadata, projection, read_columns)?; let mut scan_input = ScanInput::new(self.sst_layer, mapper) .with_files(self.inputs.to_vec()) diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index b9eed37cb30..4baa707f4e7 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -38,7 +38,7 @@ use store_api::storage::{RegionId, ScanRequest, TimeSeriesDistribution}; use crate::config::MitoConfig; use crate::error::Error; -use crate::read::read_columns::{ReadColumn, ReadColumns}; +use crate::read::read_columns::ReadColumns; use crate::read::scan_region::Scanner; use crate::test_util; use crate::test_util::{CreateRequestBuilder, TestEnv}; @@ -108,18 +108,19 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul }; assert_eq!( seq_scan.input().read_cols, - ReadColumns::from_deduped_column_ids([1, 0]) + ReadColumns::new([1, 0]) + .with_json_target_types(BTreeMap::from([(1, JsonNativeType::Variant)])) ); let stream = scanner.scan().await?; let batches = RecordBatches::try_collect(stream).await?; let expected = r#" -+------------------------------------------+-------+ -| field_0 | tag_0 | -+------------------------------------------+-------+ -| {a: {x: 10, y: ignored-a}, b: ignored-b} | tag-1 | -| {a: {x: 20, y: ignored-c}, b: ignored-d} | tag-2 | -+------------------------------------------+-------+ ++------------------------------------------------+-------+ +| field_0 | tag_0 | ++------------------------------------------------+-------+ +| {"a":{"x":10,"y":"ignored-a"},"b":"ignored-b"} | tag-1 | +| {"a":{"x":20,"y":"ignored-c"},"b":"ignored-d"} | tag-2 | ++------------------------------------------------+-------+ "#; assert_eq!(batches.pretty_print()?, expected.trim()); @@ -149,19 +150,8 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul // whole JSON2 struct. tag_0 is still read as a normal root column. assert_eq!( seq_scan.input().read_cols, - ReadColumns { - cols: vec![ - ReadColumn::new( - 1, - vec![vec![ - "field_0".to_string(), - "a".to_string(), - "x".to_string() - ]] - ), - ReadColumn::new(0, vec![]), - ] - } + ReadColumns::new([1, 0]) + .with_json_target_types(BTreeMap::from([(1, json_type_hint["field_0"].clone())])) ); // The scanner should still return a valid RecordBatch in the requested logical projection. @@ -189,10 +179,7 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul let Scanner::Seq(seq_scan) = &scanner else { unreachable!(); }; - assert_eq!( - seq_scan.input().read_cols, - ReadColumns::from_deduped_column_ids([0]) - ); + assert_eq!(seq_scan.input().read_cols, ReadColumns::new([0])); let stream = scanner.scan().await?; let batches = RecordBatches::try_collect(stream).await?; @@ -205,6 +192,19 @@ async fn test_json_type_hint_pushdown_scanner_returns_batches() -> WhateverResul +-------+ "#; assert_eq!(batches.pretty_print()?, expected.trim()); + + // JSON type hints are derived from the whole query plan, so they may include + // columns unrelated to this scan. Mito2 should ignore hints for unknown + // columns (`other_json`) instead of failing. + let request = ScanRequest { + projection: Some(vec![0]), + json_type_hint: HashMap::from([("other_json".to_string(), JsonNativeType::String)]), + ..Default::default() + }; + let scanner = engine.scanner(region_id, request).await?; + let stream = scanner.scan().await?; + let batches = RecordBatches::try_collect(stream).await?; + assert_eq!(batches.pretty_print()?, expected.trim()); Ok(()) } diff --git a/src/mito2/src/memtable/bulk/context.rs b/src/mito2/src/memtable/bulk/context.rs index af38ef66c14..994d3a027d4 100644 --- a/src/mito2/src/memtable/bulk/context.rs +++ b/src/mito2/src/memtable/bulk/context.rs @@ -72,9 +72,9 @@ impl BulkIterContext { let codec = build_primary_key_codec(®ion_metadata); let read_cols = if let Some(col_ids) = projection { - ReadColumns::from_deduped_column_ids(col_ids.iter().copied()) + ReadColumns::new(col_ids.iter().copied()) } else { - ReadColumns::from_deduped_column_ids( + ReadColumns::new( region_metadata .column_metadatas .iter() diff --git a/src/mito2/src/read/compat.rs b/src/mito2/src/read/compat.rs index af0bdc58c53..3fa625694c6 100644 --- a/src/mito2/src/read/compat.rs +++ b/src/mito2/src/read/compat.rs @@ -25,11 +25,9 @@ use datatypes::arrow::compute::{TakeOptions, take}; use datatypes::arrow::datatypes::{FieldRef, Schema, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::data_type::ConcreteDataType; -use datatypes::extension::json::is_json2_extension_type; use datatypes::prelude::DataType; use datatypes::value::Value; use datatypes::vectors::VectorRef; -use datatypes::vectors::json::array::JsonArray; use mito_codec::row_converter::{ CompositeValues, PrimaryKeyCodec, SortField, build_primary_key_codec, build_primary_key_codec_with_fields, @@ -40,8 +38,8 @@ use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::storage::ColumnId; use crate::error::{ - CompatReaderSnafu, ComputeArrowSnafu, ConvertValueSnafu, CreateDefaultSnafu, DecodeSnafu, - EncodeSnafu, NewRecordBatchSnafu, Result, UnexpectedSnafu, UnsupportedOperationSnafu, + CompatReaderSnafu, ComputeArrowSnafu, CreateDefaultSnafu, DecodeSnafu, EncodeSnafu, + NewRecordBatchSnafu, Result, UnexpectedSnafu, UnsupportedOperationSnafu, }; use crate::read::flat_projection::{FlatProjectionMapper, flat_projected_columns}; use crate::sst::parquet::flat_format::{FlatReadFormat, primary_key_column_index}; @@ -99,20 +97,12 @@ impl FlatCompatBatch { let actual = read_format.metadata(); let format_projection = read_format.format_projection(); let mut actual_schema = flat_projected_columns(actual, format_projection); - if read_format - .arrow_schema() - .fields() - .iter() - .any(is_json2_extension_type) - { - for field in read_format.arrow_schema().fields() { - if is_json2_extension_type(field) - && let Some(column_id) = - actual.column_by_name(field.name()).map(|x| x.column_id) - && let Some(i) = actual_schema.iter().position(|x| x.0 == column_id) - { - actual_schema[i].1 = ConcreteDataType::from_arrow_type(field.data_type()); - } + for (column_id, target_type) in read_format.json_target_types().iter() { + if let Some(i) = actual_schema + .iter() + .position(|(actual_column_id, _)| actual_column_id == column_id) + { + actual_schema[i].1 = ConcreteDataType::json2(target_type.clone()); } } @@ -270,16 +260,9 @@ impl FlatCompatBatch { let old_column = batch.column(*pos); if let Some(ty) = cast_type { - let casted = if let Some(json_type) = ty.as_json() - && json_type.is_json2() - { - JsonArray::from(old_column) - .project_to(&json_type.as_arrow_type()) - .context(ConvertValueSnafu)? - } else { + let casted = datatypes::arrow::compute::cast(old_column, &ty.as_arrow_type()) - .context(ComputeArrowSnafu)? - }; + .context(ComputeArrowSnafu)?; Ok(casted) } else { Ok(old_column.clone()) @@ -777,7 +760,7 @@ mod tests { let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap(); let read_format = FlatReadFormat::new( actual_metadata.clone(), - ReadColumns::from_deduped_column_ids([0, 1, 2, 3]), + ReadColumns::new([0, 1, 2, 3]), None, "test", false, @@ -864,13 +847,12 @@ mod tests { let mapper = FlatProjectionMapper::new_with_read_columns( &expected_metadata, vec![1, 2], - ReadColumns::from_deduped_column_ids([1, 2, 3]), - None, + ReadColumns::new([1, 2, 3]), ) .unwrap(); let read_format = FlatReadFormat::new( actual_metadata.clone(), - ReadColumns::from_deduped_column_ids([1, 2, 3]), + ReadColumns::new([1, 2, 3]), None, "test", false, @@ -959,7 +941,7 @@ mod tests { let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap(); let read_format = FlatReadFormat::new( actual_metadata.clone(), - ReadColumns::from_deduped_column_ids([0, 1, 2, 3]), + ReadColumns::new([0, 1, 2, 3]), None, "test", false, @@ -1041,7 +1023,7 @@ mod tests { let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap(); let read_format = FlatReadFormat::new( actual_metadata, - ReadColumns::from_deduped_column_ids([0, 1, 2]), + ReadColumns::new([0, 1, 2]), None, "test", false, @@ -1109,7 +1091,7 @@ mod tests { let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap(); let read_format = FlatReadFormat::new( actual_metadata.clone(), - ReadColumns::from_deduped_column_ids([0, 2, 3]), + ReadColumns::new([0, 2, 3]), None, "test", true, diff --git a/src/mito2/src/read/flat_projection.rs b/src/mito2/src/read/flat_projection.rs index fbd6088f670..e631cf45db4 100644 --- a/src/mito2/src/read/flat_projection.rs +++ b/src/mito2/src/read/flat_projection.rs @@ -14,7 +14,6 @@ //! Utilities for projection on flat format. -use std::collections::HashMap; use std::sync::Arc; use api::v1::SemanticType; @@ -27,7 +26,7 @@ use datatypes::arrow::array::Array; use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field}; use datatypes::extension::json::is_json2_extension_type; use datatypes::prelude::{ConcreteDataType, DataType}; -use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; +use datatypes::schema::{Schema, SchemaRef}; use datatypes::types::JsonType; use datatypes::types::json_type::JsonNativeType; use datatypes::value::Value; @@ -86,8 +85,8 @@ impl FlatProjectionMapper { ) -> Result { let projection: Vec<_> = projection.into_iter().collect(); let read_column_ids = read_column_ids_from_projection(metadata, &projection)?; - let read_cols = ReadColumns::from_deduped_column_ids(read_column_ids); - Self::new_with_read_columns(metadata, projection, read_cols, None) + let read_cols = ReadColumns::new(read_column_ids); + Self::new_with_read_columns(metadata, projection, read_cols) } /// Returns a new mapper with output projection and explicit read columns. @@ -95,7 +94,6 @@ impl FlatProjectionMapper { metadata: &RegionMetadataRef, projection: Vec, read_cols: ReadColumns, - json_type_hint: Option<&HashMap>, ) -> Result { // If the original projection is empty. let is_empty_projection = projection.is_empty(); @@ -115,7 +113,11 @@ impl FlatProjectionMapper { output_col_ids.push(col.column_id); let mut schema = col.column_schema.clone(); - maybe_concretize_json2_datatype(&mut schema, json_type_hint); + if let Some(data_type) = + json2_read_datatype(col.column_id, &schema.data_type, &read_cols) + { + schema.data_type = data_type; + } col_schemas.push(schema); } @@ -124,6 +126,7 @@ impl FlatProjectionMapper { // TODO(yingwen): Support different flat schema options. let format_projection = FormatProjection::compute_format_projection( + metadata, &id_to_index, // All columns with internal columns. metadata.column_metadatas.len() + 3, @@ -133,21 +136,8 @@ impl FlatProjectionMapper { let mut batch_schema = flat_projected_columns(metadata, &format_projection); for (column_id, data_type) in batch_schema.iter_mut() { - if let Some(json_type) = data_type.as_json() - && json_type.is_json2() - { - if let Some(concretized) = metadata - .column_by_id(*column_id) - .and_then(|metadata| { - json_type_hint.and_then(|x| x.get(&metadata.column_schema.name).cloned()) - }) - .map(ConcreteDataType::json2) - { - *data_type = concretized; - } else if is_empty_json2_type(json_type) { - // see `merge_scan::maybe_amend_json2_field` - *data_type = ConcreteDataType::json2(JsonNativeType::Variant); - } + if let Some(updated) = json2_read_datatype(*column_id, data_type, &read_cols) { + *data_type = updated; } } @@ -401,24 +391,25 @@ impl FlatProjectionMapper { } } -fn maybe_concretize_json2_datatype( - schema: &mut ColumnSchema, - json_type_hint: Option<&HashMap>, -) { - if let Some(json_type) = schema.data_type.as_json() - && json_type.is_json2() - { - if let Some(concretized) = json_type_hint - .and_then(|x| x.get(&schema.name)) - .cloned() - .map(ConcreteDataType::json2) - { - schema.data_type = concretized; - } else if is_empty_json2_type(json_type) { - // see `merge_scan::maybe_amend_json2_field` - schema.data_type = ConcreteDataType::json2(JsonNativeType::Variant); - } +fn json2_read_datatype( + column_id: ColumnId, + data_type: &ConcreteDataType, + read_cols: &ReadColumns, +) -> Option { + let json_type = data_type.as_json()?; + if !json_type.is_json2() { + return None; } + + if let Some(concretized) = read_cols.json_target_type(column_id).cloned() { + return Some(ConcreteDataType::json2(concretized)); + } + + if is_empty_json2_type(json_type) { + return Some(ConcreteDataType::json2(JsonNativeType::Variant)); + } + + None } fn is_empty_json2_type(json_type: &JsonType) -> bool { @@ -547,9 +538,8 @@ impl CompactionProjectionMapper { .collect::>(); let read_col_ids = metadata.column_metadatas.iter().map(|col| col.column_id); - let read_cols = ReadColumns::from_deduped_column_ids(read_col_ids); - let mapper = - FlatProjectionMapper::new_with_read_columns(metadata, projection, read_cols, None)?; + let read_cols = ReadColumns::new(read_col_ids); + let mapper = FlatProjectionMapper::new_with_read_columns(metadata, projection, read_cols)?; let assembler = DfBatchAssembler::new(mapper.output_schema()); Ok(Self { mapper, assembler }) @@ -617,7 +607,6 @@ impl DfBatchAssembler { #[cfg(test)] mod tests { use datatypes::schema::ColumnSchema; - use datatypes::types::json_type::JsonObjectType; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; use store_api::storage::RegionId; @@ -687,18 +676,10 @@ mod tests { #[test] fn test_json_type_hint_does_not_concretize_legacy_json() { let metadata = metadata_with_legacy_json(); - let hint = HashMap::from([( - "j".to_string(), - JsonNativeType::Object(JsonObjectType::from([( - "a".to_string(), - JsonNativeType::i64(), - )])), - )]); let mapper = FlatProjectionMapper::new_with_read_columns( &metadata, vec![0, 1], - ReadColumns::from_deduped_column_ids([0, 1]), - Some(&hint), + ReadColumns::new([0, 1]), ) .unwrap(); diff --git a/src/mito2/src/read/last_row.rs b/src/mito2/src/read/last_row.rs index 6e1a62d806e..c2316525f70 100644 --- a/src/mito2/src/read/last_row.rs +++ b/src/mito2/src/read/last_row.rs @@ -29,6 +29,7 @@ use crate::cache::{ selector_result_cache_hit, selector_result_cache_miss, }; use crate::error::{ComputeArrowSnafu, Result}; +use crate::read::read_columns::JsonTargetTypes; use crate::read::{ Batch, BatchReader, BoxedBatchReader, BoxedRecordBatchStream, timestamp_array_to_i64_slice, }; @@ -137,6 +138,7 @@ impl FlatRowGroupLastRowCachedReader { row_group_idx: usize, cache_strategy: CacheStrategy, read_cols: &ParquetReadColumns, + json_target_types: JsonTargetTypes, reader: FlatRowGroupReader, ) -> Self { let key = SelectorResultKey { @@ -148,13 +150,14 @@ impl FlatRowGroupLastRowCachedReader { if let Some(value) = cache_strategy.get_selector_result(&key) { let is_flat = matches!(&value.result, SelectorResult::Flat(_)); let schema_matches = value.read_cols == *read_cols; - if is_flat && schema_matches { + let json_target_types_matches = value.json_target_types == json_target_types; + if is_flat && schema_matches && json_target_types_matches { Self::new_hit(value) } else { - Self::new_miss(key, read_cols, reader, cache_strategy) + Self::new_miss(key, read_cols, json_target_types, reader, cache_strategy) } } else { - Self::new_miss(key, read_cols, reader, cache_strategy) + Self::new_miss(key, read_cols, json_target_types, reader, cache_strategy) } } @@ -174,6 +177,7 @@ impl FlatRowGroupLastRowCachedReader { fn new_miss( key: SelectorResultKey, read_cols: &ParquetReadColumns, + json_target_types: JsonTargetTypes, reader: FlatRowGroupReader, cache_strategy: CacheStrategy, ) -> Self { @@ -181,6 +185,7 @@ impl FlatRowGroupLastRowCachedReader { Self::Miss(FlatRowGroupLastRowReader::new( key, read_cols.clone(), + json_target_types, reader, cache_strategy, )) @@ -260,6 +265,7 @@ pub(crate) struct FlatRowGroupLastRowReader { yielded_batches: Vec, cache_strategy: CacheStrategy, read_cols: ParquetReadColumns, + json_target_types: JsonTargetTypes, /// Accumulates small selector-output batches before concatenating. pending: BatchBuffer, } @@ -268,6 +274,7 @@ impl FlatRowGroupLastRowReader { fn new( key: SelectorResultKey, read_cols: ParquetReadColumns, + json_target_types: JsonTargetTypes, reader: FlatRowGroupReader, cache_strategy: CacheStrategy, ) -> Self { @@ -278,6 +285,7 @@ impl FlatRowGroupLastRowReader { yielded_batches: vec![], cache_strategy, read_cols, + json_target_types, pending: BatchBuffer::new(), } } @@ -326,6 +334,7 @@ impl FlatRowGroupLastRowReader { let value = Arc::new(SelectorResultValue::new_flat( batches, self.read_cols.clone(), + self.json_target_types.clone(), )); self.cache_strategy.put_selector_result(self.key, value); } diff --git a/src/mito2/src/read/projection.rs b/src/mito2/src/read/projection.rs index 27fe76345b6..c6d4baee3ac 100644 --- a/src/mito2/src/read/projection.rs +++ b/src/mito2/src/read/projection.rs @@ -295,8 +295,7 @@ mod tests { let mapper = FlatProjectionMapper::new_with_read_columns( &metadata, vec![4, 1], - ReadColumns::from_deduped_column_ids([4, 1, 3]), - None, + ReadColumns::new([4, 1, 3]), ) .unwrap(); assert_eq!(&[4, 1, 3], mapper.read_columns().column_ids().as_slice()); diff --git a/src/mito2/src/read/range_cache.rs b/src/mito2/src/read/range_cache.rs index 036a574c1f5..ee0e421f86e 100644 --- a/src/mito2/src/read/range_cache.rs +++ b/src/mito2/src/read/range_cache.rs @@ -852,7 +852,7 @@ pub fn bench_cache_flat_range_stream( let cache_strategy = CacheStrategy::EnableAll(cache_manager); let fingerprint = ScanRequestFingerprintBuilder { - read_columns: ReadColumns::from_deduped_column_ids(std::iter::empty()), + read_columns: ReadColumns::new(std::iter::empty()), read_column_types: vec![], filters: vec![], time_filters: vec![], @@ -916,7 +916,7 @@ mod tests { filter_deleted: bool, partition_expr_version: u64, ) -> ScanRequestFingerprint { - let read_columns = ReadColumns::from_deduped_column_ids([1, 2]); + let read_columns = ReadColumns::new([1, 2]); ScanRequestFingerprintBuilder { read_columns, read_column_types: vec![None, None], diff --git a/src/mito2/src/read/read_columns.rs b/src/mito2/src/read/read_columns.rs index 21699b8c439..42a9f71ddcb 100644 --- a/src/mito2/src/read/read_columns.rs +++ b/src/mito2/src/read/read_columns.rs @@ -12,119 +12,72 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::BTreeMap; use std::mem; +use std::sync::Arc; -use store_api::storage::{ColumnId, NestedPath}; +use datatypes::types::json_type::JsonNativeType; +use store_api::storage::ColumnId; + +pub(crate) type JsonTargetTypes = Arc>; /// Logical columns to read from a region. /// -/// Read columns describe which logical columns and nested fields should be read -/// from storage. Each read column is identified by its [`ColumnId`], -/// which represents the root column in the storage schema. -/// -/// Nested fields under the column are specified by [`NestedPath`] entries. -/// Each path includes the root column name as its first element. -/// -/// For example, assume column id `9` corresponds to a root column named `j` -/// with nested fields: -/// -/// ```text -/// j -/// ├── a -/// └── b -/// └── c -/// ``` -/// -/// The following SQL: -/// -/// SELECT j.a, j.b.c FROM t -/// -/// may produce read columns like: -/// -/// ```text -/// ReadColumn { -/// column_id: 9, -/// nested_paths: [ -/// ["j", "a"], -/// ["j", "b", "c"], -/// ] -/// } -/// ``` -/// -/// If `nested_paths` is empty, the whole column will be read. +/// Read columns describe which logical root columns should be read from storage. +/// JSON2 columns can carry query-time target types that are later translated to +/// physical nested parquet paths by the parquet reader. #[derive(Debug, Default, Clone, PartialEq, Eq, Hash)] pub struct ReadColumns { - pub cols: Vec, + pub col_ids: Vec, + json_target_types: JsonTargetTypes, } impl ReadColumns { - pub fn from_deduped_column_ids(column_ids: I) -> Self + /// Creates read columns from logical column ids. + /// + /// This preserves the input order and duplicate ids. + pub fn new(col_ids: I) -> Self where I: IntoIterator, { - let cols = column_ids - .into_iter() - .map(|col_id| ReadColumn::new(col_id, vec![])) - .collect(); - ReadColumns { cols } + Self { + col_ids: col_ids.into_iter().collect(), + json_target_types: Arc::default(), + } + } + + pub fn with_json_target_types( + mut self, + json_target_types: BTreeMap, + ) -> Self { + self.json_target_types = Arc::new(json_target_types); + self } pub fn is_empty(&self) -> bool { - self.cols.is_empty() + self.col_ids.is_empty() } pub fn column_ids_iter(&self) -> impl Iterator + '_ { - self.cols.iter().map(|column| column.column_id) + self.col_ids.iter().copied() } pub fn column_ids(&self) -> Vec { self.column_ids_iter().collect() } - pub fn columns(&self) -> &[ReadColumn] { - &self.cols + pub fn json_target_types(&self) -> &JsonTargetTypes { + &self.json_target_types + } + + pub fn json_target_type(&self, column_id: ColumnId) -> Option<&JsonNativeType> { + self.json_target_types.get(&column_id) } pub fn estimated_size(&self) -> usize { - self.cols.capacity() * mem::size_of::() - + self - .cols - .iter() - .map(ReadColumn::estimated_size) - .sum::() - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct ReadColumn { - pub column_id: ColumnId, - /// Nested field paths under this column. - /// Empty means reading the whole column. - pub nested_paths: Vec, -} - -impl ReadColumn { - pub fn new(column_id: ColumnId, nested_paths: Vec) -> Self { - Self { - column_id, - nested_paths, - } - } - - pub fn nested_paths(&self) -> &[NestedPath] { - &self.nested_paths - } - - pub fn estimated_size(&self) -> usize { - mem::size_of::() - + self.nested_paths.capacity() * mem::size_of::() - + self - .nested_paths - .iter() - .map(|path| { - path.capacity() * mem::size_of::() - + path.iter().map(|node| node.capacity()).sum::() - }) - .sum::() + self.col_ids.capacity() * mem::size_of::() + + self.col_ids.len() * mem::size_of::() + + self.json_target_types.len() + * (mem::size_of::() + mem::size_of::()) } } diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 250f0d7fc27..c0340464ecc 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -14,7 +14,7 @@ //! Scans a region according to the scan request. -use std::collections::HashSet; +use std::collections::{BTreeMap, HashSet}; use std::fmt; use std::num::NonZeroU64; use std::sync::Arc; @@ -35,18 +35,18 @@ use datafusion_expr::Expr; use datafusion_expr::utils::expr_to_columns; use datatypes::arrow::array::{ArrayRef, BooleanArray, UInt64Array}; use datatypes::extension::json::is_json2_extension_type; +use datatypes::prelude::ConcreteDataType; use datatypes::types::json_type::JsonNativeType; use datatypes::value::timestamp_to_scalar_value; use futures::StreamExt; -use itertools::Itertools; use partition::expr::PartitionExpr; use smallvec::SmallVec; -use snafu::{OptionExt, ResultExt}; +use snafu::{OptionExt, ResultExt, ensure}; use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::region_engine::{PartitionRange, RegionScannerRef}; use store_api::storage::{ - ColumnId, NestedPath, RegionId, ScanRequest, SequenceNumber, SequenceRange, - TimeSeriesDistribution, TimeSeriesRowSelector, + ColumnId, RegionId, ScanRequest, SequenceNumber, SequenceRange, TimeSeriesDistribution, + TimeSeriesRowSelector, }; use table::predicate::{Predicate, build_time_range_predicate, extract_time_range_from_expr}; use tokio::sync::{Semaphore, mpsc}; @@ -64,7 +64,7 @@ use crate::read::compat::{self, FlatCompatBatch}; use crate::read::flat_projection::FlatProjectionMapper; use crate::read::range::{FileRangeBuilder, MemRangeBuilder, RangeMeta, RowGroupIndex}; use crate::read::range_cache::{ScanRequestFingerprint, implied_time_range_from_exprs}; -use crate::read::read_columns::{ReadColumn, ReadColumns}; +use crate::read::read_columns::ReadColumns; use crate::read::seq_scan::SeqScan; use crate::read::series_scan::SeriesScan; use crate::read::stream::ScanBatchStream; @@ -420,23 +420,7 @@ impl ScanRegion { let read_col_ids = self.build_read_col_ids(self.request.projection.as_deref(), &predicate)?; - - // Narrow JSON2 columns to avoid reading unnecessary nested fields. - // - // `read_col_ids` selects the root columns required by the projection and predicates, - // while nested projection is currently only applied to JSON2 columns, whose type hints - // further narrow them to the requested nested fields. - let has_structured_json = metadata - .schema - .arrow_schema() - .fields() - .iter() - .any(is_json2_extension_type); - let read_cols = if has_structured_json { - self.read_columns_with_json_type_hint(&read_col_ids) - } else { - ReadColumns::from_deduped_column_ids(read_col_ids.iter().copied()) - }; + let read_cols = self.build_read_columns(&read_col_ids)?; // The mapper always computes projected column ids as the schema of SSTs may change. let projection = self @@ -444,23 +428,7 @@ impl ScanRegion { .projection .clone() .unwrap_or_else(|| (0..metadata.column_metadatas.len()).collect()); - let json_type_hint = has_structured_json - .then_some(&self.request.json_type_hint) - .inspect(|json_type_hint| { - debug!( - "Concretized JSON type: {{{}}}", - json_type_hint - .iter() - .map(|(k, v)| format!("{}: {}", k, v)) - .join(", ") - ); - }); - let mapper = FlatProjectionMapper::new_with_read_columns( - metadata, - projection, - read_cols, - json_type_hint, - )?; + let mapper = FlatProjectionMapper::new_with_read_columns(metadata, projection, read_cols)?; let mapper = if self.request.preserve_pk_dictionary_encoding { mapper.with_pk_dictionary_encoding() } else { @@ -609,7 +577,8 @@ impl ScanRegion { Ok(input) } - /// Builds the ordered root column ids required by the pushed-down projection and predicate. + /// Builds the deduplicated root column ids required by the projection and + /// predicate. fn build_read_col_ids( &self, projection: Option<&[usize]>, @@ -688,27 +657,50 @@ impl ScanRegion { Ok(read_col_ids) } - /// Builds read columns with nested paths derived from JSON type hints. - fn read_columns_with_json_type_hint(&self, col_ids: &[ColumnId]) -> ReadColumns { - let cols = col_ids + /// Builds logical read columns and attaches JSON2 target types when needed. + /// + /// The behavior about JSON2 is as follows: + /// - JSON2 columns without hints use Variant to read the whole column. + /// - Hints targeting non-JSON2 read columns are rejected. + fn build_read_columns(&self, col_ids: &[ColumnId]) -> Result { + let metadata = &self.version.metadata; + let json_type_hint = &self.request.json_type_hint; + + let has_json2 = metadata + .schema + .arrow_schema() + .fields() .iter() - .map(|&col_id| { - let nested_paths = self - .version - .metadata - .column_by_id(col_id) - .and_then(|column| { - let col_name = &column.column_schema.name; - self.request - .json_type_hint - .get(col_name) - .map(|json_type| json_nested_paths(col_name, json_type)) - }) - .unwrap_or_default(); - ReadColumn::new(col_id, nested_paths) - }) - .collect(); - ReadColumns { cols } + .any(is_json2_extension_type); + + if !has_json2 && json_type_hint.is_empty() { + return Ok(ReadColumns::new(col_ids.iter().copied())); + } + + let mut json_target_types = BTreeMap::new(); + for &col_id in col_ids { + let Some(col) = metadata.column_by_id(col_id) else { + continue; + }; + let col_name = &col.column_schema.name; + let hint = json_type_hint.get(col_name); + if !col.column_schema.data_type.is_json2() { + ensure!( + hint.is_none(), + InvalidRequestSnafu { + region_id: metadata.region_id, + reason: format!( + "JSON type hint targets non-JSON2 column {} (id: {}, type: {})", + col_name, col_id, col.column_schema.data_type + ), + } + ); + continue; + } + let target_type = hint.cloned().unwrap_or(JsonNativeType::Variant); + json_target_types.insert(col_id, target_type); + } + Ok(ReadColumns::new(col_ids.iter().copied()).with_json_target_types(json_target_types)) } fn region_id(&self) -> RegionId { @@ -1583,30 +1575,6 @@ fn pre_filter_mode(append_mode: bool, merge_mode: MergeMode) -> PreFilterMode { } } -fn json_nested_paths(column_name: &str, json_type: &JsonNativeType) -> Vec { - let mut paths = Vec::new(); - let mut current = vec![column_name.to_string()]; - collect_json_nested_paths(json_type, &mut current, &mut paths); - paths -} - -fn collect_json_nested_paths( - json_type: &JsonNativeType, - current: &mut NestedPath, - paths: &mut Vec, -) { - match json_type { - JsonNativeType::Object(fields) if !fields.is_empty() => { - for (field, child) in fields { - current.push(field.clone()); - collect_json_nested_paths(child, current, paths); - current.pop(); - } - } - _ => paths.push(current.clone()), - } -} - /// Output of [build_scan_fingerprint]: the cache fingerprint plus the derived /// implied time range used to decide whether the cache key can drop the time /// predicates for a given partition (see `build_range_cache_key`). @@ -1703,9 +1671,15 @@ pub(crate) fn build_scan_fingerprint(input: &ScanInput) -> Option Result<()> { - let hint = JsonNativeType::Object(JsonObjectType::from([ - ("a".to_string(), JsonNativeType::i64()), - ( - "b".to_string(), - JsonNativeType::Object(JsonObjectType::from([( - "c".to_string(), - JsonNativeType::String, - )])), - ), - ])); - - fn nested_path(parts: &[&str]) -> NestedPath { - parts.iter().map(|part| part.to_string()).collect() - } - - assert_eq!( - json_nested_paths("j", &hint), - vec![nested_path(&["j", "a"]), nested_path(&["j", "b", "c"])] - ); - Ok(()) - } - #[tokio::test] async fn test_build_scan_fingerprint_for_eligible_scan() { let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); @@ -2370,6 +2321,158 @@ mod tests { assert_ne!(0, metadata.partition_expr_version); } + #[tokio::test] + async fn test_build_scan_fingerprint_uses_json_target_types() { + let mut builder = RegionMetadataBuilder::new(RegionId::new(123, 456)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "k0".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 1, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "j".to_string(), + ConcreteDataType::json2(JsonNativeType::Variant), + true, + ), + semantic_type: SemanticType::Field, + column_id: 2, + }) + .primary_key(vec![0]); + let metadata = Arc::new(builder.build().unwrap()); + + let make_input = |target_type| async { + let env = SchedulerEnv::new().await; + let read_cols = ReadColumns::new([0, 1, 2]) + .with_json_target_types(BTreeMap::from([(2, target_type)])); + let mapper = + FlatProjectionMapper::new_with_read_columns(&metadata, vec![0, 1, 2], read_cols) + .unwrap(); + let predicate = + PredicateGroup::new(metadata.as_ref(), &[col("k0").eq(lit("foo"))]).unwrap(); + let file = FileHandle::new( + FileMeta::default(), + Arc::new(crate::sst::file_purger::NoopFilePurger), + ); + ScanInput::new(env.access_layer.clone(), mapper) + .with_predicate(predicate) + .with_cache(CacheStrategy::EnableAll(Arc::new( + CacheManager::builder() + .range_result_cache_size(1024) + .build(), + ))) + .with_files(vec![file]) + }; + + let int_target = JsonNativeType::i64(); + let string_target = JsonNativeType::String; + let int_fingerprint = build_scan_fingerprint(&make_input(int_target.clone()).await) + .unwrap() + .fingerprint; + let string_fingerprint = build_scan_fingerprint(&make_input(string_target).await) + .unwrap() + .fingerprint; + + assert_ne!(int_fingerprint, string_fingerprint); + assert_eq!( + Some(&Some(ConcreteDataType::json2(int_target))), + int_fingerprint.read_column_types().get(2) + ); + } + + #[tokio::test] + async fn test_scan_input_rejects_json_type_hint_for_non_json2_column() { + let mut builder = RegionMetadataBuilder::new(RegionId::new(123, 456)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "k0".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + semantic_type: SemanticType::Tag, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 1, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "j".to_string(), + ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::from([( + "a".to_string(), + JsonNativeType::i64(), + )]))), + true, + ), + semantic_type: SemanticType::Field, + column_id: 2, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "v0".to_string(), + ConcreteDataType::int64_datatype(), + true, + ), + semantic_type: SemanticType::Field, + column_id: 3, + }) + .primary_key(vec![0]); + let metadata = Arc::new(builder.build().unwrap()); + let mutable = Arc::new(crate::memtable::time_partition::TimePartitions::new( + metadata.clone(), + Arc::new(crate::test_util::memtable_util::EmptyMemtableBuilder::default()), + 0, + None, + )); + let version = Arc::new( + crate::region::version::VersionBuilder::new(metadata.clone(), mutable).build(), + ); + let env = SchedulerEnv::new().await; + let request = ScanRequest { + projection: Some(vec![0, 1, 2, 3]), + json_type_hint: std::collections::HashMap::from([( + "v0".to_string(), + JsonNativeType::i64(), + )]), + ..Default::default() + }; + + let err = ScanRegion::new( + version, + env.access_layer.clone(), + request, + CacheStrategy::Disabled, + ) + .scan_input() + .await; + let Err(err) = err else { + panic!("scan input should reject JSON type hint for non-JSON2 column"); + }; + + assert!(err.to_string().contains("non-JSON2 column v0")); + } + #[test] fn test_update_dyn_filters_with_empty_base_predicates() { let metadata = Arc::new(metadata_with_primary_key(vec![0, 1], false)); diff --git a/src/mito2/src/sst/parquet/file_range.rs b/src/mito2/src/sst/parquet/file_range.rs index 2b9e800fb23..a3a94ef25df 100644 --- a/src/mito2/src/sst/parquet/file_range.rs +++ b/src/mito2/src/sst/parquet/file_range.rs @@ -235,6 +235,7 @@ impl FileRange { self.row_group_idx, cache_strategy, self.context.read_format().parquet_read_columns(), + self.context.read_format().json_target_types().clone(), flat_row_group_reader, ); FlatPruneReader::new_with_last_row_reader(self.context.clone(), reader, skip_fields) @@ -827,9 +828,7 @@ mod tests { let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, @@ -895,9 +894,7 @@ mod tests { let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, diff --git a/src/mito2/src/sst/parquet/flat_format.rs b/src/mito2/src/sst/parquet/flat_format.rs index 237b7c0f2c6..e381ea76f49 100644 --- a/src/mito2/src/sst/parquet/flat_format.rs +++ b/src/mito2/src/sst/parquet/flat_format.rs @@ -51,12 +51,11 @@ use crate::error::{ ComputeArrowSnafu, DecodeSnafu, InvalidParquetSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, Result, }; -use crate::read::read_columns::ReadColumns; +use crate::read::read_columns::{JsonTargetTypes, ReadColumns}; use crate::sst::parquet::format::{ FIXED_POS_COLUMN_NUM, FormatProjection, INTERNAL_COLUMN_NUM, PrimaryKeyArray, PrimaryKeyReadFormat, StatValues, column_null_counts, column_values, }; -use crate::sst::parquet::json_align::align_schema_by_nested_paths; use crate::sst::parquet::read_columns::ParquetReadColumns; use crate::sst::{ FlatSchemaOptions, flat_sst_arrow_schema_column_num, tag_maybe_to_dictionary_field, @@ -181,6 +180,8 @@ pub(crate) fn field_column_start(metadata: &RegionMetadata, num_columns: usize) pub struct FlatReadFormat { /// Sequence number to override the sequence read from the SST. override_sequence: Option, + /// Logical columns requested by this read. + read_cols: ReadColumns, /// Parquet format adapter. parquet_adapter: ParquetAdapter, /// Output schema to wrap binary `__primary_key` back to a dictionary; `None` disables wrapping. @@ -210,22 +211,25 @@ impl FlatReadFormat { // Only skip auto convert when the primary key encoding is sparse. ParquetAdapter::PrimaryKeyToFlat(ParquetPrimaryKeyToFlat::new( metadata, - read_cols, + read_cols.clone(), skip_auto_convert, )) } else { ParquetAdapter::PrimaryKeyToFlat(ParquetPrimaryKeyToFlat::new( - metadata, read_cols, false, + metadata, + read_cols.clone(), + false, )) } } else { let file_schema = file_schema .unwrap_or_else(|| to_flat_sst_arrow_schema(&metadata, &Default::default())); - ParquetAdapter::Flat(ParquetFlat::new(metadata, read_cols, file_schema)) + ParquetAdapter::Flat(ParquetFlat::new(metadata, read_cols.clone(), file_schema)) }; Ok(FlatReadFormat { override_sequence: None, + read_cols, parquet_adapter, pk_dict_wrap_schema: None, }) @@ -294,22 +298,51 @@ impl FlatReadFormat { } } - /// Gets the projected output schema produced by parquet reading. + /// Gets the projected output schema expected by the scan. pub(crate) fn output_arrow_schema(&self) -> Result { - let read_columns = self.parquet_read_columns(); - let projection = read_columns.root_indices(); + let projection = self.parquet_read_columns().root_indices(); let mut schema = self .arrow_schema() .project(projection) .context(ComputeArrowSnafu)?; - if read_columns.has_nested() { - debug_assert_eq!(schema.fields().len(), read_columns.columns().len()); - let nested_paths = read_columns.columns().iter().map(|x| x.nested_paths()); - align_schema_by_nested_paths(&mut schema, nested_paths); + let mut fields = schema.fields().iter().cloned().collect::>(); + for (column_id, target_type) in self.json_target_types().iter() { + let Some(index) = self.parquet_projected_index_by_id(*column_id) else { + continue; + }; + let Some(field) = schema.fields().get(index) else { + continue; + }; + fields[index] = Arc::new( + field + .as_ref() + .clone() + .with_data_type(ConcreteDataType::json2(target_type.clone()).as_arrow_type()), + ); } + schema.fields = fields.into(); Ok(Arc::new(schema)) } + /// Index of a column in the projected schema produced directly by parquet + /// reading, before any primary-key-to-flat conversion. + fn parquet_projected_index_by_id(&self, column_id: ColumnId) -> Option { + match &self.parquet_adapter { + ParquetAdapter::Flat(p) => p + .format_projection + .column_id_to_projected_index + .get(&column_id) + .copied(), + // `format_projection` addresses the post-conversion flat batch here. + // This helper needs the raw primary-key projection used by parquet reading. + ParquetAdapter::PrimaryKeyToFlat(p) => p + .format + .field_id_to_projected_index() + .get(&column_id) + .copied(), + } + } + /// Gets the metadata of the SST. pub(crate) fn metadata(&self) -> &RegionMetadataRef { match &self.parquet_adapter { @@ -326,6 +359,11 @@ impl FlatReadFormat { } } + /// Gets JSON2 target types keyed by column id. + pub(crate) fn json_target_types(&self) -> &JsonTargetTypes { + self.read_cols.json_target_types() + } + /// Gets the projection in the flat format. /// /// When `skip_auto_convert` is enabled (primary-key format read), this returns the @@ -491,6 +529,7 @@ impl ParquetPrimaryKeyToFlat { } else { // Computes the format projection for the new format. let format_projection = FormatProjection::compute_format_projection( + &metadata, &id_to_index, sst_column_num, read_cols.clone(), @@ -540,8 +579,12 @@ impl ParquetFlat { let id_to_index = sst_column_id_indices(&metadata); let sst_column_num = flat_sst_arrow_schema_column_num(&metadata, &FlatSchemaOptions::default()); - let format_projection = - FormatProjection::compute_format_projection(&id_to_index, sst_column_num, read_cols); + let format_projection = FormatProjection::compute_format_projection( + &metadata, + &id_to_index, + sst_column_num, + read_cols, + ); Self { metadata, @@ -841,9 +884,7 @@ impl FlatReadFormat { pub fn new_with_all_columns(metadata: RegionMetadataRef) -> FlatReadFormat { Self::new( Arc::clone(&metadata), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", false, @@ -965,7 +1006,7 @@ mod tests { .collect(); let mut read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids(column_ids), + ReadColumns::new(column_ids), None, "test", false, @@ -1045,7 +1086,7 @@ mod tests { let metadata = Arc::new(build_metadata(1, 2, PrimaryKeyEncoding::Dense)); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids([0_u32, 2_u32]), + ReadColumns::new([0_u32, 2_u32]), None, "test", false, diff --git a/src/mito2/src/sst/parquet/format.rs b/src/mito2/src/sst/parquet/format.rs index a6338c1a1c9..f3f0a34a682 100644 --- a/src/mito2/src/sst/parquet/format.rs +++ b/src/mito2/src/sst/parquet/format.rs @@ -39,6 +39,7 @@ use datatypes::arrow::array::{ use datatypes::arrow::datatypes::{SchemaRef, UInt32Type}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::DataType; +use datatypes::types::json_type::JsonNativeType; use datatypes::vectors::Helper; use mito_codec::row_converter::{ CompositeValues, PrimaryKeyCodec, SortField, build_primary_key_codec_with_fields, @@ -47,12 +48,12 @@ use parquet::file::metadata::{ParquetMetaData, RowGroupMetaData}; use parquet::file::statistics::Statistics; use snafu::{OptionExt, ResultExt, ensure}; use store_api::metadata::{ColumnMetadata, RegionMetadataRef}; -use store_api::storage::{ColumnId, SequenceNumber}; +use store_api::storage::{ColumnId, NestedPath, SequenceNumber}; use crate::error::{ ConvertVectorSnafu, DecodeSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, Result, }; -use crate::read::read_columns::ReadColumns; +use crate::read::read_columns::{JsonTargetTypes, ReadColumns}; use crate::read::{Batch, BatchBuilder, BatchColumn}; use crate::sst::file::{FileMeta, FileTimeRange}; use crate::sst::parquet::read_columns::{ParquetReadColumn, ParquetReadColumns}; @@ -233,6 +234,7 @@ impl PrimaryKeyReadFormat { let arrow_schema = to_sst_arrow_schema(&metadata); let format_projection = FormatProjection::compute_format_projection( + &metadata, &field_id_to_index, arrow_schema.fields.len(), read_cols, @@ -592,18 +594,21 @@ impl FormatProjection { /// /// `id_to_index` is a mapping from column id to the index of the column in the SST. pub(crate) fn compute_format_projection( + metadata: &RegionMetadataRef, id_to_index: &HashMap, sst_column_num: usize, cols: ReadColumns, ) -> Self { + let json_target_types = cols.json_target_types().clone(); let mut projected_columns: Vec<_> = cols - .cols + .col_ids .into_iter() - .filter_map(|col| { - id_to_index - .get(&col.column_id) - .copied() - .map(|index_of_sst| (col.column_id, index_of_sst, col.nested_paths)) + .filter_map(|col_id| { + id_to_index.get(&col_id).copied().map(|index_of_sst| { + let nested_paths = + json_target_nested_paths(metadata, &json_target_types, col_id); + (col_id, index_of_sst, nested_paths) + }) }) .collect(); // Sorts columns by their indices in the SST. SST uses a bitmap for projection. @@ -681,6 +686,45 @@ impl FormatProjection { } } +fn json_target_nested_paths( + metadata: &RegionMetadataRef, + json_target_types: &JsonTargetTypes, + column_id: ColumnId, +) -> Vec { + let Some(target_type) = json_target_types.get(&column_id) else { + return Vec::new(); + }; + let Some(column) = metadata.column_by_id(column_id) else { + return Vec::new(); + }; + + json_nested_paths(&column.column_schema.name, target_type) +} + +fn json_nested_paths(column_name: &str, json_type: &JsonNativeType) -> Vec { + let mut paths = Vec::new(); + let mut current = vec![column_name.to_string()]; + collect_json_nested_paths(json_type, &mut current, &mut paths); + paths +} + +fn collect_json_nested_paths( + json_type: &JsonNativeType, + current: &mut NestedPath, + paths: &mut Vec, +) { + match json_type { + JsonNativeType::Object(fields) if !fields.is_empty() => { + for (field, child) in fields { + current.push(field.clone()); + collect_json_nested_paths(child, current, paths); + current.pop(); + } + } + _ => paths.push(current.clone()), + } +} + /// Values of column statistics of the SST. /// /// It also distinguishes the case that a column is not found and @@ -710,9 +754,7 @@ impl PrimaryKeyReadFormat { pub fn new_with_all_columns(metadata: RegionMetadataRef) -> PrimaryKeyReadFormat { Self::new( Arc::clone(&metadata), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), ) } } @@ -818,6 +860,7 @@ pub(crate) fn need_override_sequence(parquet_meta: &ParquetMetaData) -> bool { #[cfg(test)] mod tests { + use std::collections::BTreeMap; use std::sync::Arc; use api::v1::OpType; @@ -840,7 +883,6 @@ mod tests { use super::*; use crate::error::InvalidMetadataSnafu; - use crate::read::read_columns::ReadColumn; use crate::sst::parquet::flat_format::{ FlatReadFormat, FlatWriteFormat, sequence_column_index, sst_column_id_indices, }; @@ -986,29 +1028,25 @@ mod tests { fn test_projection_indices() { let metadata = build_test_region_metadata(); // Only read tag1 - let read_format = - PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::from_deduped_column_ids([3])); + let read_format = PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::new([3])); assert_eq!( &[2, 3, 4, 5], read_format.parquet_read_columns().root_indices() ); // Only read field1 - let read_format = - PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::from_deduped_column_ids([4])); + let read_format = PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::new([4])); assert_eq!( &[0, 2, 3, 4, 5], read_format.parquet_read_columns().root_indices() ); // Only read ts - let read_format = - PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::from_deduped_column_ids([5])); + let read_format = PrimaryKeyReadFormat::new(metadata.clone(), ReadColumns::new([5])); assert_eq!( &[2, 3, 4, 5], read_format.parquet_read_columns().root_indices() ); // Read field0, tag0, ts - let read_format = - PrimaryKeyReadFormat::new(metadata, ReadColumns::from_deduped_column_ids([2, 1, 5])); + let read_format = PrimaryKeyReadFormat::new(metadata, ReadColumns::new([2, 1, 5])); assert_eq!( &[1, 2, 3, 4, 5], read_format.parquet_read_columns().root_indices() @@ -1049,14 +1087,16 @@ mod tests { let metadata = Arc::new(builder.build().context(InvalidMetadataSnafu)?); let column_id_to_parquet_index = sst_column_id_indices(&metadata); let projection = FormatProjection::compute_format_projection( + &metadata, &column_id_to_parquet_index, metadata.column_metadatas.len() + FIXED_POS_COLUMN_NUM, - ReadColumns { - cols: vec![ReadColumn::new( - 4, - vec![vec!["j".to_string(), "a".to_string()]], - )], - }, + ReadColumns::new([4]).with_json_target_types(BTreeMap::from([( + 4, + JsonNativeType::Object(JsonObjectType::from([( + "a".to_string(), + JsonNativeType::i64(), + )])), + )])), ); let columns = projection.parquet_read_cols.columns(); @@ -1111,8 +1151,7 @@ mod tests { .iter() .map(|col| col.column_id) .collect(); - let read_format = - PrimaryKeyReadFormat::new(metadata, ReadColumns::from_deduped_column_ids(column_ids)); + let read_format = PrimaryKeyReadFormat::new(metadata, ReadColumns::new(column_ids)); assert_eq!(arrow_schema, *read_format.arrow_schema()); let record_batch = RecordBatch::new_empty(arrow_schema); @@ -1131,8 +1170,7 @@ mod tests { .iter() .map(|col| col.column_id) .collect(); - let read_format = - PrimaryKeyReadFormat::new(metadata, ReadColumns::from_deduped_column_ids(column_ids)); + let read_format = PrimaryKeyReadFormat::new(metadata, ReadColumns::new(column_ids)); let columns: Vec = vec![ Arc::new(Int64Array::from(vec![1, 1, 10, 10])), // field1 @@ -1160,9 +1198,7 @@ mod tests { let metadata = build_test_region_metadata(); let read_format = PrimaryKeyReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), ); let columns: Vec = vec![ @@ -1330,56 +1366,36 @@ mod tests { // The projection includes all "fixed position" columns: ts(4), __primary_key(5), __sequence(6), __op_type(7) // Only read tag1 (column_id=3, index=1) + fixed columns - let read_format = FlatReadFormat::new( - metadata.clone(), - ReadColumns::from_deduped_column_ids([3]), - None, - "test", - false, - ) - .unwrap(); + let read_format = + FlatReadFormat::new(metadata.clone(), ReadColumns::new([3]), None, "test", false) + .unwrap(); assert_eq!( &[1, 4, 5, 6, 7], read_format.parquet_read_columns().root_indices() ); // Only read field1 (column_id=4, index=2) + fixed columns - let read_format = FlatReadFormat::new( - metadata.clone(), - ReadColumns::from_deduped_column_ids([4]), - None, - "test", - false, - ) - .unwrap(); + let read_format = + FlatReadFormat::new(metadata.clone(), ReadColumns::new([4]), None, "test", false) + .unwrap(); assert_eq!( &[2, 4, 5, 6, 7], read_format.parquet_read_columns().root_indices() ); // Only read ts (column_id=5, index=4) + fixed columns (ts is already included in fixed) - let read_format = FlatReadFormat::new( - metadata.clone(), - ReadColumns::from_deduped_column_ids([5]), - None, - "test", - false, - ) - .unwrap(); + let read_format = + FlatReadFormat::new(metadata.clone(), ReadColumns::new([5]), None, "test", false) + .unwrap(); assert_eq!( &[4, 5, 6, 7], read_format.parquet_read_columns().root_indices() ); // Read field0(column_id=2, index=3), tag0(column_id=1, index=0), ts(column_id=5, index=4) + fixed columns - let read_format = FlatReadFormat::new( - metadata, - ReadColumns::from_deduped_column_ids([2, 1, 5]), - None, - "test", - false, - ) - .unwrap(); + let read_format = + FlatReadFormat::new(metadata, ReadColumns::new([2, 1, 5]), None, "test", false) + .unwrap(); assert_eq!( &[0, 3, 4, 5, 6, 7], read_format.parquet_read_columns().root_indices() @@ -1391,7 +1407,7 @@ mod tests { let metadata = build_test_region_metadata(); let mut format = FlatReadFormat::new( metadata, - ReadColumns::from_deduped_column_ids(std::iter::once(1)), // Just read tag0 + ReadColumns::new(std::iter::once(1)), // Just read tag0 Some(build_test_flat_sst_schema()), "test", false, @@ -1608,7 +1624,7 @@ mod tests { .collect(); let format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids(column_ids), + ReadColumns::new(column_ids), Some(build_test_arrow_schema()), "test", false, @@ -1674,7 +1690,7 @@ mod tests { .collect(); let format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids(column_ids.clone()), + ReadColumns::new(column_ids.clone()), None, "test", false, @@ -1744,7 +1760,7 @@ mod tests { let format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids(column_ids), + ReadColumns::new(column_ids), None, "test", true, diff --git a/src/mito2/src/sst/parquet/json_align/mod.rs b/src/mito2/src/sst/parquet/json_align/mod.rs index 89b993d3d09..54c97bdbb89 100644 --- a/src/mito2/src/sst/parquet/json_align/mod.rs +++ b/src/mito2/src/sst/parquet/json_align/mod.rs @@ -15,10 +15,8 @@ use datatypes::arrow::record_batch::RecordBatch; use futures::stream::BoxStream; -mod schema; mod stream; -pub(crate) use schema::align_schema_by_nested_paths; pub(crate) use stream::NestedSchemaAligner; use crate::error::Result; diff --git a/src/mito2/src/sst/parquet/json_align/schema.rs b/src/mito2/src/sst/parquet/json_align/schema.rs deleted file mode 100644 index 86bd4d4d0d7..00000000000 --- a/src/mito2/src/sst/parquet/json_align/schema.rs +++ /dev/null @@ -1,250 +0,0 @@ -// Copyright 2023 Greptime Team -// -// 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. - -use std::collections::BTreeMap; -use std::sync::Arc; - -use arrow_schema::{DataType as ArrowDataType, FieldRef}; -use datatypes::arrow::datatypes::Schema; -use datatypes::extension::json::is_json2_extension_type; -use store_api::storage::NestedPath; - -/// Aligns nested struct fields according to the requested nested paths. -/// -/// For each root field: -/// - An empty path list keeps the whole field unchanged. -/// - Non-JSON root fields ignore nested paths and keep the whole field unchanged. -/// - JSON2 root fields are rebuilt from `nested_paths`. -/// - Existing schema fields are preserved only when they are the requested leaf. -/// - Requested paths missing from the schema are synthesized with JSONB (`Binary`) -/// leaves. -/// -/// For example, if the schema has `j: struct>` -/// and `nested_paths` requests `j.a.x` and `j.a.z`, the result is -/// `j: struct>`. -pub(crate) fn align_schema_by_nested_paths<'a, I>(schema: &mut Schema, nested_paths: I) -where - I: IntoIterator, -{ - let fields = schema - .fields - .into_iter() - .zip(nested_paths) - .map(|(field, paths)| { - if !paths.is_empty() && is_json2_extension_type(field) { - let child_paths = paths - .iter() - .map(|path| { - if path.first().is_some_and(|root| root == field.name()) { - &path[1..] - } else { - path - } - }) - .collect::>(); - rebuild_field_by_nested_paths(field, &child_paths) - } else { - field.clone() - } - }) - .collect::>(); - schema.fields = fields.into() -} - -fn rebuild_field_by_nested_paths(field: &FieldRef, nested_paths: &[&[String]]) -> FieldRef { - if nested_paths.iter().any(|path| path.is_empty()) { - return field.clone(); - }; - - let fields = group_child_paths(nested_paths) - .into_iter() - .map(|(name, paths)| { - let existing = find_struct_child(field, &name); - build_field_from_nested_paths(&name, existing, &paths) - }) - .collect::>(); - - Arc::new( - field - .as_ref() - .clone() - .with_data_type(ArrowDataType::Struct(fields.into())), - ) -} - -fn group_child_paths<'a>(nested_paths: &[&'a [String]]) -> BTreeMap> { - let mut child_paths = BTreeMap::>::new(); - for path in nested_paths { - let Some((name, remaining)) = path.split_first() else { - continue; - }; - child_paths.entry(name.clone()).or_default().push(remaining); - } - child_paths -} - -fn find_struct_child<'a>(field: &'a FieldRef, name: &str) -> Option<&'a FieldRef> { - let ArrowDataType::Struct(fields) = field.data_type() else { - return None; - }; - fields.iter().find(|field| field.name() == name) -} - -fn build_field_from_nested_paths( - name: &str, - existing: Option<&FieldRef>, - nested_paths: &[&[String]], -) -> FieldRef { - if nested_paths.iter().any(|path| path.is_empty()) { - return existing.cloned().unwrap_or_else(|| new_jsonb_field(name)); - } - - let fields = group_child_paths(nested_paths) - .into_iter() - .map(|(name, paths)| { - let existing_child = existing.and_then(|field| find_struct_child(field, &name)); - build_field_from_nested_paths(&name, existing_child, &paths) - }) - .collect::>(); - - let field = existing - .map(|field| field.as_ref().clone()) - .unwrap_or_else(|| arrow_schema::Field::new(name, ArrowDataType::Binary, true)); - Arc::new(field.with_data_type(ArrowDataType::Struct(fields.into()))) -} - -fn new_jsonb_field(name: &str) -> FieldRef { - Arc::new(arrow_schema::Field::new(name, ArrowDataType::Binary, true)) -} - -#[cfg(test)] -mod tests { - use arrow_schema::Field; - use datatypes::extension::json::Json2ExtensionType; - - use super::*; - - #[test] - fn test_align_schema_by_nested_paths() { - fn new_field(name: &str, data_type: ArrowDataType) -> FieldRef { - Arc::new(Field::new(name, data_type, true)) - } - - fn struct_field(name: &str, fields: impl IntoIterator) -> FieldRef { - new_field(name, ArrowDataType::Struct(fields.into_iter().collect())) - } - - fn json_struct_field(name: &str, fields: impl IntoIterator) -> FieldRef { - Arc::new( - Field::new( - name, - ArrowDataType::Struct(fields.into_iter().collect()), - true, - ) - .with_extension_type(Json2ExtensionType::default()), - ) - } - - let mut schema = Schema::new([ - json_struct_field( - "j", - [ - struct_field( - "a", - [ - new_field("x", ArrowDataType::Int64), - new_field("y", ArrowDataType::Utf8), - struct_field( - "z", - [ - new_field("q", ArrowDataType::Boolean), - new_field("r", ArrowDataType::Float64), - ], - ), - ], - ), - new_field("b", ArrowDataType::Utf8), - struct_field( - "c", - vec![ - new_field("d", ArrowDataType::Int64), - new_field("e", ArrowDataType::Utf8), - ], - ), - ], - ), - new_field("tag", ArrowDataType::Utf8), - struct_field( - "k", - [ - new_field("k_0", ArrowDataType::Int64), - new_field("k_1", ArrowDataType::Utf8), - ], - ), - ]); - - let nested_paths = [ - vec![ - ["j", "a", "x"].iter().map(|x| x.to_string()).collect(), - ["j", "a", "z", "q"].iter().map(|x| x.to_string()).collect(), - ["j", "a", "m"].iter().map(|x| x.to_string()).collect(), - ["j", "b", "x"].iter().map(|x| x.to_string()).collect(), - ["j", "c"].iter().map(|x| x.to_string()).collect(), - ["j", "d", "e"].iter().map(|x| x.to_string()).collect(), - ], - vec![["tag", "ignored"].iter().map(|x| x.to_string()).collect()], - vec![], - ]; - - align_schema_by_nested_paths( - &mut schema, - nested_paths.iter().map(|paths| paths.as_slice()), - ); - - let expected = Schema::new([ - json_struct_field( - "j", - [ - struct_field( - "a", - [ - new_field("m", ArrowDataType::Binary), - new_field("x", ArrowDataType::Int64), - struct_field("z", vec![new_field("q", ArrowDataType::Boolean)]), - ], - ), - struct_field("b", [new_field("x", ArrowDataType::Binary)]), - struct_field( - "c", - [ - new_field("d", ArrowDataType::Int64), - new_field("e", ArrowDataType::Utf8), - ], - ), - struct_field("d", [new_field("e", ArrowDataType::Binary)]), - ], - ), - new_field("tag", ArrowDataType::Utf8), - struct_field( - "k", - [ - new_field("k_0", ArrowDataType::Int64), - new_field("k_1", ArrowDataType::Utf8), - ], - ), - ]); - - assert_eq!(schema, expected); - } -} diff --git a/src/mito2/src/sst/parquet/prefilter.rs b/src/mito2/src/sst/parquet/prefilter.rs index 5b08a0e4957..fcaaae4a54d 100644 --- a/src/mito2/src/sst/parquet/prefilter.rs +++ b/src/mito2/src/sst/parquet/prefilter.rs @@ -1457,9 +1457,7 @@ mod tests { Arc::new(sst_region_metadata_with_encoding(PrimaryKeyEncoding::Dense)); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", false, @@ -1495,9 +1493,7 @@ mod tests { )); let legacy_read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "memtable", false, @@ -1524,9 +1520,7 @@ mod tests { let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); let raw_pk_read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "memtable", true, @@ -1560,9 +1554,7 @@ mod tests { let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); let full_read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, @@ -1594,7 +1586,7 @@ mod tests { let ts = metric_metadata.time_index_column().column_id; let projected_read_format = FlatReadFormat::new( metric_metadata.clone(), - ReadColumns::from_deduped_column_ids([field_0, ts]), + ReadColumns::new([field_0, ts]), None, "test", true, @@ -1646,9 +1638,7 @@ mod tests { )); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", false, @@ -1686,9 +1676,7 @@ mod tests { let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, @@ -1734,9 +1722,7 @@ mod tests { Arc::new(sst_region_metadata_with_encoding(PrimaryKeyEncoding::Dense)); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", false, diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 97a4abfac02..4f890cb83e8 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -451,7 +451,7 @@ impl ParquetReaderBuilder { } else { let expected_meta = self.expected_metadata.as_ref().unwrap_or(®ion_meta); // Lists all column ids to read, we always use the expected metadata if possible. - ReadColumns::from_deduped_column_ids( + ReadColumns::new( expected_meta .column_metadatas .iter() @@ -2472,9 +2472,7 @@ mod tests { let metadata = Arc::new(sst_region_metadata()); let format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, @@ -2679,7 +2677,7 @@ mod tests { let region_metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); let read_format = FlatReadFormat::new( region_metadata.clone(), - ReadColumns::from_deduped_column_ids( + ReadColumns::new( region_metadata .column_metadatas .iter() @@ -2946,9 +2944,7 @@ mod tests { let expected_metadata = expected_metadata_with_reused_tag_name(metadata.as_ref()); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, @@ -2970,9 +2966,7 @@ mod tests { let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); let read_format = FlatReadFormat::new( metadata.clone(), - ReadColumns::from_deduped_column_ids( - metadata.column_metadatas.iter().map(|c| c.column_id), - ), + ReadColumns::new(metadata.column_metadatas.iter().map(|c| c.column_id)), None, "test", true, diff --git a/src/query/src/optimizer/json_type_concretize.rs b/src/query/src/optimizer/json_type_concretize.rs index 4210433eb3a..4625b619784 100644 --- a/src/query/src/optimizer/json_type_concretize.rs +++ b/src/query/src/optimizer/json_type_concretize.rs @@ -16,12 +16,14 @@ use std::collections::HashMap; use arrow_schema::DataType; use common_function::scalars::json::json_get::JsonGetWithType; -use datafusion::datasource::DefaultTableSource; +use datafusion::datasource::{DefaultTableSource, TableProvider}; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion_common::{Result, plan_datafusion_err, plan_err}; use datafusion_expr::{Expr, LogicalPlan}; use datafusion_optimizer::{OptimizerConfig, OptimizerRule}; +use datatypes::extension::json::is_json2_extension_type; use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; +use table::table::adapter::DfTableProviderAdapter; use crate::dummy_catalog::DummyTableProvider; @@ -56,22 +58,56 @@ impl OptimizerRule for JsonTypeConcretizeRule { return Ok(Transformed::no(plan)); }; - let Some(adapter) = source - .table_provider - .as_any() - .downcast_ref::() - else { - return Ok(Transformed::no(plan)); - }; - - adapter.with_json_type_hint(json_types.clone()); - Ok(Transformed::yes(plan)) + if apply_json_type_hint(source.table_provider.as_ref(), &json_types) { + Ok(Transformed::yes(plan)) + } else { + Ok(Transformed::no(plan)) + } } _ => Ok(Transformed::no(plan)), }) } } +// FIXME: `json_types` is keyed only by unqualified column name. In joins with +// same-named JSON2 columns, a hint deduced from one scan can be applied to +// another scan. Carry the originating relation/scan when deducing hints. +/// Applies JSON type hints to providers that can carry scan request hints. +/// +/// Returns `true` if at least one JSON2 hint is retained and written to the provider. +fn apply_json_type_hint( + provider: &dyn TableProvider, + json_types: &HashMap, +) -> bool { + let schema = provider.schema(); + let json_types = json_types + .iter() + .filter(|(column, _)| { + schema + .fields() + .iter() + .any(|field| field.name() == *column && is_json2_extension_type(field)) + }) + .map(|(column, json_type)| (column.clone(), json_type.clone())) + .collect::>(); + + if json_types.is_empty() { + return false; + } + + if let Some(adapter) = provider.as_any().downcast_ref::() { + adapter.with_json_type_hint(json_types); + return true; + } + + if let Some(adapter) = provider.as_any().downcast_ref::() { + adapter.with_json_type_hint(json_types); + return true; + } + + false +} + fn deduce_json_types(plan: &LogicalPlan) -> Result> { let mut json_types = HashMap::::new(); @@ -228,11 +264,11 @@ mod tests { #[test] fn test_json_type_concretize_rule_rewrite() -> Result<()> { let exprs = vec![ - json_get_expr(col("k0"), path_expr("a.b"), Some(DataType::Int64))?.alias("ab"), - json_get_expr(col("k0"), path_expr("a.c"), None)?.alias("ac"), - json_get_expr(col("k0"), path_expr("d"), Some(DataType::Boolean))?.alias("d"), + json_get_expr(col("j"), path_expr("a.b"), Some(DataType::Int64))?.alias("ab"), + json_get_expr(col("j"), path_expr("a.c"), None)?.alias("ac"), + json_get_expr(col("j"), path_expr("d"), Some(DataType::Boolean))?.alias("d"), ]; - let (provider, plan) = build_plan(exprs)?; + let (provider, plan) = build_json2_plan(exprs)?; assert!( JsonTypeConcretizeRule @@ -253,17 +289,17 @@ mod tests { let request = provider.scan_request(); assert_eq!(1, request.json_type_hint.len()); - assert_eq!(Some(&expected), request.json_type_hint.get("k0")); + assert_eq!(Some(&expected), request.json_type_hint.get("j")); Ok(()) } #[test] fn test_json_type_concretize_rule_conflict_to_variant() -> Result<()> { let exprs = vec![ - json_get_expr(col("k0"), path_expr("a"), Some(DataType::Int64))?.alias("a_num"), - json_get_expr(col("k0"), path_expr("a.b"), Some(DataType::Boolean))?.alias("a_obj"), + json_get_expr(col("j"), path_expr("a"), Some(DataType::Int64))?.alias("a_num"), + json_get_expr(col("j"), path_expr("a.b"), Some(DataType::Boolean))?.alias("a_obj"), ]; - let (provider, plan) = build_plan(exprs)?; + let (provider, plan) = build_json2_plan(exprs)?; assert!( JsonTypeConcretizeRule @@ -277,11 +313,26 @@ mod tests { )])); assert_eq!( Some(&expected), - provider.scan_request().json_type_hint.get("k0") + provider.scan_request().json_type_hint.get("j") ); Ok(()) } + #[test] + fn test_json_type_concretize_rule_ignores_non_json2_columns() -> Result<()> { + let exprs = + vec![json_get_expr(col("k0"), path_expr("a.b"), Some(DataType::Int64))?.alias("ab")]; + let (provider, plan) = build_plan(exprs)?; + + assert!( + !JsonTypeConcretizeRule + .rewrite(plan, &OptimizerContext::default())? + .transformed + ); + assert!(provider.scan_request().json_type_hint.is_empty()); + Ok(()) + } + #[test] fn test_json_type_concretize_rule_no_json_get() -> Result<()> { let (provider, plan) = build_plan(vec![col("k0"), col("v0")])?; diff --git a/src/table/src/table/adapter.rs b/src/table/src/table/adapter.rs index 0ae56ce1c31..ac71b448e94 100644 --- a/src/table/src/table/adapter.rs +++ b/src/table/src/table/adapter.rs @@ -28,6 +28,7 @@ use datafusion_expr::TableProviderFilterPushDown as DfTableProviderFilterPushDow use datafusion_expr::expr::Expr; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; +use datatypes::types::json_type::JsonNativeType; use store_api::storage::{ScanRequest, VectorSearchRequest}; use crate::table::{TableRef, TableType}; @@ -63,6 +64,10 @@ impl DfTableProviderAdapter { self.scan_req.lock().unwrap().vector_search = Some(hint); } + pub fn with_json_type_hint(&self, hint: std::collections::HashMap) { + self.scan_req.lock().unwrap().json_type_hint = hint; + } + pub fn get_vector_search_hint(&self) -> Option { self.scan_req.lock().unwrap().vector_search.clone() } diff --git a/tests/cases/standalone/common/types/json/json2_limit.result b/tests/cases/standalone/common/types/json/json2_limit.result index 3cb462a8a7e..035c4ed07bf 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.result +++ b/tests/cases/standalone/common/types/json/json2_limit.result @@ -77,6 +77,75 @@ drop table json2_whole_and_path_read; Affected Rows: 0 +create table json2_join_same_name_left ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +Affected Rows: 0 + +create table json2_join_same_name_right ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +Affected Rows: 0 + +insert into json2_join_same_name_left values + (1, 'a', '{"a": 1, "left_only": "kept"}'); + +Affected Rows: 1 + +insert into json2_join_same_name_right values + (1, 'a', '{"a": "right", "right_only": "should be kept"}'); + +Affected Rows: 1 + +admin flush_table('json2_join_same_name_left'); + ++------------------------------------------------+ +| ADMIN flush_table('json2_join_same_name_left') | ++------------------------------------------------+ +| 0 | ++------------------------------------------------+ + +admin flush_table('json2_join_same_name_right'); + ++-------------------------------------------------+ +| ADMIN flush_table('json2_join_same_name_right') | ++-------------------------------------------------+ +| 0 | ++-------------------------------------------------+ + +-- FIXME: This should return `right` and `1`. The current NULL values are caused +-- by JSON type hints losing the table qualifier in joins. +select json_get(r.j, 'a')::string, json_get(l.j, 'a')::int64 +from json2_join_same_name_left l +join json2_join_same_name_right r +on l.k = r.k; + ++---------------------------------------------+---------------------------------------------------+ +| json_get(r.j,Utf8("a")) | arrow_cast(json_get(l.j,Utf8("a")),Utf8("Int64")) | ++---------------------------------------------+---------------------------------------------------+ +| {"a":"right","right_only":"should be kept"} | | ++---------------------------------------------+---------------------------------------------------+ + +drop table json2_join_same_name_left; + +Affected Rows: 0 + +drop table json2_join_same_name_right; + +Affected Rows: 0 + create table json2_without_append_mode ( ts timestamp time index, j json2 diff --git a/tests/cases/standalone/common/types/json/json2_limit.sql b/tests/cases/standalone/common/types/json/json2_limit.sql index 4917b9690e5..986ed58eb2a 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.sql +++ b/tests/cases/standalone/common/types/json/json2_limit.sql @@ -44,6 +44,45 @@ select j from json2_whole_and_path_read where j.a.b = 1; drop table json2_whole_and_path_read; +create table json2_join_same_name_left ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +create table json2_join_same_name_right ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +insert into json2_join_same_name_left values + (1, 'a', '{"a": 1, "left_only": "kept"}'); + +insert into json2_join_same_name_right values + (1, 'a', '{"a": "right", "right_only": "should be kept"}'); + +admin flush_table('json2_join_same_name_left'); + +admin flush_table('json2_join_same_name_right'); + +-- FIXME: This should return `right` and `1`. The current NULL values are caused +-- by JSON type hints losing the table qualifier in joins. +select json_get(r.j, 'a')::string, json_get(l.j, 'a')::int64 +from json2_join_same_name_left l +join json2_join_same_name_right r +on l.k = r.k; + +drop table json2_join_same_name_left; + +drop table json2_join_same_name_right; + create table json2_without_append_mode ( ts timestamp time index, j json2 From 76924c2d36b8d6b3d2afb111536e7bda04d0bda8 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Fri, 14 Aug 2026 06:32:17 +0000 Subject: [PATCH 47/92] feat(mito2): introduce two-phase metric series scans (#8826) * feat(mito2): add two-phase series scan Signed-off-by: evenyag * docs: regenerate configuration reference Signed-off-by: evenyag * test(sqlness): update series scan explain results Signed-off-by: evenyag * test: update config API expectation Signed-off-by: evenyag * fix(mito2): bound two-phase series discovery Signed-off-by: evenyag * fix(mito2): avoid candidate distribution deadlock Signed-off-by: evenyag * chore(mito2): remove obsolete dead code allowances Signed-off-by: evenyag * fix(mito2): share series scan memory pool Signed-off-by: evenyag --------- Signed-off-by: evenyag --- config/config.md | 2 + config/datanode.example.toml | 3 + config/standalone.example.toml | 3 + src/metric-engine/src/engine/read.rs | 173 +++++++- src/mito2/src/config.rs | 12 + src/mito2/src/engine.rs | 32 ++ src/mito2/src/engine/scan_test.rs | 196 ++++++++- src/mito2/src/read/scan_region.rs | 35 +- src/mito2/src/read/series_candidate.rs | 13 +- src/mito2/src/read/series_reader.rs | 68 ++- src/mito2/src/read/series_scan.rs | 408 ++++++++++++++++-- tests-integration/tests/http.rs | 1 + .../explain/step_aggr_advance.result | 28 +- .../tql-explain-analyze/analyze.result | 14 +- .../promql/histogram_multi_partition.result | 4 +- .../promql/tsid_binary_join_regression.result | 56 +-- .../tsid_histogram_quantile_regression.result | 2 +- .../common/select/tql_filter.result | 6 +- .../standalone/common/tql/partition.result | 10 +- .../standalone/optimizer/windowed_sort.result | 2 +- .../tql-explain-analyze/analyze.result | 22 +- .../tql-explain-analyze/column_pruning.result | 4 +- .../tql-explain-analyze/tsid_column.result | 12 +- 23 files changed, 963 insertions(+), 143 deletions(-) diff --git a/config/config.md b/config/config.md index c7737a164b8..123eaf90a20 100644 --- a/config/config.md +++ b/config/config.md @@ -185,6 +185,7 @@ | `region_engine.mito.min_compaction_interval` | String | `0m` | Minimum time interval between two compactions.
To align with the old behavior, the default value is 0 (no restrictions). | | `region_engine.mito.schedule_compaction_after_edit` | Bool | `true` | Whether to allow to schedule a compaction after a successful region edit.

Setting this to "true" is a necessary but not sufficient condition for scheduling compaction after a region edit.
Other constraints, such as "min_compaction_interval", may still prevent compaction from being scheduled.
Setting this to "false", however, guarantees that compaction will not be scheduled after a region edit. | | `region_engine.mito.default_flat_format` | Bool | `true` | Whether to enable flat format as the default SST format. | +| `region_engine.mito.experimental_series_scan_v2` | Bool | `true` | Whether to enable the experimental two-phase mode for eligible metric series scans. | | `region_engine.mito.index` | -- | -- | The options for index in Mito engine. | | `region_engine.mito.index.aux_path` | String | `""` | Auxiliary directory path for the index in filesystem, used to store intermediate files for
creating the index and staging files for searching the index, defaults to `{data_home}/index_intermediate`.
The default name for this directory is `index_intermediate` for backward compatibility.

This path contains two subdirectories:
- `__intm`: for storing intermediate files used during creating index.
- `staging`: for storing staging files used during searching index. | | `region_engine.mito.index.staging_size` | String | `2GB` | The max capacity of the staging directory. | @@ -606,6 +607,7 @@ | `region_engine.mito.min_compaction_interval` | String | `0m` | Minimum time interval between two compactions.
To align with the old behavior, the default value is 0 (no restrictions). | | `region_engine.mito.schedule_compaction_after_edit` | Bool | `true` | Whether to allow to schedule a compaction after a successful region edit.

Setting this to "true" is a necessary but not sufficient condition for scheduling compaction after a region edit.
Other constraints, such as "min_compaction_interval", may still prevent compaction from being scheduled.
Setting this to "false", however, guarantees that compaction will not be scheduled after a region edit. | | `region_engine.mito.default_flat_format` | Bool | `true` | Whether to enable flat format as the default SST format. | +| `region_engine.mito.experimental_series_scan_v2` | Bool | `true` | Whether to enable the experimental two-phase mode for eligible metric series scans. | | `region_engine.mito.index` | -- | -- | The options for index in Mito engine. | | `region_engine.mito.index.aux_path` | String | `""` | Auxiliary directory path for the index in filesystem, used to store intermediate files for
creating the index and staging files for searching the index, defaults to `{data_home}/index_intermediate`.
The default name for this directory is `index_intermediate` for backward compatibility.

This path contains two subdirectories:
- `__intm`: for storing intermediate files used during creating index.
- `staging`: for storing staging files used during searching index. | | `region_engine.mito.index.staging_size` | String | `2GB` | The max capacity of the staging directory. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 3426630d7b6..1f1377d91cc 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -571,6 +571,9 @@ schedule_compaction_after_edit = true ## Whether to enable flat format as the default SST format. default_flat_format = true +## Whether to enable the experimental two-phase mode for eligible metric series scans. +experimental_series_scan_v2 = true + ## The options for index in Mito engine. [region_engine.mito.index] diff --git a/config/standalone.example.toml b/config/standalone.example.toml index c8dff4944c1..068025e2a28 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -729,6 +729,9 @@ schedule_compaction_after_edit = true ## Whether to enable flat format as the default SST format. default_flat_format = true +## Whether to enable the experimental two-phase mode for eligible metric series scans. +experimental_series_scan_v2 = true + ## The options for index in Mito engine. [region_engine.mito.index] diff --git a/src/metric-engine/src/engine/read.rs b/src/metric-engine/src/engine/read.rs index 53dd086d225..ee086184c23 100644 --- a/src/metric-engine/src/engine/read.rs +++ b/src/metric-engine/src/engine/read.rs @@ -306,15 +306,72 @@ impl MetricEngineInner { #[cfg(test)] mod test { + use std::fmt; + + use api::v1::Rows; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; - use store_api::region_request::RegionRequest; + use datafusion::physical_plan::DisplayFormatType; + use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; + use futures_util::TryStreamExt; + use futures_util::future::try_join_all; + use mito2::config::MitoConfig; + use store_api::region_engine::{PrepareRequest, QueryScanContext}; + use store_api::region_request::{RegionFlushRequest, RegionPutRequest, RegionRequest}; + use store_api::storage::TimeSeriesDistribution; use super::*; + use crate::config::EngineConfig; use crate::test_util::{ - TestEnv, alter_logical_region_add_tag_columns, create_logical_region_request, + self, TestEnv, alter_logical_region_add_tag_columns, create_logical_region_request, }; + struct ScannerDisplay<'a>(&'a dyn store_api::region_engine::RegionScanner); + + impl fmt::Display for ScannerDisplay<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt_as(DisplayFormatType::Default, f) + } + } + + async fn count_scanner_rows(scanner: &mut RegionScannerRef, partitions: usize) -> usize { + let ranges = scanner + .properties() + .partitions + .iter() + .flatten() + .copied() + .collect::>(); + let mut prepared = vec![Vec::new(); partitions]; + prepared[0] = ranges; + scanner + .prepare( + PrepareRequest::default() + .with_ranges(prepared) + .with_target_partitions(partitions), + ) + .unwrap(); + + let metrics = ExecutionPlanMetricsSet::default(); + let context = QueryScanContext::default(); + let streams = (0..partitions) + .map(|partition| { + scanner + .scan_partition(&context, &metrics, partition) + .unwrap() + }) + .collect::>(); + try_join_all(streams.into_iter().map(|stream| async move { + stream + .try_fold(0, |rows, batch| async move { Ok(rows + batch.num_rows()) }) + .await + })) + .await + .unwrap() + .into_iter() + .sum() + } + #[tokio::test] async fn test_invalid_logical_projection() { let env = TestEnv::new().await; @@ -406,4 +463,116 @@ mod test { &[11, 10, 9, 8, 0, 1, 4] ); } + + #[tokio::test] + async fn test_two_phase_series_scan_reads_metric_region() { + let env = TestEnv::with_mito_config( + "test_two_phase_series_scan", + MitoConfig { + experimental_series_scan_v2: true, + ..Default::default() + }, + EngineConfig::default(), + ) + .await; + env.init_metric_region().await; + + let physical_region_id = env.default_physical_region_id(); + let logical_region_id = env.default_logical_region_id(); + let logical_region_id_2 = RegionId::new(1024, logical_region_id.region_number()); + env.metric() + .handle_request( + logical_region_id_2, + RegionRequest::Create(create_logical_region_request( + &["job"], + physical_region_id, + "test_metric_region_2", + )), + ) + .await + .unwrap(); + + let schema = test_util::row_schema_with_tags(&["job"]); + let put = |rows| { + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: schema.clone(), + rows: test_util::build_rows(1, rows), + }, + hint: None, + partition_expr_version: None, + }) + }; + env.metric() + .handle_request(logical_region_id, put(3)) + .await + .unwrap(); + env.metric() + .handle_request( + physical_region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + env.metric() + .handle_request(logical_region_id, put(2)) + .await + .unwrap(); + env.metric() + .handle_request(logical_region_id_2, put(4)) + .await + .unwrap(); + + let data_region_id = utils::to_data_region_id(physical_region_id); + let request = ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + ..Default::default() + }; + let mut scanner = env + .mito() + .handle_query(data_region_id, request) + .await + .unwrap(); + + let rows = tokio::time::timeout( + std::time::Duration::from_secs(30), + count_scanner_rows(&mut scanner, 2), + ) + .await + .expect("two-phase series scan should not deadlock"); + assert_eq!(7, rows); + + let explain = ScannerDisplay(scanner.as_ref()).to_string(); + assert!(explain.contains("\"mode\":\"two_phase\""), "{explain}"); + } + + #[tokio::test] + async fn test_series_scan_v2_flag_disables_two_phase_mode() { + let env = TestEnv::with_mito_config( + "test_disable_two_phase_series_scan", + MitoConfig { + experimental_series_scan_v2: false, + ..Default::default() + }, + EngineConfig::default(), + ) + .await; + env.init_metric_region().await; + + let data_region_id = utils::to_data_region_id(env.default_physical_region_id()); + let scanner = env + .mito() + .handle_query( + data_region_id, + ScanRequest { + distribution: Some(TimeSeriesDistribution::PerSeries), + ..Default::default() + }, + ) + .await + .unwrap(); + + let explain = ScannerDisplay(scanner.as_ref()).to_string(); + assert!(explain.contains("\"mode\":\"legacy\""), "{explain}"); + } } diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index f7489709ee3..54445549a05 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -188,6 +188,9 @@ pub struct MitoConfig { /// When enabled, forces using BulkMemtable and BulkMemtableBuilder. pub default_flat_format: bool, + /// Whether to enable the experimental two-phase mode for eligible metric series scans. + pub experimental_series_scan_v2: bool, + pub gc: GcConfig, } @@ -239,6 +242,7 @@ impl Default for MitoConfig { min_compaction_interval: Duration::from_secs(0), schedule_compaction_after_edit: true, default_flat_format: true, + experimental_series_scan_v2: true, gc: GcConfig::default(), }; @@ -389,6 +393,14 @@ mod tests { assert_eq!(ReadableSize::mb(512), config.sst_meta_cache_size); assert_eq!(ReadableSize::mb(128), config.prefilter_result_cache_size); } + + #[test] + fn test_experimental_series_scan_v2_config() { + assert!(MitoConfig::default().experimental_series_scan_v2); + + let config: MitoConfig = toml::from_str("experimental_series_scan_v2 = false").unwrap(); + assert!(!config.experimental_series_scan_v2); + } } /// Index build mode. diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index a128531b535..ddd34669eb7 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -102,6 +102,7 @@ use common_recordbatch::{QueryMemoryTracker, SendableRecordBatchStream}; use common_stat::get_total_memory_bytes; use common_telemetry::{info, tracing, warn}; use common_wal::options::WalOptions; +use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, UnboundedMemoryPool}; use futures::future::{join_all, try_join_all}; use futures::stream::{self, Stream, StreamExt}; use object_store::manager::ObjectStoreManagerRef; @@ -236,6 +237,7 @@ impl<'a, S: LogStore> MitoEngineBuilder<'a, S> { let wal_raw_entry_reader = Arc::new(LogStoreRawEntryReader::new(self.log_store)); let total_memory = get_total_memory_bytes().max(0) as u64; let scan_memory_limit = config.scan_memory_limit.resolve(total_memory) as usize; + let scan_memory_pool = new_scan_memory_pool(scan_memory_limit); let scan_memory_tracker = QueryMemoryTracker::builder(scan_memory_limit, config.scan_memory_on_exhausted) .on_update(|usage| { @@ -254,6 +256,7 @@ impl<'a, S: LogStore> MitoEngineBuilder<'a, S> { config, wal_raw_entry_reader, scan_memory_tracker, + scan_memory_pool, region_hook, #[cfg(feature = "enterprise")] extension_range_provider_factory: None, @@ -739,6 +742,8 @@ struct EngineInner { wal_raw_entry_reader: Arc, /// Memory tracker for table scans. scan_memory_tracker: QueryMemoryTracker, + /// Memory pool shared by internal scan operators across all queries. + scan_memory_pool: Arc, /// The region hook (if any) registered via plugins; exposed for the GC worker /// to fire [`RegionHook::on_region_gc`]. region_hook: Option, @@ -1097,6 +1102,8 @@ impl EngineInner { ) .with_query_stat_counters(region.region_stats.query_stat_counters()) .with_max_concurrent_scan_files(self.config.max_concurrent_scan_files) + .with_scan_memory_pool(self.scan_memory_pool.clone()) + .with_experimental_series_scan_v2(self.config.experimental_series_scan_v2) .with_ignore_inverted_index(self.config.inverted_index.apply_on_query.disabled()) .with_ignore_fulltext_index(self.config.fulltext_index.apply_on_query.disabled()) .with_ignore_bloom_filter(self.config.bloom_filter_index.apply_on_query.disabled()) @@ -1443,6 +1450,7 @@ impl MitoEngine { let wal_raw_entry_reader = Arc::new(LogStoreRawEntryReader::new(log_store.clone())); let total_memory = get_total_memory_bytes().max(0) as u64; let scan_memory_limit = config.scan_memory_limit.resolve(total_memory) as usize; + let scan_memory_pool = new_scan_memory_pool(scan_memory_limit); let scan_memory_tracker = QueryMemoryTracker::builder(scan_memory_limit, config.scan_memory_on_exhausted) .on_update(|usage| { @@ -1472,6 +1480,7 @@ impl MitoEngine { config, wal_raw_entry_reader, scan_memory_tracker, + scan_memory_pool, region_hook: None, #[cfg(feature = "enterprise")] extension_range_provider_factory: None, @@ -1485,10 +1494,20 @@ impl MitoEngine { } } +fn new_scan_memory_pool(scan_memory_limit: usize) -> Arc { + if scan_memory_limit == 0 { + Arc::new(UnboundedMemoryPool::default()) + } else { + Arc::new(GreedyMemoryPool::new(scan_memory_limit)) + } +} + #[cfg(test)] mod tests { use std::time::Duration; + use datafusion::execution::memory_pool::MemoryConsumer; + use super::*; use crate::sst::file::FileMeta; @@ -1574,4 +1593,17 @@ mod tests { }; assert!(!is_valid_region_edit(&edit)); } + + #[test] + fn test_scan_memory_pool_is_shared_across_consumers() { + let pool = new_scan_memory_pool(100); + let cloned_pool = pool.clone(); + let first = MemoryConsumer::new("first-scan").register(&pool); + let second = MemoryConsumer::new("second-scan").register(&cloned_pool); + + first.try_grow(60).unwrap(); + assert!(second.try_grow(50).is_err()); + second.try_grow(40).unwrap(); + assert_eq!(100, pool.reserved()); + } } diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index 4baa707f4e7..a53b77c8253 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -13,11 +13,12 @@ // limitations under the License. use std::collections::{BTreeMap, HashMap}; +use std::sync::Arc; use api::helper::encode_json_value; -use api::v1::Rows; use api::v1::helper::row; use api::v1::value::ValueData; +use api::v1::{ColumnDataType, Rows, SemanticType, WriteHint}; use common_base::readable_size::ReadableSize; use common_error::ext::{ErrorExt, WhateverResult}; use common_error::status_code::StatusCode; @@ -26,14 +27,18 @@ use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_common::ScalarValue; use datafusion_expr::{col, lit}; use datatypes::arrow::array::AsArray; -use datatypes::arrow::datatypes::{Float64Type, TimestampMillisecondType}; +use datatypes::arrow::datatypes::{Float64Type, TimestampMillisecondType, UInt64Type}; use datatypes::json::value::JsonValue; use datatypes::prelude::ConcreteDataType; use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; use futures::TryStreamExt; +use futures::future::try_join_all; use serde_json::json; +use store_api::codec::PrimaryKeyEncoding; +use store_api::metric_engine_consts::PRIMARY_KEY_ENCODING; use store_api::region_engine::{PrepareRequest, RegionEngine, RegionScanner}; -use store_api::region_request::RegionRequest; +use store_api::region_request::{RegionPutRequest, RegionRequest}; +use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; use store_api::storage::{RegionId, ScanRequest, TimeSeriesDistribution}; use crate::config::MitoConfig; @@ -41,6 +46,7 @@ use crate::error::Error; use crate::read::read_columns::ReadColumns; use crate::read::scan_region::Scanner; use crate::test_util; +use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding}; use crate::test_util::{CreateRequestBuilder, TestEnv}; #[tokio::test] @@ -753,6 +759,7 @@ async fn test_series_scan_with_format(flat_format: bool) { let Scanner::Series(mut scanner) = scanner else { panic!("Scanner should be series scan"); }; + assert_eq!("legacy", scanner.mode()); // 3 partition ranges for 3 time window. assert_eq!( 3, @@ -789,6 +796,189 @@ async fn test_series_scan_with_format(flat_format: bool) { assert_eq!(expected_rows, actual_rows); } +#[tokio::test] +async fn test_two_phase_series_scan() { + let mut env = TestEnv::with_prefix("test_two_phase_series_scan").await; + let engine = env + .create_engine(MitoConfig { + experimental_series_scan_v2: true, + ..Default::default() + }) + .await; + + let region_id = RegionId::new(1, 1); + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let mut request = CreateRequestBuilder::new().build(); + request.column_metadatas = metadata.column_metadatas.clone(); + request.primary_key = metadata.primary_key.clone(); + request + .options + .insert(PRIMARY_KEY_ENCODING.to_string(), "sparse".to_string()); + request + .options + .insert("memtable.type".to_string(), "bulk".to_string()); + request + .options + .insert("sst_format".to_string(), "flat".to_string()); + let full_row_schema = test_util::rows_schema(&request); + let mut encoded_primary_key_schema = full_row_schema[0].clone(); + encoded_primary_key_schema.column_name = PRIMARY_KEY_COLUMN_NAME.to_string(); + encoded_primary_key_schema.datatype = ColumnDataType::Binary.into(); + encoded_primary_key_schema.semantic_type = SemanticType::Tag.into(); + let row_schema = vec![ + encoded_primary_key_schema, + full_row_schema[5].clone(), + full_row_schema[4].clone(), + ]; + + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + let rows = |values: &[(u32, u64, &str, &str, u64, i64)]| Rows { + schema: row_schema.clone(), + rows: values + .iter() + .map(|(table_id, tsid, tag_0, tag_1, field, ts)| { + row(vec![ + ValueData::BinaryValue(new_sparse_primary_key( + &[*tag_0, *tag_1], + &metadata, + *table_id, + *tsid, + )), + ValueData::TimestampMillisecondValue(*ts), + ValueData::U64Value(*field), + ]) + }) + .collect(), + }; + let put = |rows| { + RegionRequest::Put(RegionPutRequest { + rows, + hint: Some(WriteHint { + primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse.into(), + }), + partition_expr_version: None, + }) + }; + + engine + .handle_request( + region_id, + put(rows(&[ + (10, 0, "a", "x", 10, 1000), + (10, u64::MAX, "b", "y", 20, 1000), + (20, 0, "c", "z", 30, 1000), + ])), + ) + .await + .unwrap(); + test_util::flush_region(&engine, region_id, None).await; + engine + .handle_request( + region_id, + put(rows(&[ + // Replaces the flushed row for this series and timestamp. + (10, 0, "a", "x", 11, 1000), + (10, 0, "a", "x", 12, 2000), + (10, u64::MAX, "b", "y", 21, 2000), + (20, u64::MAX, "d", "w", 40, 1000), + ])), + ) + .await + .unwrap(); + + let scanner = engine + .scanner( + region_id, + ScanRequest { + // Internal metric identifiers must still be read for candidate discovery. + projection: Some(vec![2, 4, 5]), + distribution: Some(TimeSeriesDistribution::PerSeries), + ..Default::default() + }, + ) + .await + .unwrap(); + let Scanner::Series(mut scanner) = scanner else { + panic!("Scanner should be series scan"); + }; + assert_eq!("two_phase", scanner.mode()); + + let ranges = scanner + .properties() + .partitions + .iter() + .flatten() + .copied() + .collect::>(); + scanner + .prepare( + PrepareRequest::default() + .with_ranges(vec![ranges, Vec::new(), Vec::new()]) + .with_target_partitions(3), + ) + .unwrap(); + + let metrics_set = ExecutionPlanMetricsSet::default(); + let context = store_api::region_engine::QueryScanContext::default(); + let partition_batches = try_join_all((0..3).map(|partition| { + let stream = scanner + .scan_partition(&context, &metrics_set, partition) + .unwrap(); + async move { stream.try_collect::>().await } + })) + .await + .unwrap(); + + let mut series_to_partition = BTreeMap::new(); + let mut actual_rows = Vec::new(); + for (partition, batches) in partition_batches.into_iter().enumerate() { + for batch in batches { + let tags = batch.column_by_name("tag_0").unwrap(); + let fields = batch + .column_by_name("field_0") + .unwrap() + .as_primitive::(); + let timestamps = batch + .column_by_name("ts") + .unwrap() + .as_primitive::(); + for row in 0..batch.num_rows() { + let tag = datatypes::arrow_array::string_array_value_at_index(tags, row) + .unwrap() + .to_string(); + let previous = series_to_partition.insert(tag.clone(), partition); + if let Some(previous) = previous { + assert_eq!(previous, partition, "series {tag} moved partitions"); + } + actual_rows.push((tag, fields.value(row), timestamps.value(row))); + } + } + } + actual_rows.sort(); + assert_eq!( + vec![ + ("a".to_string(), 11, 1000), + ("a".to_string(), 12, 2000), + ("b".to_string(), 20, 1000), + ("b".to_string(), 21, 2000), + ("c".to_string(), 30, 1000), + ("d".to_string(), 40, 1000), + ], + actual_rows + ); + assert_eq!(4, series_to_partition.len()); + assert_eq!(Some(&0), series_to_partition.get("a")); + assert_eq!(Some(&0), series_to_partition.get("c")); + assert_eq!(Some(&2), series_to_partition.get("b")); + assert_eq!(Some(&2), series_to_partition.get("d")); +} + /// Scans all partitions in round-robin fashion and returns rows sorted by (tag, ts). /// Also asserts that each series appears in only one partition. async fn collect_partition_rows_round_robin( diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index c0340464ecc..df199d4bac2 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -28,6 +28,7 @@ use common_recordbatch::filter::SimpleFilterEvaluator; use common_telemetry::tracing::Instrument; use common_telemetry::{debug, error, tracing, warn}; use common_time::range::TimestampRange; +use datafusion::execution::memory_pool::{MemoryPool, UnboundedMemoryPool}; use datafusion::physical_plan::expressions::DynamicFilterPhysicalExpr; use datafusion_common::pruning::PruningStatistics; use datafusion_common::{Column, ScalarValue}; @@ -238,6 +239,10 @@ pub(crate) struct ScanRegion { cache_strategy: CacheStrategy, /// Maximum number of SST files to scan concurrently. max_concurrent_scan_files: usize, + /// Memory pool shared by internal scan operators across all queries. + scan_memory_pool: Arc, + /// Whether to enable the experimental two-phase metric series scan. + experimental_series_scan_v2: bool, /// Whether to ignore inverted index. ignore_inverted_index: bool, /// Whether to ignore fulltext index. @@ -269,6 +274,8 @@ impl ScanRegion { request, cache_strategy, max_concurrent_scan_files: DEFAULT_MAX_CONCURRENT_SCAN_FILES, + scan_memory_pool: Arc::new(UnboundedMemoryPool::default()), + experimental_series_scan_v2: false, ignore_inverted_index: false, ignore_fulltext_index: false, ignore_bloom_filter: false, @@ -297,6 +304,20 @@ impl ScanRegion { self } + /// Sets the memory pool shared by internal scan operators. + #[must_use] + pub(crate) fn with_scan_memory_pool(mut self, scan_memory_pool: Arc) -> Self { + self.scan_memory_pool = scan_memory_pool; + self + } + + /// Sets whether eligible metric series scans use the two-phase implementation. + #[must_use] + pub(crate) fn with_experimental_series_scan_v2(mut self, enabled: bool) -> Self { + self.experimental_series_scan_v2 = enabled; + self + } + /// Sets whether to ignore inverted index. #[must_use] pub(crate) fn with_ignore_inverted_index(mut self, ignore: bool) -> Self { @@ -387,8 +408,9 @@ impl ScanRegion { /// Scans by series. #[tracing::instrument(skip_all, fields(region_id = %self.region_id()))] pub(crate) async fn series_scan(self) -> Result { + let experimental_series_scan_v2 = self.experimental_series_scan_v2; let input = self.scan_input().await?; - Ok(SeriesScan::new(input)) + Ok(SeriesScan::new(input, experimental_series_scan_v2)) } /// Returns true if the region can use unordered scan for current request. @@ -541,6 +563,7 @@ impl ScanRegion { .with_bloom_filter_index_appliers(bloom_filter_appliers) .with_fulltext_index_appliers(fulltext_index_appliers) .with_max_concurrent_scan_files(self.max_concurrent_scan_files) + .with_scan_memory_pool(self.scan_memory_pool) .with_start_time(self.start_time) .with_append_mode(self.version.options.append_mode) .with_filter_deleted(self.filter_deleted) @@ -899,6 +922,8 @@ pub struct ScanInput { ignore_file_not_found: bool, /// Maximum number of SST files to scan concurrently. pub(crate) max_concurrent_scan_files: usize, + /// Memory pool shared by internal scan operators across all queries. + pub(crate) scan_memory_pool: Arc, /// Index appliers. inverted_index_appliers: [Option; 2], bloom_filter_index_appliers: [Option; 2], @@ -950,6 +975,7 @@ impl ScanInput { cache_strategy: CacheStrategy::Disabled, ignore_file_not_found: false, max_concurrent_scan_files: DEFAULT_MAX_CONCURRENT_SCAN_FILES, + scan_memory_pool: Arc::new(UnboundedMemoryPool::default()), inverted_index_appliers: [None, None], bloom_filter_index_appliers: [None, None], fulltext_index_appliers: [None, None], @@ -1037,6 +1063,13 @@ impl ScanInput { self } + /// Sets the memory pool shared by internal scan operators. + #[must_use] + pub(crate) fn with_scan_memory_pool(mut self, scan_memory_pool: Arc) -> Self { + self.scan_memory_pool = scan_memory_pool; + self + } + /// Sets inverted index appliers. #[must_use] pub(crate) fn with_inverted_index_appliers( diff --git a/src/mito2/src/read/series_candidate.rs b/src/mito2/src/read/series_candidate.rs index 03d76d2adf6..2933327a38e 100644 --- a/src/mito2/src/read/series_candidate.rs +++ b/src/mito2/src/read/series_candidate.rs @@ -71,7 +71,6 @@ pub(crate) struct MetricSeriesId { pub(crate) type MetricSeriesIdStream = BoxStream<'static, Result>>; /// Builds candidate metric series from the ranges assigned to a [`SeriesScan`](super::series_scan::SeriesScan). -#[allow(dead_code)] pub(crate) struct SeriesCandidateScanner { stream_ctx: Arc, partitions: Vec>, @@ -82,7 +81,6 @@ pub(crate) struct SeriesCandidateScanner { part_metrics: PartitionMetrics, } -#[allow(dead_code)] impl SeriesCandidateScanner { /// Creates a candidate-series scanner for native memtable and SST ranges. /// @@ -346,8 +344,7 @@ impl SeriesCandidateRangeBuilder { } } -pub(crate) fn validate_metric_metadata(stream_ctx: &StreamContext) -> Result<()> { - let metadata = stream_ctx.input.region_metadata(); +pub(crate) fn is_sparse_metric_metadata(metadata: &store_api::metadata::RegionMetadataRef) -> bool { let valid_prefix = metadata .primary_key .starts_with(&[ReservedColumnId::table_id(), ReservedColumnId::tsid()]); @@ -358,8 +355,14 @@ pub(crate) fn validate_metric_metadata(stream_ctx: &StreamContext) -> Result<()> table_id.column_schema.data_type == ConcreteDataType::uint32_datatype() && tsid.column_schema.data_type == ConcreteDataType::uint64_datatype() }); + + metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse && valid_prefix && valid_types +} + +pub(crate) fn validate_metric_metadata(stream_ctx: &StreamContext) -> Result<()> { + let metadata = stream_ctx.input.region_metadata(); ensure!( - metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse && valid_prefix && valid_types, + is_sparse_metric_metadata(metadata), InvalidRequestSnafu { region_id: metadata.region_id, reason: "candidate-series scan requires sparse (__table_id, __tsid) primary keys", diff --git a/src/mito2/src/read/series_reader.rs b/src/mito2/src/read/series_reader.rs index d16d13f2dac..0a3e008f27a 100644 --- a/src/mito2/src/read/series_reader.rs +++ b/src/mito2/src/read/series_reader.rs @@ -84,12 +84,16 @@ impl SeriesRange { pub(crate) struct AssignedSeriesBatch { range: SeriesRange, series: Vec, + enable_range_cache: bool, } -#[allow(dead_code)] impl AssignedSeriesBatch { - fn new(range: SeriesRange, series: Vec) -> Self { - Self { range, series } + fn new(range: SeriesRange, series: Vec, enable_range_cache: bool) -> Self { + Self { + range, + series, + enable_range_cache, + } } pub(crate) fn range(&self) -> SeriesRange { @@ -99,23 +103,28 @@ impl AssignedSeriesBatch { pub(crate) fn series(&self) -> &[MetricSeriesId] { &self.series } + + pub(crate) fn enable_range_cache(&self) -> bool { + self.enable_range_cache + } } /// Collects candidate batches and assigns every TSID by its integer range. -#[allow(dead_code)] pub(crate) struct SeriesBatchCollector { assignments: Vec>, + num_series: usize, } -#[allow(dead_code)] impl SeriesBatchCollector { pub(crate) fn new(partitions: usize) -> Option { (partitions > 0).then(|| Self { assignments: (0..partitions).map(|_| Vec::new()).collect(), + num_series: 0, }) } pub(crate) fn push(&mut self, batch: Vec) { + self.num_series += batch.len(); let partitions = self.assignments.len(); for series in batch { let partition = SeriesRange::partition_for(series.tsid, partitions); @@ -123,13 +132,21 @@ impl SeriesBatchCollector { } } - pub(crate) fn finish(self) -> Vec { + pub(crate) fn len(&self) -> usize { + self.num_series + } + + pub(crate) fn finish(self, enable_range_cache: bool) -> Vec { let partitions = self.assignments.len(); self.assignments .into_iter() .enumerate() .map(|(partition, series)| { - AssignedSeriesBatch::new(SeriesRange::new(partition, partitions).unwrap(), series) + AssignedSeriesBatch::new( + SeriesRange::new(partition, partitions).unwrap(), + series, + enable_range_cache, + ) }) .collect() } @@ -140,14 +157,16 @@ impl SeriesBatchCollector { struct MetricSeriesFilter { range: SeriesRange, series: Arc>, + enable_range_cache: bool, } impl MetricSeriesFilter { fn new(assigned: &AssignedSeriesBatch) -> Self { - let series = assigned.series.iter().copied().collect(); + let series = assigned.series().iter().copied().collect(); Self { - range: assigned.range, + range: assigned.range(), series: Arc::new(series), + enable_range_cache: assigned.enable_range_cache(), } } @@ -232,7 +251,6 @@ fn filter_flat_stream_by_series( } /// Reads all collected metric series assigned to one partition. -#[allow(dead_code)] pub(crate) struct SeriesReader { stream_ctx: Arc, partition_ranges: Vec, @@ -244,7 +262,6 @@ pub(crate) struct SeriesReader { part_metrics: PartitionMetrics, } -#[allow(dead_code)] impl SeriesReader { /// Creates a reader for the series assigned to one data partition. /// @@ -354,7 +371,10 @@ async fn build_series_partition_range( partition_pruner: Arc, part_metrics: PartitionMetrics, ) -> Result<(BoxedRecordBatchStream, usize)> { - let cache_key = build_series_range_cache_key(&stream_ctx, &part_range, range); + let cache_key = filter + .enable_range_cache + .then(|| build_series_range_cache_key(&stream_ctx, &part_range, range)) + .flatten(); if let Some(key) = cache_key.as_ref() { if let Some(value) = stream_ctx.input.cache_strategy.get_range_result(key) { part_metrics.inc_range_cache_hit(); @@ -557,7 +577,7 @@ mod tests { ) -> AssignedSeriesBatch { let mut collector = SeriesBatchCollector::new(partitions).unwrap(); collector.push(series); - collector.finish().remove(partition) + collector.finish(true).remove(partition) } #[test] @@ -572,13 +592,13 @@ mod tests { ]; let mut first = SeriesBatchCollector::new(4).unwrap(); first.push(input.clone()); - let first = first.finish(); + let first = first.finish(true); let mut second = SeriesBatchCollector::new(4).unwrap(); for chunk in input.chunks(2) { second.push(chunk.to_vec()); } - let second = second.finish(); + let second = second.finish(true); assert_eq!( first @@ -601,6 +621,24 @@ mod tests { ); } + #[test] + fn collector_tracks_size_and_cache_policy() { + let mut collector = SeriesBatchCollector::new(2).unwrap(); + collector.push(vec![series(1, 1), series(1, u64::MAX)]); + collector.push(vec![series(2, 2)]); + assert_eq!(3, collector.len()); + + let assignments = collector.finish(false); + assert_eq!( + 3, + assignments + .iter() + .map(|batch| batch.series().len()) + .sum::() + ); + assert!(assignments.iter().all(|batch| !batch.enable_range_cache())); + } + #[test] fn series_ranges_cover_the_tsid_domain() { let ranges = (0..3) diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs index adefae91173..1a1369830f3 100644 --- a/src/mito2/src/read/series_scan.rs +++ b/src/mito2/src/read/series_scan.rs @@ -15,6 +15,7 @@ //! Per-series scan implementation. use std::fmt; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -45,13 +46,15 @@ use crate::error::{ ScanSeriesSnafu, TooManyFilesToReadSnafu, }; use crate::read::ScannerMetrics; -use crate::read::pruner::{PartitionPruner, Pruner}; +use crate::read::pruner::{PartitionPruner, Pruner, PrunerOptions}; use crate::read::scan_region::{ScanInput, StreamContext}; use crate::read::scan_util::{ PartitionMetrics, PartitionMetricsList, SeriesDistributorMetrics, compute_average_batch_size, compute_parallel_channel_size, }; use crate::read::seq_scan::SeqScan; +use crate::read::series_candidate::{SeriesCandidateScanner, is_sparse_metric_metadata}; +use crate::read::series_reader::{AssignedSeriesBatch, SeriesBatchCollector, SeriesReader}; use crate::read::stream::{ConvertBatchStream, ScanBatch, ScanBatchStream}; use crate::sst::parquet::flat_format::primary_key_column_index; use crate::sst::parquet::format::PrimaryKeyArray; @@ -59,8 +62,57 @@ use crate::sst::parquet::format::PrimaryKeyArray; /// Timeout to send a batch to a sender. const SEND_TIMEOUT: Duration = Duration::from_micros(100); -/// List of receivers. -type ReceiverList = Vec>>>; +/// Maximum number of candidate series retained before distributing assignments. +const CANDIDATE_SERIES_ASSIGNMENT_THRESHOLD: usize = 1_000_000; + +/// Legacy data receivers for output partitions. +type LegacyReceiverList = Vec>>>; + +/// Candidate assignment receivers for output partitions. +type CandidateReceiverList = Vec>; + +/// Candidate assignment senders for output partitions. +type CandidateSenderList = Vec>>>; + +struct CandidateReceiver { + receiver: mpsc::UnboundedReceiver>, + active_receivers: Arc, +} + +impl CandidateReceiver { + async fn recv(&mut self) -> Option> { + self.receiver.recv().await + } +} + +impl Drop for CandidateReceiver { + fn drop(&mut self) { + let previous = self.active_receivers.fetch_sub(1, Ordering::Relaxed); + debug_assert!(previous > 0); + } +} + +/// Input required by a partition-local series reader. +struct SeriesReaderInput { + assigned_series: AssignedSeriesBatch, + partition_pruner: Arc, + range_semaphore: Arc, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SeriesScanMode { + Legacy, + TwoPhase, +} + +impl SeriesScanMode { + fn as_str(self) -> &'static str { + match self { + Self::Legacy => "legacy", + Self::TwoPhase => "two_phase", + } + } +} /// Scans a region and returns sorted rows of a series in the same partition. /// @@ -68,14 +120,18 @@ type ReceiverList = Vec>>>; /// partition. /// Always returns the same series (primary key) to the same partition. pub struct SeriesScan { + /// Implementation used by this scan. + mode: SeriesScanMode, /// Properties of the scanner. properties: ScannerProperties, /// Context of streams. stream_ctx: Arc, /// Shared pruner for file range building. pruner: Arc, - /// Receivers of each partition. - receivers: Mutex, + /// Legacy data receivers for each partition. + legacy_receivers: Mutex, + /// Candidate assignment receivers for each partition. + candidate_receivers: Mutex, /// Metrics for each partition. /// The scanner only sets in query and keeps it empty during compaction. metrics_list: Arc, @@ -83,7 +139,12 @@ pub struct SeriesScan { impl SeriesScan { /// Creates a new [SeriesScan]. - pub(crate) fn new(input: ScanInput) -> Self { + pub(crate) fn new(input: ScanInput, experimental_series_scan_v2: bool) -> Self { + let mode = if experimental_series_scan_v2 && Self::supports_two_phase(&input) { + SeriesScanMode::TwoPhase + } else { + SeriesScanMode::Legacy + }; let mut properties = ScannerProperties::default() .with_append_mode(input.append_mode) .with_total_rows(input.total_rows()); @@ -95,17 +156,40 @@ impl SeriesScan { // Create the shared pruner with number of workers equal to CPU cores. let num_workers = common_stat::get_total_cpu_cores().max(1); - let pruner = Arc::new(Pruner::new(stream_ctx.clone(), num_workers)); + let pruner = match mode { + SeriesScanMode::Legacy => Arc::new(Pruner::new(stream_ctx.clone(), num_workers)), + SeriesScanMode::TwoPhase => Arc::new(Pruner::new_with_options( + stream_ctx.clone(), + num_workers, + PrunerOptions { + retain_builders: true, + enable_predicate_prefilter: false, + }, + )), + }; Self { + mode, properties, stream_ctx, pruner, - receivers: Mutex::new(Vec::new()), + legacy_receivers: Mutex::new(Vec::new()), + candidate_receivers: Mutex::new(Vec::new()), metrics_list: Arc::new(PartitionMetricsList::default()), } } + fn supports_two_phase(input: &ScanInput) -> bool { + if !is_sparse_metric_metadata(input.region_metadata()) { + return false; + } + #[cfg(feature = "enterprise")] + if !input.extension_ranges().is_empty() { + return false; + } + true + } + #[tracing::instrument( skip_all, fields( @@ -174,60 +258,128 @@ impl SeriesScan { } ); - self.maybe_start_distributor(metrics_set, &self.metrics_list, ctx.explain_verbose); + match self.mode { + SeriesScanMode::Legacy => self.scan_legacy_batch_in_partition( + partition, + part_metrics, + metrics_set, + ctx.explain_verbose, + ), + SeriesScanMode::TwoPhase => self.scan_two_phase_batch_in_partition( + partition, + part_metrics, + metrics_set, + ctx.explain_verbose, + ), + } + } - let mut receiver = self.take_receiver(partition)?; + fn scan_legacy_batch_in_partition( + &self, + partition: usize, + part_metrics: PartitionMetrics, + metrics_set: &ExecutionPlanMetricsSet, + explain_verbose: bool, + ) -> Result { + self.maybe_start_legacy_distributor(metrics_set, explain_verbose); + let mut receiver = self.legacy_receivers.lock().unwrap()[partition] + .take() + .context(ScanMultiTimesSnafu { partition })?; let stream = try_stream! { part_metrics.on_first_poll(); let mut fetch_start = Instant::now(); + let mut metrics = ScannerMetrics::default(); while let Some(series) = receiver.recv().await { let series = series?; - - let mut metrics = ScannerMetrics::default(); metrics.scan_cost += fetch_start.elapsed(); - fetch_start = Instant::now(); - metrics.num_batches += series.num_batches(); metrics.num_rows += series.num_rows(); let yield_start = Instant::now(); yield ScanBatch::Series(series); metrics.yield_cost += yield_start.elapsed(); - - part_metrics.merge_metrics(&metrics); + fetch_start = Instant::now(); } + part_metrics.merge_metrics(&metrics); part_metrics.on_finish(); }; Ok(Box::pin(stream)) } - /// Takes the receiver for the partition. - fn take_receiver(&self, partition: usize) -> Result>> { - let mut rx_list = self.receivers.lock().unwrap(); - rx_list[partition] + fn scan_two_phase_batch_in_partition( + &self, + partition: usize, + part_metrics: PartitionMetrics, + metrics_set: &ExecutionPlanMetricsSet, + explain_verbose: bool, + ) -> Result { + self.maybe_start_candidate_distributor(metrics_set, explain_verbose); + // Safety: `scan_batch_in_partition` validates the index, and the receiver + // list always contains one entry for each output partition. + let mut receiver = self.candidate_receivers.lock().unwrap()[partition] .take() - .context(ScanMultiTimesSnafu { partition }) + .context(ScanMultiTimesSnafu { partition })?; + let stream_ctx = self.stream_ctx.clone(); + let partition_ranges = self + .properties + .partitions + .iter() + .flatten() + .copied() + .collect::>(); + let stream = try_stream! { + part_metrics.on_first_poll(); + + let mut fetch_start = Instant::now(); + let mut metrics = ScannerMetrics::default(); + while let Some(input) = receiver.recv().await { + let input = input?; + metrics.scan_cost += fetch_start.elapsed(); + + let build_start = Instant::now(); + let reader = SeriesReader::try_new( + stream_ctx.clone(), + partition_ranges.clone(), + input.assigned_series, + input.partition_pruner, + input.range_semaphore, + part_metrics.clone(), + )?; + let mut reader_stream = reader.build_stream().await?; + metrics.scan_cost += build_start.elapsed(); + fetch_start = Instant::now(); + while let Some(record_batch) = reader_stream.try_next().await? { + metrics.scan_cost += fetch_start.elapsed(); + metrics.num_batches += 1; + metrics.num_rows += record_batch.num_rows(); + + let yield_start = Instant::now(); + yield ScanBatch::RecordBatch(record_batch); + metrics.yield_cost += yield_start.elapsed(); + fetch_start = Instant::now(); + } + } + metrics.scan_cost += fetch_start.elapsed(); + + part_metrics.merge_metrics(&metrics); + part_metrics.on_finish(); + }; + Ok(Box::pin(stream)) } - /// Starts the distributor if the receiver list is empty. - #[tracing::instrument( - skip(self, metrics_set, metrics_list), - fields(region_id = %self.stream_ctx.input.mapper.metadata().region_id) - )] - fn maybe_start_distributor( + fn maybe_start_legacy_distributor( &self, metrics_set: &ExecutionPlanMetricsSet, - metrics_list: &Arc, explain_verbose: bool, ) { - let mut rx_list = self.receivers.lock().unwrap(); + let mut rx_list = self.legacy_receivers.lock().unwrap(); if !rx_list.is_empty() { return; } - let (senders, receivers) = new_channel_list(self.properties.num_partitions()); + let (senders, receivers) = new_legacy_channel_list(self.properties.num_partitions()); let mut distributor = SeriesDistributor { stream_ctx: self.stream_ctx.clone(), range_semaphore: Some(Arc::new(Semaphore::new(self.properties.num_partitions()))), @@ -236,7 +388,7 @@ impl SeriesScan { pruner: self.pruner.clone(), senders, metrics_set: metrics_set.clone(), - metrics_list: metrics_list.clone(), + metrics_list: self.metrics_list.clone(), explain_verbose, }; let region_id = distributor.stream_ctx.input.mapper.metadata().region_id; @@ -251,6 +403,41 @@ impl SeriesScan { *rx_list = receivers; } + fn maybe_start_candidate_distributor( + &self, + metrics_set: &ExecutionPlanMetricsSet, + explain_verbose: bool, + ) { + let mut rx_list = self.candidate_receivers.lock().unwrap(); + if !rx_list.is_empty() { + return; + } + + let (senders, receivers, active_receivers) = + new_candidate_channel_list(self.properties.num_partitions()); + let mut distributor = SeriesCandidateDistributor { + stream_ctx: self.stream_ctx.clone(), + range_semaphore: Arc::new(Semaphore::new(self.properties.num_partitions())), + partitions: self.properties.partitions.clone(), + pruner: self.pruner.clone(), + senders, + active_receivers, + metrics_set: metrics_set.clone(), + metrics_list: self.metrics_list.clone(), + explain_verbose, + }; + let region_id = distributor.stream_ctx.input.mapper.metadata().region_id; + let span = tracing::info_span!("SeriesScan::candidate_distributor", region_id = %region_id); + common_runtime::spawn_query( + async move { + distributor.execute().await; + } + .instrument(span), + ); + + *rx_list = receivers; + } + /// Scans the region and returns a stream. #[tracing::instrument( skip_all, @@ -319,7 +506,7 @@ impl SeriesScan { } } -fn new_channel_list(num_partitions: usize) -> (SenderList, ReceiverList) { +fn new_legacy_channel_list(num_partitions: usize) -> (SenderList, LegacyReceiverList) { let (senders, receivers): (Vec<_>, Vec<_>) = (0..num_partitions) .map(|_| { let (sender, receiver) = mpsc::channel(1); @@ -329,6 +516,26 @@ fn new_channel_list(num_partitions: usize) -> (SenderList, ReceiverList) { (SenderList::new(senders), receivers) } +fn new_candidate_channel_list( + num_partitions: usize, +) -> (CandidateSenderList, CandidateReceiverList, Arc) { + let active_receivers = Arc::new(AtomicUsize::new(num_partitions)); + let (senders, receivers) = (0..num_partitions) + .map(|_| { + // Partition streams can be consumed sequentially. A bounded channel for an unpolled + // partition could block the distributor and prevent it from sending more assignments + // to the partition currently being consumed, causing a deadlock. + let (sender, receiver) = mpsc::unbounded_channel(); + let receiver = CandidateReceiver { + receiver, + active_receivers: active_receivers.clone(), + }; + (Some(sender), Some(receiver)) + }) + .unzip(); + (senders, receivers, active_receivers) +} + impl RegionScanner for SeriesScan { fn name(&self) -> &str { "SeriesScan" @@ -402,19 +609,24 @@ impl DisplayAs for SeriesScan { )?; match t { DisplayFormatType::Default | DisplayFormatType::TreeRender => { - self.stream_ctx.format_for_explain(false, f) + self.stream_ctx.format_for_explain(false, f)?; } DisplayFormatType::Verbose => { self.stream_ctx.format_for_explain(true, f)?; - self.metrics_list.format_verbose_metrics(f) } } + write!(f, ", \"mode\":\"{}\"", self.mode.as_str())?; + if matches!(t, DisplayFormatType::Verbose) { + self.metrics_list.format_verbose_metrics(f)?; + } + Ok(()) } } impl fmt::Debug for SeriesScan { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("SeriesScan") + .field("mode", &self.mode) .field("num_ranges", &self.stream_ctx.ranges.len()) .finish() } @@ -426,6 +638,11 @@ impl SeriesScan { pub(crate) fn input(&self) -> &ScanInput { &self.stream_ctx.input } + + /// Returns the scan mode for tests. + pub(crate) fn mode(&self) -> &'static str { + self.mode.as_str() + } } /// The distributor scans series and distributes them to different partitions. @@ -463,9 +680,7 @@ impl SeriesDistributor { fields(region_id = %self.stream_ctx.input.mapper.metadata().region_id) )] async fn execute(&mut self) { - let result = self.scan_partitions_flat().await; - - if let Err(e) = result { + if let Err(e) = self.scan_partitions_flat().await { self.senders.send_error(e).await; } } @@ -605,6 +820,113 @@ impl SeriesDistributor { } } +/// Discovers metric series and sends one reader assignment to each partition. +struct SeriesCandidateDistributor { + stream_ctx: Arc, + range_semaphore: Arc, + partitions: Vec>, + pruner: Arc, + senders: CandidateSenderList, + active_receivers: Arc, + metrics_set: ExecutionPlanMetricsSet, + metrics_list: Arc, + explain_verbose: bool, +} + +impl SeriesCandidateDistributor { + async fn execute(&mut self) { + if let Err(e) = self.distribute().await { + self.send_error(e); + } + } + + async fn distribute(&mut self) -> Result<()> { + let part_metrics = new_partition_metrics( + &self.stream_ctx, + self.explain_verbose, + &self.metrics_set, + self.partitions.len(), + &self.metrics_list, + ); + part_metrics.on_first_poll(); + + let candidate_scanner = SeriesCandidateScanner::try_new( + self.stream_ctx.clone(), + self.partitions.clone(), + self.pruner.clone(), + self.range_semaphore.clone(), + self.stream_ctx.input.scan_memory_pool.clone(), + self.metrics_set.clone(), + part_metrics.clone(), + )?; + let partition_pruner = candidate_scanner.partition_pruner(); + let mut candidates = candidate_scanner.build_stream().await?; + let mut collector = + SeriesBatchCollector::new(self.partitions.len()).context(InvalidSenderSnafu)?; + let mut chunked = false; + while let Some(batch) = candidates.try_next().await? { + if !self.should_fetch_candidates() { + part_metrics.on_finish(); + return Ok(()); + } + collector.push(batch); + if collector.len() >= CANDIDATE_SERIES_ASSIGNMENT_THRESHOLD { + chunked = true; + self.send_assignments(collector.finish(false), &partition_pruner); + if !self.should_fetch_candidates() { + part_metrics.on_finish(); + return Ok(()); + } + collector = + SeriesBatchCollector::new(self.partitions.len()).context(InvalidSenderSnafu)?; + } + } + + if collector.len() > 0 { + self.send_assignments(collector.finish(!chunked), &partition_pruner); + } + part_metrics.on_finish(); + Ok(()) + } + + fn send_assignments( + &mut self, + assignments: Vec, + partition_pruner: &Arc, + ) { + for (partition, assigned_series) in assignments.into_iter().enumerate() { + if assigned_series.series().is_empty() { + continue; + } + let Some(sender) = self.senders[partition].as_ref() else { + continue; + }; + let sent = sender + .send(Ok(SeriesReaderInput { + assigned_series, + partition_pruner: partition_pruner.clone(), + range_semaphore: self.range_semaphore.clone(), + })) + .is_ok(); + if !sent { + self.senders[partition] = None; + } + } + } + + fn should_fetch_candidates(&self) -> bool { + self.active_receivers.load(Ordering::Relaxed) > 0 + } + + fn send_error(&mut self, error: Error) { + let error = Arc::new(error); + for sender in self.senders.iter_mut().filter_map(Option::take) { + let result = Err(error.clone()).context(ScanSeriesSnafu); + let _ = sender.send(result); + } + } +} + /// Batches of the same series. #[derive(Debug)] pub enum SeriesBatch { @@ -893,6 +1215,18 @@ mod tests { use super::*; + #[test] + fn candidate_distributor_stops_after_all_receivers_close() { + let (_senders, mut receivers, active_receivers) = new_candidate_channel_list(2); + assert_eq!(2, active_receivers.load(Ordering::Relaxed)); + + drop(receivers[0].take()); + assert_eq!(1, active_receivers.load(Ordering::Relaxed)); + + drop(receivers[1].take()); + assert_eq!(0, active_receivers.load(Ordering::Relaxed)); + } + fn new_test_record_batch( primary_keys: &[&[u8]], timestamps: &[i64], diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index b2c8fe394e2..530b581a2ca 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2258,6 +2258,7 @@ scan_memory_on_exhausted = "fail" min_compaction_interval = "0s" schedule_compaction_after_edit = true default_flat_format = true +experimental_series_scan_v2 = true [region_engine.mito.index] aux_path = "" diff --git a/tests/cases/distributed/explain/step_aggr_advance.result b/tests/cases/distributed/explain/step_aggr_advance.result index 9fc8dfe5a1c..59f2946e8bc 100644 --- a/tests/cases/distributed/explain/step_aggr_advance.result +++ b/tests/cases/distributed/explain/step_aggr_advance.result @@ -67,7 +67,7 @@ tql analyze (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST], preserve_partitioning=[true] REDACTED @@ -79,7 +79,7 @@ tql analyze (1752591864, 1752592164, '30s') max by (a, b, c) (max_over_time(aggr |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -144,7 +144,7 @@ tql analyze (1752591864, 1752592164, '30s') sum by (a, b) (max_over_time(aggr_op |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, greptime_timestamp@2 as greptime_timestamp], aggr=[__sum_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED @@ -154,7 +154,7 @@ tql analyze (1752591864, 1752592164, '30s') sum by (a, b) (max_over_time(aggr_op |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -219,7 +219,7 @@ tql analyze (1752591864, 1752592164, '30s') avg by (a) (max_over_time(aggr_optim |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, greptime_timestamp@1 as greptime_timestamp], aggr=[__avg_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED @@ -229,7 +229,7 @@ tql analyze (1752591864, 1752592164, '30s') avg by (a) (max_over_time(aggr_optim |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -284,7 +284,7 @@ tql analyze (1752591864, 1752592164, '30s') count by (a, b, c, d) (max_over_time |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, d@3 ASC NULLS LAST, greptime_timestamp@4 ASC NULLS LAST], preserve_partitioning=[true] REDACTED @@ -294,7 +294,7 @@ tql analyze (1752591864, 1752592164, '30s') count by (a, b, c, d) (max_over_time |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -359,7 +359,7 @@ tql analyze (1752591864, 1752592164, '30s') min by (b, c, d) (max_over_time(aggr |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, c@1 as c, d@2 as d, greptime_timestamp@3 as greptime_timestamp], aggr=[__min_state(prom_max_over_time(greptime_timestamp_range,greptime_value))] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED @@ -369,7 +369,7 @@ tql analyze (1752591864, 1752592164, '30s') min by (b, c, d) (max_over_time(aggr |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -429,7 +429,7 @@ tql analyze sum(aggr_optimize_not); |_|_|_ProjectionExec: expr=[greptime_timestamp@4 as greptime_timestamp, greptime_value@5 as greptime_value] REDACTED |_|_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[300000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[__sum_state(aggr_optimize_not.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED @@ -437,7 +437,7 @@ tql analyze sum(aggr_optimize_not); |_|_|_ProjectionExec: expr=[greptime_timestamp@4 as greptime_timestamp, greptime_value@5 as greptime_value] REDACTED |_|_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[300000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -563,7 +563,7 @@ tql analyze (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c, greptime_timestamp@3 as greptime_timestamp, sum(prom_rate(greptime_timestamp_range,greptime_value,aggr_optimize_not.greptime_timestamp,Int64(120000)))@4 as sum(prom_rate(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(120000)))] REDACTED |_|_|_SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST, c@2 ASC NULLS LAST, greptime_timestamp@3 ASC NULLS LAST] REDACTED @@ -576,7 +576,7 @@ tql analyze (1752591864, 1752592164, '30s') sum by (a, b, c) (rate(aggr_optimize |_|_|_PromRangeManipulateExec: req range=[1752591864000..1752592164000], interval=[30000], eval range=[120000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[greptime_timestamp], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["a", "b", "c", "d"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/distributed/tql-explain-analyze/analyze.result b/tests/cases/distributed/tql-explain-analyze/analyze.result index bdd7c4811e1..059cf55df64 100644 --- a/tests/cases/distributed/tql-explain-analyze/analyze.result +++ b/tests/cases/distributed/tql-explain-analyze/analyze.result @@ -26,7 +26,7 @@ TQL ANALYZE (0, 10, '5s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -51,7 +51,7 @@ TQL ANALYZE (0, 10, '1s', '2s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[2000], interval=[1000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -75,7 +75,7 @@ TQL ANALYZE ('1970-01-01T00:00:00'::timestamp, '1970-01-01T00:00:00'::timestamp | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -134,12 +134,12 @@ TQL ANALYZE (0, 10, '5s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -166,7 +166,7 @@ TQL ANALYZE (0, 10, '5s') rate(test[10s]); |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_ProjectionExec: expr=[j@0 as j, prom_rate(j_range,i,test.j,Int64(10000))@1 as prom_rate(j_range,i,j,Int64(10000)), k@2 as k, l@3 as l] REDACTED |_|_|_FilterExec: prom_rate(j_range,i,test.j,Int64(10000))@1 IS NOT NULL REDACTED @@ -175,7 +175,7 @@ TQL ANALYZE (0, 10, '5s') rate(test[10s]); |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/standalone/common/promql/histogram_multi_partition.result b/tests/cases/standalone/common/promql/histogram_multi_partition.result index 3d5490f38a1..93168ce449f 100644 --- a/tests/cases/standalone/common/promql/histogram_multi_partition.result +++ b/tests/cases/standalone/common/promql/histogram_multi_partition.result @@ -56,7 +56,7 @@ tql analyze (0, 10, '10s') histogram_quantile(0.5, sum by (le) (histogram_gap_bu |_|_|_ProjectionExec: expr=[ts@0 as ts, le@1 as le, val@3 as val] REDACTED |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[10000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["shard", "le"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[le@0 as le, ts@1 as ts], aggr=[__sum_state(histogram_gap_bucket.val)] REDACTED |_|_|_RepartitionExec: partitioning=Hash([le@0, ts@1],REDACTED @@ -64,7 +64,7 @@ tql analyze (0, 10, '10s') histogram_quantile(0.5, sum by (le) (histogram_gap_bu |_|_|_ProjectionExec: expr=[ts@0 as ts, le@1 as le, val@3 as val] REDACTED |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[10000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["shard", "le"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ diff --git a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result index 72a6dc9c5b4..a3983260e33 100644 --- a/tests/cases/standalone/common/promql/tsid_binary_join_regression.result +++ b/tests/cases/standalone/common/promql/tsid_binary_join_regression.result @@ -127,13 +127,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / tsid_binary_join_right; |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -167,13 +167,13 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left + tsid_binary_join_right) / tsid |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -210,19 +210,19 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left + tsid_binary_join_right) * (ts |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -255,13 +255,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / ignoring(host) tsid_binary_join |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -294,13 +294,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / on(job) tsid_binary_join_right_ |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, job@3 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -334,13 +334,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > tsid_binary_join_right; |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -373,13 +373,13 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > bool tsid_binary_join_right; |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -418,19 +418,19 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > tsid_binary_join_right) / tsi |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -473,25 +473,25 @@ TQL ANALYZE (0, 5, '5s') ((tsid_binary_join_left > bool tsid_binary_join_right) |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -529,7 +529,7 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[ts@2 as ts, __tsid@3 as __tsid, greptime_value@4 as greptime_value, host@5 as host, job@6 as job, CASE WHEN __common_expr_1@0 IS NOT NULL THEN __common_expr_1@0 ELSE_END as __promql_or_match_0, CASE WHEN __common_expr_2@1 IS NOT NULL THEN __common_expr_2@1 ELSE_END as __promql_or_match_1] REDACTED |_|_|_ProjectionExec: expr=[CAST(host@1 AS Utf8) as __common_expr_1, CAST(job@2 AS Utf8) as __common_expr_2, ts@4 as ts, __tsid@3 as __tsid, greptime_value@0 as greptime_value, host@1 as host, job@2 as job] REDACTED @@ -537,13 +537,13 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left or tsid_binary_join_right) / tsi |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -582,19 +582,19 @@ TQL ANALYZE (0, 5, '5s') (tsid_binary_join_left / ignoring(host) group_left tsid |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ diff --git a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result index ba2e97bcc2d..af5f3a9a4f8 100644 --- a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result +++ b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result @@ -118,7 +118,7 @@ TQL ANALYZE (0, 10, '5s') histogram_quantile(0.5, tsid_no_aggr_histogram_bucket) |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, job@3 as job, le@4 as le, __tsid@2 as __tsid, ts@0 as ts] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ diff --git a/tests/cases/standalone/common/select/tql_filter.result b/tests/cases/standalone/common/select/tql_filter.result index 345671376f7..b4fbfdb5862 100644 --- a/tests/cases/standalone/common/select/tql_filter.result +++ b/tests/cases/standalone/common/select/tql_filter.result @@ -24,7 +24,7 @@ tql analyze (1, 3, '1s') t1{ a = "a" }; | 1_| 0_|_PromInstantManipulateExec: range=[1000..3000], lookback=[300000], interval=[1000], time index=[b] REDACTED |_|_|_PromSeriesDivideExec: tags=["a"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ @@ -47,7 +47,7 @@ tql analyze (1, 3, '1s') t1{ a =~ ".*" }; | 1_| 0_|_PromInstantManipulateExec: range=[1000..3000], lookback=[300000], interval=[1000], time index=[b] REDACTED |_|_|_PromSeriesDivideExec: tags=["a"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 6_| +-+-+-+ @@ -70,7 +70,7 @@ tql analyze (1, 3, '1s') t1{ a =~ "a.*" }; | 1_| 0_|_PromInstantManipulateExec: range=[1000..3000], lookback=[300000], interval=[1000], time index=[b] REDACTED |_|_|_PromSeriesDivideExec: tags=["a"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ diff --git a/tests/cases/standalone/common/tql/partition.result b/tests/cases/standalone/common/tql/partition.result index f1072edbb32..e32d13043bd 100644 --- a/tests/cases/standalone/common/tql/partition.result +++ b/tests/cases/standalone/common/tql/partition.result @@ -32,7 +32,7 @@ tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); |_|_|_PromRangeManipulateExec: req range=[0..10000], interval=[1000], eval range=[60000], time index=[j] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -82,7 +82,7 @@ tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); |_|_|_PromRangeManipulateExec: req range=[0..10000], interval=[1000], eval range=[60000], time index=[j] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[k@0 as k, j@1 as j], aggr=[__avg_state(prom_irate(j_range,i))] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED @@ -92,7 +92,7 @@ tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); |_|_|_PromRangeManipulateExec: req range=[0..10000], interval=[1000], eval range=[60000], time index=[j] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -141,10 +141,10 @@ tql analyze (0, 10, '1s') 100 - (avg by (k) (irate(t[1m])) * 100); |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [k@2 ASC, l@3 ASC, j@1 ASC] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_SortPreservingMergeExec: [k@2 ASC, l@3 ASC, j@1 ASC] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/standalone/optimizer/windowed_sort.result b/tests/cases/standalone/optimizer/windowed_sort.result index 544563ad1a7..3be0dd1c3e4 100644 --- a/tests/cases/standalone/optimizer/windowed_sort.result +++ b/tests/cases/standalone/optimizer/windowed_sort.result @@ -290,7 +290,7 @@ TQL ANALYZE('2023-10-01 00:00:00.000'::TIMESTAMP, '2023-10-01 00:00:05.000'::TIM |_|_|_ProjectionExec: expr=[greptime_timestamp@0 as greptime_timestamp, greptime_value@1 as greptime_value] REDACTED |_|_|_PromInstantManipulateExec: range=[1696118400000..1696118405000], lookback=[300000], interval=[1000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesDivideExec: tags=["instance", "job"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 5_| +-+-+-+ diff --git a/tests/cases/standalone/tql-explain-analyze/analyze.result b/tests/cases/standalone/tql-explain-analyze/analyze.result index 203debf90bf..954020ea4cf 100644 --- a/tests/cases/standalone/tql-explain-analyze/analyze.result +++ b/tests/cases/standalone/tql-explain-analyze/analyze.result @@ -26,7 +26,7 @@ TQL ANALYZE (0, 10, '5s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -51,7 +51,7 @@ TQL ANALYZE (0, 10, '1s', '2s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[2000], interval=[1000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -75,7 +75,7 @@ TQL ANALYZE ('1970-01-01T00:00:00'::timestamp, '1970-01-01T00:00:00'::timestamp | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 4_| +-+-+-+ @@ -134,12 +134,12 @@ TQL ANALYZE (0, 10, '5s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -165,7 +165,7 @@ TQL ANALYZE (0, 10, '5s') rate(test[10s]); |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_FilterExec: prom_rate(j_range,i,j,Int64(10000))@1 IS NOT NULL REDACTED |_|_|_ProjectionExec: expr=[j@1 as j, prom_rate(j_range@4, i@0, j@1, 10000) as prom_rate(j_range,i,j,Int64(10000)), k@2 as k, l@3 as l] REDACTED @@ -173,7 +173,7 @@ TQL ANALYZE (0, 10, '5s') rate(test[10s]); |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[j], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -238,12 +238,12 @@ TQL ANALYZE FORMAT TEXT (0, 10, '5s') test; | 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED |_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED |_|_|_CooperativeExec REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ @@ -297,7 +297,7 @@ TQL ANALYZE sum(test2); |_|_|_ProjectionExec: expr=[greptime_timestamp@0 as greptime_timestamp, greptime_value@1 as greptime_value] REDACTED |_|_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[300000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesDivideExec: tags=["shard"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| | 1_| 1_|_AggregateExec: mode=FinalPartitioned, gby=[greptime_timestamp@0 as greptime_timestamp], aggr=[__sum_state(test2.greptime_value)] REDACTED |_|_|_RepartitionExec: partitioning=REDACTED @@ -305,7 +305,7 @@ TQL ANALYZE sum(test2); |_|_|_ProjectionExec: expr=[greptime_timestamp@0 as greptime_timestamp, greptime_value@1 as greptime_value] REDACTED |_|_|_PromInstantManipulateExec: range=[0..0], lookback=[300000], interval=[300000], time index=[greptime_timestamp] REDACTED |_|_|_PromSeriesDivideExec: tags=["shard"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/standalone/tql-explain-analyze/column_pruning.result b/tests/cases/standalone/tql-explain-analyze/column_pruning.result index 4bce77d7832..cecb72fcee1 100644 --- a/tests/cases/standalone/tql-explain-analyze/column_pruning.result +++ b/tests/cases/standalone/tql-explain-analyze/column_pruning.result @@ -45,7 +45,7 @@ TQL ANALYZE (0, 10, '5s') sum(promql_column_pruning); |_|_|_ProjectionExec: expr=[ts@0 as ts, greptime_value@4 as greptime_value] REDACTED |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["job", "instance", "region"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ @@ -75,7 +75,7 @@ TQL ANALYZE (0, 10, '5s') sum(rate(promql_column_pruning[5s])); |_|_|_PromRangeManipulateExec: req range=[0..10000], interval=[5000], eval range=[5000], time index=[ts] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[ts], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["job", "instance", "region"] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"legacy" REDACTED |_|_|_| |_|_| Total rows: 0_| +-+-+-+ diff --git a/tests/cases/standalone/tql-explain-analyze/tsid_column.result b/tests/cases/standalone/tql-explain-analyze/tsid_column.result index 8d27a55ab38..cf2deb76647 100644 --- a/tests/cases/standalone/tql-explain-analyze/tsid_column.result +++ b/tests/cases/standalone/tql-explain-analyze/tsid_column.result @@ -54,7 +54,7 @@ TQL ANALYZE (0, 10, '5s') sum(tsid_metric); |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 3_| +-+-+-+ @@ -82,7 +82,7 @@ TQL ANALYZE (0, 10, '5s') sum by (job, instance) (tsid_metric); |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, instance@3 as instance, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 6_| +-+-+-+ @@ -120,7 +120,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(count(tsid |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, job@3 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ts@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED @@ -133,7 +133,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(count(tsid |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[ts], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ @@ -173,7 +173,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(sum(tsid_m |_|_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, job@3 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| | 1_| 0_|_SortPreservingMergeExec: [ts@0 ASC NULLS LAST] REDACTED |_|_|_SortExec: expr=[ts@0 ASC NULLS LAST], preserve_partitioning=[true] REDACTED @@ -186,7 +186,7 @@ TQL ANALYZE (0, 10, '5s') sum(irate(tsid_metric[1h])) / scalar(count(sum(tsid_m |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[ts], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED |_|_|_ProjectionExec: expr=[val@1 as val, __tsid@2 as __tsid, ts@0 as ts] REDACTED -|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "distribution":"PerSeries", "mode":"two_phase" REDACTED |_|_|_| |_|_| Total rows: 2_| +-+-+-+ From c3cd186997538e2ed665ed32ce1755baf62bd857 Mon Sep 17 00:00:00 2001 From: Morax <100508620+fzlzjerry@users.noreply.github.com> Date: Fri, 14 Aug 2026 07:04:51 +0000 Subject: [PATCH 48/92] fix(query): respect query timezone in timestamp casts (#8859) * fix(query): respect timezone in insert values Signed-off-by: Morax * fix(query): normalize timestamp casts with query timezone Signed-off-by: Morax * fix(query): scope timestamp conversion to insert assignments Signed-off-by: Morax * style(query): simplify Arc usage Signed-off-by: Morax --------- Signed-off-by: Morax --- src/query/src/optimizer/type_conversion.rs | 284 +++++++++++++++--- src/query/src/planner.rs | 70 ++++- .../insert/insert_default_timezone.result | 39 +++ .../common/insert/insert_default_timezone.sql | 18 ++ 4 files changed, 370 insertions(+), 41 deletions(-) diff --git a/src/query/src/optimizer/type_conversion.rs b/src/query/src/optimizer/type_conversion.rs index 1702aacce83..366a7d8a31d 100644 --- a/src/query/src/optimizer/type_conversion.rs +++ b/src/query/src/optimizer/type_conversion.rs @@ -12,14 +12,16 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; + use common_time::Timezone; use common_time::timestamp::{TimeUnit, Timestamp}; use datafusion::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; use datafusion_common::{DFSchemaRef, DataFusionError, Result, ScalarValue}; -use datafusion_expr::expr::InList; +use datafusion_expr::expr::{Cast, InList}; use datafusion_expr::{ - Between, BinaryExpr, Expr, ExprSchemable, Filter, LogicalPlan, Operator, TableScan, + Between, BinaryExpr, Expr, ExprSchemable, Filter, LogicalPlan, Operator, TableScan, WriteOp, }; use datatypes::arrow::compute; use datatypes::arrow::datatypes::DataType; @@ -32,7 +34,7 @@ use crate::plan::ExtractExpr; /// TypeConversionRule converts some literal values in logical plan to other types according /// to data type of corresponding columns. /// Specifically: -/// - string literal of timestamp is converted to `Expr::Literal(ScalarValue::TimestampMillis)` +/// - string literal of timestamp is converted to the target timestamp type /// - string literal of boolean is converted to `Expr::Literal(ScalarValue::Boolean)` pub struct TypeConversionRule; @@ -45,10 +47,8 @@ impl ExtensionAnalyzerRule for TypeConversionRule { ) -> Result { plan.transform_up_with_subqueries(|plan| match plan { LogicalPlan::Filter(filter) => { - let mut converter = TypeConverter { - schema: filter.input.schema().clone(), - query_ctx: ctx.query_ctx(), - }; + let mut converter = + TypeConverter::new(filter.input.schema().clone(), ctx.query_ctx()); let rewritten = filter.predicate.clone().rewrite(&mut converter)?.data; Ok(Transformed::yes(LogicalPlan::Filter(Filter::try_new( rewritten, @@ -63,10 +63,7 @@ impl ExtensionAnalyzerRule for TypeConversionRule { filters, fetch, }) => { - let mut converter = TypeConverter { - schema: projected_schema.clone(), - query_ctx: ctx.query_ctx(), - }; + let mut converter = TypeConverter::new(projected_schema.clone(), ctx.query_ctx()); let rewrite_filters = filters .into_iter() .map(|e| e.rewrite(&mut converter).map(|x| x.data)) @@ -80,8 +77,18 @@ impl ExtensionAnalyzerRule for TypeConversionRule { fetch, }))) } - LogicalPlan::Projection { .. } - | LogicalPlan::Window { .. } + LogicalPlan::Projection { .. } => { + let mut converter = TypeConverter::new(plan.schema().clone(), ctx.query_ctx()); + let inputs = plan.inputs().into_iter().cloned().collect::>(); + let expr = plan + .expressions_consider_join() + .into_iter() + .map(|e| e.rewrite(&mut converter).map(|x| x.data)) + .collect::>>()?; + + plan.with_new_exprs(expr, inputs).map(Transformed::yes) + } + LogicalPlan::Window { .. } | LogicalPlan::Aggregate { .. } | LogicalPlan::Repartition { .. } | LogicalPlan::Extension { .. } @@ -89,10 +96,7 @@ impl ExtensionAnalyzerRule for TypeConversionRule { | LogicalPlan::Union { .. } | LogicalPlan::Values { .. } | LogicalPlan::Analyze { .. } => { - let mut converter = TypeConverter { - schema: plan.schema().clone(), - query_ctx: ctx.query_ctx(), - }; + let mut converter = TypeConverter::new(plan.schema().clone(), ctx.query_ctx()); let inputs = plan.inputs().into_iter().cloned().collect::>(); let expr = plan .expressions_consider_join() @@ -107,10 +111,7 @@ impl ExtensionAnalyzerRule for TypeConversionRule { let Ok(schema) = join.left.schema().join(join.right.schema()) else { return Ok(Transformed::no(LogicalPlan::Join(join))); }; - let mut converter = TypeConverter { - schema: std::sync::Arc::new(schema), - query_ctx: ctx.query_ctx(), - }; + let mut converter = TypeConverter::new(Arc::new(schema), ctx.query_ctx()); let plan = LogicalPlan::Join(join); let inputs = plan.inputs().into_iter().cloned().collect::>(); let expr = plan @@ -122,6 +123,14 @@ impl ExtensionAnalyzerRule for TypeConversionRule { plan.with_new_exprs(expr, inputs).map(Transformed::yes) } + LogicalPlan::Dml(mut dml) if matches!(dml.op, WriteOp::Insert(_)) => { + dml.input = Arc::new(rewrite_insert_assignments( + dml.input.as_ref().clone(), + ctx.query_ctx(), + )?); + Ok(Transformed::yes(LogicalPlan::Dml(dml))) + } + LogicalPlan::Distinct { .. } | LogicalPlan::Limit { .. } | LogicalPlan::Subquery { .. } @@ -146,6 +155,10 @@ struct TypeConverter { } impl TypeConverter { + fn new(schema: DFSchemaRef, query_ctx: QueryContextRef) -> Self { + Self { query_ctx, schema } + } + fn column_type(&self, expr: &Expr) -> Option { if let Expr::Column(_) = expr && let Ok(v) = expr.get_type(&self.schema) @@ -162,7 +175,7 @@ impl TypeConverter { ) -> Result { match (target_type, value) { (DataType::Timestamp(_, _), ScalarValue::Utf8(Some(v))) => { - string_to_timestamp_ms(v, Some(&self.query_ctx.timezone())) + parse_string_to_timestamp(v, Some(&self.query_ctx.timezone())) } (DataType::Boolean, ScalarValue::Utf8(Some(v))) => match v.to_lowercase().as_str() { "true" => Ok(ScalarValue::Boolean(Some(true))), @@ -294,6 +307,120 @@ impl TreeNodeRewriter for TypeConverter { } } +/// Rewrites timestamp assignment casts at the INSERT boundary. The source query +/// is intentionally left untouched so explicit casts keep DataFusion semantics. +fn rewrite_insert_assignments( + plan: LogicalPlan, + query_ctx: QueryContextRef, +) -> Result { + let LogicalPlan::Projection(projection) = plan else { + return Ok(plan); + }; + + // VALUES type coercion is performed inside the Values node. Only rewrite + // the direct source of the assignment projection, not nested query plans. + let input = rewrite_insert_values(projection.input.as_ref().clone(), query_ctx.clone())?; + let input_literals = projected_literals(&input); + let mut converter = InsertAssignmentConverter { + query_ctx, + input_literals, + }; + let plan = LogicalPlan::Projection(projection); + let expr = plan + .expressions_consider_join() + .into_iter() + .map(|expr| expr.rewrite(&mut converter).map(|x| x.data)) + .collect::>>()?; + + plan.with_new_exprs(expr, vec![input]) +} + +fn rewrite_insert_values(plan: LogicalPlan, query_ctx: QueryContextRef) -> Result { + let LogicalPlan::Values(mut values) = plan else { + return Ok(plan); + }; + let mut converter = InsertAssignmentConverter { + query_ctx, + input_literals: None, + }; + values.values = values + .values + .into_iter() + .map(|row| { + row.into_iter() + .map(|expr| expr.rewrite(&mut converter).map(|x| x.data)) + .collect::>>() + }) + .collect::>>()?; + + Ok(LogicalPlan::Values(values)) +} + +fn projected_literals(plan: &LogicalPlan) -> Option<(DFSchemaRef, Vec>)> { + let LogicalPlan::Projection(projection) = plan else { + return None; + }; + let literals = projection.expr.iter().map(projected_literal).collect(); + Some((projection.schema.clone(), literals)) +} + +fn projected_literal(expr: &Expr) -> Option { + match expr { + Expr::Literal(value, _) => Some(value.clone()), + Expr::Alias(alias) => projected_literal(&alias.expr), + _ => None, + } +} + +struct InsertAssignmentConverter { + query_ctx: QueryContextRef, + input_literals: Option<(DFSchemaRef, Vec>)>, +} + +impl InsertAssignmentConverter { + fn assignment_literal(&self, expr: &Expr) -> Option { + let Expr::Column(column) = expr else { + return expr.as_literal().cloned(); + }; + let (schema, literals) = self.input_literals.as_ref()?; + let index = schema.maybe_index_of_column(column)?; + literals.get(index)?.clone() + } +} + +impl TreeNodeRewriter for InsertAssignmentConverter { + type Node = Expr; + + fn f_up(&mut self, expr: Expr) -> Result> { + let Expr::Cast(Cast { + expr: inner, + data_type: target_type, + }) = expr + else { + return Ok(Transformed::no(expr)); + }; + let original = || { + Expr::Cast(Cast { + expr: inner.clone(), + data_type: target_type.clone(), + }) + }; + if !matches!(&target_type, DataType::Timestamp(_, _)) { + return Ok(Transformed::no(original())); + } + let Some(ScalarValue::Utf8(Some(value))) = self.assignment_literal(inner.as_ref()) else { + return Ok(Transformed::no(original())); + }; + + match cast_string_to_timestamp(&value, &target_type, Some(&self.query_ctx.timezone())) { + Ok(value) if !value.is_null() => Ok(Transformed::yes(Expr::Literal(value, None))), + // Preserve DataFusion's existing parser for unsupported timestamp + // syntax and out-of-range values. + Ok(_) | Err(_) => Ok(Transformed::no(original())), + } + } +} + fn timestamp_to_timestamp_ms_expr(val: i64, unit: TimeUnit) -> Expr { let timestamp = match unit { TimeUnit::Second => val * 1_000, @@ -308,18 +435,34 @@ fn timestamp_to_timestamp_ms_expr(val: i64, unit: TimeUnit) -> Expr { ) } -fn string_to_timestamp_ms(string: &str, timezone: Option<&Timezone>) -> Result { +fn cast_string_to_timestamp( + string: &str, + target_type: &DataType, + timezone: Option<&Timezone>, +) -> Result { + let parsed = parse_string_to_timestamp(string, timezone)?; + cast_timestamp(parsed, target_type) +} + +fn parse_string_to_timestamp(string: &str, timezone: Option<&Timezone>) -> Result { let ts = Timestamp::from_str(string, timezone) .map_err(|e| DataFusionError::External(Box::new(e)))?; let value = Some(ts.value()); - let scalar = match ts.unit() { + Ok(match ts.unit() { TimeUnit::Second => ScalarValue::TimestampSecond(value, None), TimeUnit::Millisecond => ScalarValue::TimestampMillisecond(value, None), TimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(value, None), TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(value, None), - }; - Ok(scalar) + }) +} + +fn cast_timestamp(parsed: ScalarValue, target_type: &DataType) -> Result { + let parsed = parsed.to_array()?; + let casted = compute::cast(&parsed, target_type) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; + + ScalarValue::try_from_array(&casted, 0) } #[cfg(test)] @@ -337,19 +480,23 @@ mod tests { use super::*; #[test] - fn test_string_to_timestamp_ms() { + fn test_cast_string_to_timestamp() { + use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit; + + let target_type = DataType::Timestamp(ArrowTimeUnit::Second, None); assert_eq!( - string_to_timestamp_ms("2022-02-02 19:00:00+08:00", None).unwrap(), + cast_string_to_timestamp("2022-02-02 19:00:00+08:00", &target_type, None).unwrap(), ScalarValue::TimestampSecond(Some(1643799600), None) ); assert_eq!( - string_to_timestamp_ms("2009-02-13 23:31:30Z", None).unwrap(), + cast_string_to_timestamp("2009-02-13 23:31:30Z", &target_type, None).unwrap(), ScalarValue::TimestampSecond(Some(1234567890), None) ); assert_eq!( - string_to_timestamp_ms( + cast_string_to_timestamp( "2009-02-13 23:31:30", + &target_type, Some(&Timezone::from_tz_string("Asia/Shanghai").unwrap()) ) .unwrap(), @@ -357,13 +504,24 @@ mod tests { ); assert_eq!( - string_to_timestamp_ms( + cast_string_to_timestamp( "2009-02-13 23:31:30", + &target_type, Some(&Timezone::from_tz_string("-8:00").unwrap()) ) .unwrap(), ScalarValue::TimestampSecond(Some(1234567890 + 8 * 3600), None) ); + + assert_eq!( + cast_string_to_timestamp( + "2009-02-13 23:31:30.123456789Z", + &DataType::Timestamp(ArrowTimeUnit::Nanosecond, None), + None, + ) + .unwrap(), + ScalarValue::TimestampNanosecond(Some(1_234_567_890_123_456_789), None) + ); } #[test] @@ -421,14 +579,11 @@ mod tests { ) .unwrap(), ); - let mut converter = TypeConverter { - schema, - query_ctx: QueryContext::arc(), - }; + let mut converter = TypeConverter::new(schema, QueryContext::arc()); assert_eq!( Expr::Column(Column::from_name("ts")).gt(ScalarValue::TimestampSecond( - Some(1599514949), + Some(1_599_514_949), None ) .lit()), @@ -439,6 +594,58 @@ mod tests { ); } + #[test] + fn test_insert_assignment_uses_query_timezone_and_target_precision() { + use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit; + + let query_ctx = QueryContext::arc(); + query_ctx.set_timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()); + let mut converter = InsertAssignmentConverter { + query_ctx, + input_literals: None, + }; + + let expr = Expr::Cast(Cast { + expr: Box::new("2009-02-13 23:31:30.123456789".lit()), + data_type: DataType::Timestamp(ArrowTimeUnit::Nanosecond, None), + }); + assert_eq!( + converter.f_up(expr).unwrap().data, + ScalarValue::TimestampNanosecond(Some(1_234_539_090_123_456_789), None).lit() + ); + } + + #[test] + fn test_insert_assignment_falls_back_for_unsupported_literal() { + use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit; + + let mut converter = InsertAssignmentConverter { + query_ctx: QueryContext::arc(), + input_literals: None, + }; + for literal in ["1970-01-01", "-8-01-01 00:00:01.5"] { + let expr = Expr::Cast(Cast { + expr: Box::new(literal.lit()), + data_type: DataType::Timestamp(ArrowTimeUnit::Nanosecond, None), + }); + + assert_eq!(converter.f_up(expr.clone()).unwrap().data, expr); + } + } + + #[test] + fn test_type_converter_leaves_explicit_timestamp_cast_unchanged() { + use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit; + + let mut converter = TypeConverter::new(Arc::new(DFSchema::empty()), QueryContext::arc()); + let expr = Expr::Cast(Cast { + expr: Box::new("2009-02-13 23:31:30".lit()), + data_type: DataType::Timestamp(ArrowTimeUnit::Millisecond, None), + }); + + assert_eq!(converter.f_up(expr.clone()).unwrap().data, expr); + } + #[test] fn test_convert_bool() { let col_name = "is_valid"; @@ -452,10 +659,7 @@ mod tests { ) .unwrap(), ); - let mut converter = TypeConverter { - schema, - query_ctx: QueryContext::arc(), - }; + let mut converter = TypeConverter::new(schema, QueryContext::arc()); assert_eq!( Expr::Column(Column::from_name(col_name)).eq(true.lit()), diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index cd0366cfac4..c941839c0ad 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -664,9 +664,10 @@ mod tests { use catalog::RegisterTableRequest; use catalog::memory::MemoryCatalogManager; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_time::Timezone; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, Schema}; - use session::context::QueryContext; + use session::context::{QueryContext, QueryContextBuilder}; use store_api::metric_engine_consts::{ DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, @@ -697,6 +698,36 @@ mod tests { crate::tests::new_query_engine_with_table(table) } + async fn create_timestamp_test_engine() -> QueryEngineRef { + let columns = vec![ + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + "st", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ]; + let schema = Arc::new(Schema::new(columns)); + let table_meta = TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![]) + .value_indices(vec![0, 1]) + .next_column_id(1024) + .build() + .unwrap(); + let table_info = TableInfoBuilder::new("timestamps", table_meta) + .build() + .unwrap(); + let table = EmptyTable::from_table_info(&table_info); + + crate::tests::new_query_engine_with_table(table) + } + fn create_promql_test_engine() -> QueryEngineRef { let catalog_manager = MemoryCatalogManager::with_default_setup(); let physical_table_name = "phy"; @@ -829,6 +860,43 @@ mod tests { engine.planner().plan(&stmt, query_ctx).await.unwrap() } + #[tokio::test] + async fn test_insert_timestamp_literals_use_query_timezone() { + let query_ctx = Arc::new( + QueryContextBuilder::default() + .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()) + .build(), + ); + let engine = create_timestamp_test_engine().await; + + for (sql, expected_timestamp) in [ + ( + "INSERT INTO timestamps (ts, st) \ + VALUES ('2026-08-02 12:00:00.001', now())", + 1_785_643_200_001_i64, + ), + ( + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-03 12:00:00.001', now()", + 1_785_729_600_001_i64, + ), + ] { + let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); + let plan = engine + .planner() + .plan(&stmt, query_ctx.clone()) + .await + .unwrap() + .display_indent() + .to_string(); + + assert!( + plan.contains(&format!("TimestampMillisecond({expected_timestamp}, None)")), + "{plan}" + ); + } + } + #[tokio::test] async fn test_extract_placeholder_cast_types_multiple() { let plan = parse_sql_to_plan( diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.result b/tests/cases/standalone/common/insert/insert_default_timezone.result index 9b12e4a15d0..da8b1261e7a 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.result +++ b/tests/cases/standalone/common/insert/insert_default_timezone.result @@ -52,6 +52,41 @@ SELECT * FROM test1 ORDER BY j; | 4 | 2025-02-01T00:01:01 | +---+---------------------+ +CREATE TABLE test3 (ts TIMESTAMP TIME INDEX, st TIMESTAMP, ts_ns TIMESTAMP(9)); + +Affected Rows: 0 + +INSERT INTO test3 (ts) VALUES ('2026-08-01 12:00:00.001'); + +Affected Rows: 1 + +INSERT INTO test3 (ts, st) VALUES ('2026-08-02 12:00:00.001', now()); + +Affected Rows: 1 + +INSERT INTO test3 (ts, st) SELECT '2026-08-03 12:00:00.001', now(); + +Affected Rows: 1 + +INSERT INTO test3 (ts, st, ts_ns) VALUES ( + '2026-08-04 12:00:00.001', + now(), + '2026-08-04 12:00:00.123456789' +); + +Affected Rows: 1 + +SELECT ts, ts_ns FROM test3 ORDER BY ts; + ++-------------------------+-------------------------------+ +| ts | ts_ns | ++-------------------------+-------------------------------+ +| 2026-08-01T04:00:00.001 | | +| 2026-08-02T04:00:00.001 | | +| 2026-08-03T04:00:00.001 | | +| 2026-08-04T04:00:00.001 | 2026-08-04T04:00:00.123456789 | ++-------------------------+-------------------------------+ + SET time_zone = 'UTC'; Affected Rows: 0 @@ -64,3 +99,7 @@ DROP TABLE test2; Affected Rows: 0 +DROP TABLE test3; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.sql b/tests/cases/standalone/common/insert/insert_default_timezone.sql index 2473cdcee49..fba8c134fa0 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.sql +++ b/tests/cases/standalone/common/insert/insert_default_timezone.sql @@ -16,8 +16,26 @@ SELECT * FROM test2 ORDER BY j; SELECT * FROM test1 ORDER BY j; +CREATE TABLE test3 (ts TIMESTAMP TIME INDEX, st TIMESTAMP, ts_ns TIMESTAMP(9)); + +INSERT INTO test3 (ts) VALUES ('2026-08-01 12:00:00.001'); + +INSERT INTO test3 (ts, st) VALUES ('2026-08-02 12:00:00.001', now()); + +INSERT INTO test3 (ts, st) SELECT '2026-08-03 12:00:00.001', now(); + +INSERT INTO test3 (ts, st, ts_ns) VALUES ( + '2026-08-04 12:00:00.001', + now(), + '2026-08-04 12:00:00.123456789' +); + +SELECT ts, ts_ns FROM test3 ORDER BY ts; + SET time_zone = 'UTC'; DROP TABLE test1; DROP TABLE test2; + +DROP TABLE test3; From 4dd92c774e3cceca0b3e505f23334f4d55c18444 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Fri, 14 Aug 2026 08:24:50 +0000 Subject: [PATCH 49/92] feat: add json_object function and use it in the entity-graph derivation (#8870) * feat: add json_object scalar function Builds a JSONB object from interleaved (key, value, ...) arguments, like MySQL's JSON_OBJECT. Values are written into the binary directly, so JSON-hostile characters (quotes, backslashes, control characters) need no text-level escaping. Keys must be non-NULL strings; values may be strings, numbers, booleans, or NULL (JSON null). Signed-off-by: Dennis Zhuang * fix: build entity-graph JSON objects with json_object The derivation assembled entity_id_attrs and descriptive by concatenating a JSON text and parsing it, escaping only backslash and double quote in runtime values. A label containing a control character (e.g. a newline) produced unparseable text and failed the whole semantic_entities scan instead of one attribute. json_object assembles the JSONB binary directly from the value columns, so no text escaping is involved; NULL-to-'' stays at the call site. Signed-off-by: Dennis Zhuang * chore: trim comments and fold duplicate test coverage Signed-off-by: Dennis Zhuang * fix: json_object() returns an empty object; narrow values to integers and floats MySQL's JSON_OBJECT allows an empty pair list, so the signature accepts zero arguments and the row count falls back to number_rows. Decimals stay rejected instead of casting to Float64: JSONB numbers (i64/u64/f64) cannot represent them exactly and a silent precision loss is worse than an explicit cast. Signed-off-by: Dennis Zhuang * chore: document key-to-string conversion and align test naming Keys follow MySQL JSON_OBJECT: any castable type is converted to string. Rustdoc and the cast-failure message now say so, with a numeric-key test. Test names take the module-conventional test_ prefix. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/common/function/src/scalars/json.rs | 2 + .../function/src/scalars/json/json_object.rs | 281 ++++++++++++++++++ src/operator/src/statement/semantic_graph.rs | 66 ++-- .../common/function/json/json.result | 70 +++++ .../standalone/common/function/json/json.sql | 25 ++ 5 files changed, 406 insertions(+), 38 deletions(-) create mode 100644 src/common/function/src/scalars/json/json_object.rs diff --git a/src/common/function/src/scalars/json.rs b/src/common/function/src/scalars/json.rs index 8a06c972c9d..d0e7f412ac2 100644 --- a/src/common/function/src/scalars/json.rs +++ b/src/common/function/src/scalars/json.rs @@ -15,6 +15,7 @@ pub mod json_get; mod json_get_rewriter; mod json_is; +mod json_object; mod json_object_keys; mod json_path_exists; mod json_path_match; @@ -54,6 +55,7 @@ impl JsonFunction { registry.register_scalar(JsonIsArray::default()); registry.register_scalar(JsonIsObject::default()); + registry.register_scalar(json_object::JsonObjectFunction::default()); registry.register_scalar(json_object_keys::JsonObjectKeysFunction::default()); registry.register_scalar(json_path_exists::JsonPathExistsFunction::default()); registry.register_scalar(json_path_match::JsonPathMatchFunction::default()); diff --git a/src/common/function/src/scalars/json/json_object.rs b/src/common/function/src/scalars/json/json_object.rs new file mode 100644 index 00000000000..a9602c3abc5 --- /dev/null +++ b/src/common/function/src/scalars/json/json_object.rs @@ -0,0 +1,281 @@ +// Copyright 2023 Greptime Team +// +// 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. + +use std::fmt::{self, Display}; +use std::sync::Arc; + +use datafusion_common::DataFusionError; +use datafusion_common::arrow::array::{Array, ArrayRef, AsArray, BinaryViewBuilder}; +use datafusion_common::arrow::compute; +use datafusion_common::arrow::datatypes::{DataType, Float64Type, Int64Type, UInt64Type}; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, TypeSignature, Volatility}; + +use crate::function::Function; + +const NAME: &str = "json_object"; + +/// Builds a `JSONB` object from interleaved `(key, value, key, value, ...)` +/// arguments, like MySQL's `JSON_OBJECT`; called with no arguments it returns +/// `{}`. Values are written into the binary directly, so they need no JSON +/// text escaping. Keys must be non-NULL and are converted to strings (so a +/// numeric key like `1` becomes `"1"`, as in MySQL); values may be strings, +/// integers, floats, booleans, or NULL (rendered as JSON null). Other types — +/// including decimals, which JSONB numbers cannot represent exactly — are +/// rejected; cast them explicitly. A duplicate key keeps the last value. +#[derive(Clone, Debug)] +pub(crate) struct JsonObjectFunction { + signature: Signature, +} + +impl Default for JsonObjectFunction { + fn default() -> Self { + Self { + signature: Signature::one_of( + vec![TypeSignature::Nullary, TypeSignature::VariadicAny], + Volatility::Immutable, + ), + } + } +} + +/// A value column normalized to the canonical arrow type its JSON rendering +/// reads from. +enum ValueColumn { + Null, + Bool(ArrayRef), + Int(ArrayRef), + UInt(ArrayRef), + Float(ArrayRef), + String(ArrayRef), +} + +impl ValueColumn { + fn try_new(array: &ArrayRef) -> datafusion_common::Result { + let normalized = match array.data_type() { + DataType::Null => return Ok(ValueColumn::Null), + DataType::Boolean => ValueColumn::Bool(array.clone()), + DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => { + ValueColumn::Int(compute::cast(array, &DataType::Int64)?) + } + DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => { + ValueColumn::UInt(compute::cast(array, &DataType::UInt64)?) + } + DataType::Float16 | DataType::Float32 | DataType::Float64 => { + ValueColumn::Float(compute::cast(array, &DataType::Float64)?) + } + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { + ValueColumn::String(compute::cast(array, &DataType::Utf8View)?) + } + other => { + return Err(DataFusionError::Execution(format!( + "{NAME} does not support values of type {other}; cast the value to a string" + ))); + } + }; + Ok(normalized) + } + + fn value(&self, row: usize) -> jsonb::Value<'_> { + let array = match self { + ValueColumn::Null => return jsonb::Value::Null, + ValueColumn::Bool(array) + | ValueColumn::Int(array) + | ValueColumn::UInt(array) + | ValueColumn::Float(array) + | ValueColumn::String(array) => array, + }; + if !array.is_valid(row) { + return jsonb::Value::Null; + } + match self { + ValueColumn::Null => unreachable!(), + ValueColumn::Bool(array) => array.as_boolean().value(row).into(), + ValueColumn::Int(array) => array.as_primitive::().value(row).into(), + ValueColumn::UInt(array) => array.as_primitive::().value(row).into(), + ValueColumn::Float(array) => array.as_primitive::().value(row).into(), + ValueColumn::String(array) => array.as_string_view().value(row).into(), + } + } +} + +impl Function for JsonObjectFunction { + fn name(&self) -> &str { + NAME + } + + fn return_type(&self, _: &[DataType]) -> datafusion_common::Result { + Ok(DataType::BinaryView) + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn invoke_with_args( + &self, + args: ScalarFunctionArgs, + ) -> datafusion_common::Result { + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + if arrays.len() % 2 != 0 { + return Err(DataFusionError::Execution(format!( + "{NAME} expects (key, value) argument pairs, got {} arguments", + arrays.len() + ))); + } + let pairs = arrays + .chunks(2) + .map(|pair| { + let keys = compute::cast(&pair[0], &DataType::Utf8View).map_err(|_| { + DataFusionError::Execution(format!( + "{NAME} cannot convert keys of type {} to string", + pair[0].data_type() + )) + })?; + if keys.null_count() > 0 { + return Err(DataFusionError::Execution(format!( + "{NAME} does not allow NULL keys" + ))); + } + Ok((keys, ValueColumn::try_new(&pair[1])?)) + }) + .collect::>>()?; + + let rows = arrays.first().map_or(args.number_rows, |a| a.len()); + let mut builder = BinaryViewBuilder::with_capacity(rows); + let mut buf = Vec::new(); + for row in 0..rows { + let mut object = jsonb::Object::new(); + for (keys, values) in &pairs { + object.insert( + keys.as_string_view().value(row).to_string(), + values.value(row), + ); + } + buf.clear(); + jsonb::Value::Object(object).write_to_vec(&mut buf); + builder.append_value(&buf); + } + Ok(ColumnarValue::Array(Arc::new(builder.finish()))) + } +} + +impl Display for JsonObjectFunction { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "JSON_OBJECT") + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_schema::Field; + use datafusion_common::arrow::array::{ + Int64Array, NullArray, StringArray, TimestampMillisecondArray, + }; + + use super::*; + + fn invoke(args: Vec, rows: usize) -> datafusion_common::Result> { + let function = JsonObjectFunction::default(); + let result = function.invoke_with_args(ScalarFunctionArgs { + args, + arg_fields: vec![], + number_rows: rows, + return_field: Arc::new(Field::new("x", DataType::BinaryView, true)), + config_options: Arc::new(Default::default()), + })?; + let array = result.to_array(rows)?; + let array = array.as_binary_view(); + Ok((0..array.len()) + .map(|i| jsonb::from_slice(array.value(i)).unwrap().to_string()) + .collect()) + } + + fn key(name: &str) -> ColumnarValue { + ColumnarValue::Scalar(datafusion_common::ScalarValue::Utf8(Some(name.to_string()))) + } + + #[test] + fn test_builds_objects_from_mixed_types_without_escaping() { + let texts = invoke( + vec![ + key("host"), + ColumnarValue::Array(Arc::new(StringArray::from(vec![ + Some("we\"ird\\\nhost"), + None, + ]))), + key("pid"), + ColumnarValue::Array(Arc::new(Int64Array::from(vec![42, 7]))), + key("up"), + ColumnarValue::Scalar(datafusion_common::ScalarValue::Boolean(Some(true))), + key("v"), + ColumnarValue::Array(Arc::new(NullArray::new(2))), + ], + 2, + ) + .unwrap(); + assert_eq!( + texts, + vec![ + r#"{"host":"we\"ird\\\nhost","pid":42,"up":true,"v":null}"#, + r#"{"host":null,"pid":7,"up":true,"v":null}"#, + ] + ); + } + + #[test] + fn test_empty_call_builds_one_empty_object_per_row() { + assert_eq!(invoke(vec![], 3).unwrap(), vec!["{}", "{}", "{}"]); + } + + #[test] + fn test_numeric_key_converts_to_string() { + let texts = invoke( + vec![ + ColumnarValue::Scalar(datafusion_common::ScalarValue::Int64(Some(7))), + key("v"), + ], + 1, + ) + .unwrap(); + assert_eq!(texts, vec![r#"{"7":"v"}"#]); + } + + #[test] + fn test_rejects_odd_arguments_null_keys_and_unsupported_values() { + let err = invoke(vec![key("a")], 1).unwrap_err(); + assert!(err.to_string().contains("(key, value) argument pairs")); + + let err = invoke( + vec![ + ColumnarValue::Array(Arc::new(StringArray::from(vec![None::<&str>]))), + key("v"), + ], + 1, + ) + .unwrap_err(); + assert!(err.to_string().contains("NULL keys")); + + let err = invoke( + vec![ + key("ts"), + ColumnarValue::Array(Arc::new(TimestampMillisecondArray::from(vec![1_000]))), + ], + 1, + ) + .unwrap_err(); + assert!(err.to_string().contains("does not support values")); + } +} diff --git a/src/operator/src/statement/semantic_graph.rs b/src/operator/src/statement/semantic_graph.rs index 777c6f677ea..7014a94f8ac 100644 --- a/src/operator/src/statement/semantic_graph.rs +++ b/src/operator/src/statement/semantic_graph.rs @@ -465,43 +465,34 @@ fn null_json() -> Expr { lit(ScalarValue::Binary(None)) } -/// Renders a compile-time-known string as JSON text (quoted, fully escaped — -/// including control characters, unlike the runtime value escaping). +/// Renders a compile-time-known string as JSON text (quoted, fully escaped). fn json_quote(value: &str) -> String { serde_json::Value::from(value).to_string() } -/// Wraps a column reference in `replace` calls so its runtime value is -/// JSON-escaped (`\` then `"`); NULL becomes `''` so one NULL column does not -/// invalidate the whole JSON text (descriptive columns are nullable). -fn json_escaped_value_expr(column: &str) -> Expr { - let escaped_backslash = - string_fns::replace().call(vec![cast_string_or_empty(column), lit("\\"), lit("\\\\")]); - string_fns::replace().call(vec![escaped_backslash, lit("\""), lit("\\\"")]) -} +/// The `json_object` UDF, resolved like [`PARSE_JSON_UDF`]. It assembles the +/// JSONB binary directly from the value columns, so runtime values need no +/// JSON text escaping. +static JSON_OBJECT_UDF: LazyLock> = LazyLock::new(|| { + Arc::new( + FUNCTION_REGISTRY + .get_function("json_object") + .expect("json_object must be registered") + .provide(FunctionContext::default()), + ) +}); -/// Builds a JSONB object from `columns` by concatenating a JSON text and parsing -/// it — GreptimeDB has no struct→json function. Keys are JSON-escaped in Rust; -/// values are JSON-escaped at runtime via [`json_escaped_value_expr`]. -/// -/// TODO(entity-graph): replace the text round-trip with a UDF that assembles -/// JSONB directly from the value columns (`jsonb::ObjectBuilder`, keys baked -/// in), dropping the escaping helpers and the per-row parse cost. +/// Builds a JSONB object with one entry per column: key = the column name, +/// value = the column rendered as a string, NULL coalesced to `""` so one NULL +/// column does not null the entry (descriptive columns are nullable). Keys come +/// out sorted — JSONB objects are key-ordered regardless of input order. fn json_object_expr(columns: &[String]) -> Expr { - if columns.is_empty() { - return parse_json_expr(lit("{}")); + let mut args = Vec::with_capacity(columns.len() * 2); + for column in columns { + args.push(lit(column.as_str())); + args.push(cast_string_or_empty(column)); } - let mut parts = vec![lit("{")]; - for (i, column) in columns.iter().enumerate() { - if i > 0 { - parts.push(lit(",")); - } - parts.push(lit(format!("{}:\"", json_quote(column)))); - parts.push(json_escaped_value_expr(column)); - parts.push(lit("\"")); - } - parts.push(lit("}")); - parse_json_expr(concat_expr(parts)) + cast(JSON_OBJECT_UDF.call(args), DataType::Binary) } /// Renders pre-sorted columns as a `k=v,k=v` concatenation. `nullable` @@ -599,9 +590,7 @@ fn registry_source( let entity_id_attrs = if decl.id_columns.len() == 1 { null_json() } else { - let mut cols = decl.id_columns.clone(); - cols.sort(); - json_object_expr(&cols) + json_object_expr(&decl.id_columns) }; let scope = match decl.scope_columns.as_slice() { @@ -799,7 +788,7 @@ mod tests { Arc::new(StringArray::from(vec!["cart", "cart", "cart"])), Arc::new(Int64Array::from(vec![42, 42, 42])), Arc::new(StringArray::from(vec![ - Some(r#"we"ird\host"#), + Some("we\"ird\\\nhost"), None, Some("h2"), ])), @@ -912,7 +901,8 @@ mod tests { rows.sort(); // Composite id -> sorted `k=v,k=v` plus a JSON object of the id columns; - // descriptive JSON escapes `\` and `"` in runtime values, NULL -> "". + // descriptive JSON keeps `\`, `"` and control characters intact in + // runtime values, NULL -> "". assert_eq!( rows, vec![ @@ -929,7 +919,7 @@ mod tests { ( "pid=42,service_name=cart".to_string(), Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()), - Some(r#"{"host":"we\"ird\\host"}"#.to_string()), + Some(r#"{"host":"we\"ird\\\nhost"}"#.to_string()), ), ] ); @@ -955,7 +945,7 @@ mod tests { let batches = collect(&ctx, plan).await; let mut scopes: Vec = batches.iter().flat_map(|b| strings(b, 7)).collect(); scopes.sort(); - assert_eq!(scopes, vec!["", "h2", r#"we"ird\host"#]); + assert_eq!(scopes, vec!["", "h2", "we\"ird\\\nhost"]); // Multiple scope columns: sorted `k=v,k=v`. let mut multi = decl("service", &["service_name"]); @@ -978,7 +968,7 @@ mod tests { vec![ "host=,pid=42", "host=h2,pid=42", - r#"host=we"ird\host,pid=42"# + "host=we\"ird\\\nhost,pid=42" ] ); } diff --git a/tests/cases/standalone/common/function/json/json.result b/tests/cases/standalone/common/function/json/json.result index ab32c9c0132..c035db72e4f 100644 --- a/tests/cases/standalone/common/function/json/json.result +++ b/tests/cases/standalone/common/function/json/json.result @@ -145,3 +145,73 @@ SELECT json_path_match(parse_json('null'), '$.a == 1'); | | +------------------------------------------------------------+ +--- json_object --- +SELECT json_to_string(json_object('a', 1, 'b', 'text', 'c', true, 'd', 1.5)); + ++-----------------------------------------------------------------------------------------------------------------------+ +| json_to_string(json_object(Utf8("a"),Int64(1),Utf8("b"),Utf8("text"),Utf8("c"),Boolean(true),Utf8("d"),Float64(1.5))) | ++-----------------------------------------------------------------------------------------------------------------------+ +| {"a":1,"b":"text","c":true,"d":1.5} | ++-----------------------------------------------------------------------------------------------------------------------+ + +SELECT json_to_string(json_object()); + ++-------------------------------+ +| json_to_string(json_object()) | ++-------------------------------+ +| {} | ++-------------------------------+ + +SELECT json_to_string(json_object('nul', NULL)); + ++-----------------------------------------------+ +| json_to_string(json_object(Utf8("nul"),NULL)) | ++-----------------------------------------------+ +| {"nul":null} | ++-----------------------------------------------+ + +SELECT json_to_string(json_object('nl', concat('line1', chr(10), 'line2'), 'q', 'quote"back\slash')); + ++-------------------------------------------------------------------------------------------------------------------------------+ +| json_to_string(json_object(Utf8("nl"),concat(Utf8("line1"),chr(Int64(10)),Utf8("line2")),Utf8("q"),Utf8("quote"back\slash"))) | ++-------------------------------------------------------------------------------------------------------------------------------+ +| {"nl":"line1\nline2","q":"quote\"back\\slash"} | ++-------------------------------------------------------------------------------------------------------------------------------+ + +SELECT json_object('a'); + +Error: 3001(EngineExecuteQuery), Execution error: json_object expects (key, value) argument pairs, got 1 arguments + +SELECT json_object(NULL, 1); + +Error: 3001(EngineExecuteQuery), Execution error: json_object does not allow NULL keys + +SELECT json_object('ts', to_timestamp(0)); + +Error: 3001(EngineExecuteQuery), Execution error: json_object does not support values of type Timestamp(ns); cast the value to a string + +SELECT json_object('price', CAST('12.34' AS DECIMAL(10, 2))); + +Error: 3001(EngineExecuteQuery), Execution error: json_object does not support values of type Decimal128(10, 2); cast the value to a string + +CREATE TABLE json_object_src (ts TIMESTAMP TIME INDEX, host STRING, pid BIGINT); + +Affected Rows: 0 + +INSERT INTO json_object_src VALUES (0, 'h1', 42), (1, NULL, 7); + +Affected Rows: 2 + +SELECT json_to_string(json_object('host', host, 'pid', pid)) FROM json_object_src ORDER BY ts; + ++------------------------------------------------------------------------------------------------+ +| json_to_string(json_object(Utf8("host"),json_object_src.host,Utf8("pid"),json_object_src.pid)) | ++------------------------------------------------------------------------------------------------+ +| {"host":"h1","pid":42} | +| {"host":null,"pid":7} | ++------------------------------------------------------------------------------------------------+ + +DROP TABLE json_object_src; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/function/json/json.sql b/tests/cases/standalone/common/function/json/json.sql index e09f206df9f..16517b54f69 100644 --- a/tests/cases/standalone/common/function/json/json.sql +++ b/tests/cases/standalone/common/function/json/json.sql @@ -37,3 +37,28 @@ SELECT json_path_match(parse_json('{"a":1,"b":[1,2,3]}'), '$.b[1 to last] >= 2') SELECT json_path_match(parse_json('{"a":1,"b":[1,2,3]}'), 'null'); SELECT json_path_match(parse_json('null'), '$.a == 1'); + +--- json_object --- +SELECT json_to_string(json_object('a', 1, 'b', 'text', 'c', true, 'd', 1.5)); + +SELECT json_to_string(json_object()); + +SELECT json_to_string(json_object('nul', NULL)); + +SELECT json_to_string(json_object('nl', concat('line1', chr(10), 'line2'), 'q', 'quote"back\slash')); + +SELECT json_object('a'); + +SELECT json_object(NULL, 1); + +SELECT json_object('ts', to_timestamp(0)); + +SELECT json_object('price', CAST('12.34' AS DECIMAL(10, 2))); + +CREATE TABLE json_object_src (ts TIMESTAMP TIME INDEX, host STRING, pid BIGINT); + +INSERT INTO json_object_src VALUES (0, 'h1', 42), (1, NULL, 7); + +SELECT json_to_string(json_object('host', host, 'pid', pid)) FROM json_object_src ORDER BY ts; + +DROP TABLE json_object_src; From 2eb5f593f8dde448a81c8a3bfaac9867bc69b8aa Mon Sep 17 00:00:00 2001 From: sun Date: Fri, 14 Aug 2026 08:45:02 +0000 Subject: [PATCH 50/92] feat: update dashboard to v0.13.12 (#8882) --- src/servers/dashboard/VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/servers/dashboard/VERSION b/src/servers/dashboard/VERSION index dbbe182110e..3d431cfd5f0 100644 --- a/src/servers/dashboard/VERSION +++ b/src/servers/dashboard/VERSION @@ -1 +1 @@ -v0.13.11 +v0.13.12 From ed4271af403e6fa8898287ddd735c1c984644af7 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Fri, 14 Aug 2026 09:36:01 +0000 Subject: [PATCH 51/92] feat(procedure): record event actor (#8849) * feat(event): record procedure actor Signed-off-by: WenyXu * test: handle streamed region migration output Signed-off-by: WenyXu * test: cover procedure actors across SQL protocols Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/common/event-recorder/src/event_table.rs | 283 ------------------ .../meta/src/ddl/tests/event/database.rs | 4 +- src/common/meta/src/ddl/tests/event/flow.rs | 6 +- src/common/meta/src/ddl/tests/event/view.rs | 4 +- src/common/procedure/src/event.rs | 65 +++- src/meta-client/src/client/procedure.rs | 54 +--- src/meta-srv/src/event/gc.rs | 6 +- src/meta-srv/src/event/repartition.rs | 6 +- src/meta-srv/src/service/procedure.rs | 76 +---- tests-integration/src/test_util.rs | 55 +++- tests-integration/tests/database_ddl_event.rs | 60 +++- .../tests/event_recorder_test_util.rs | 57 ++++ tests-integration/tests/flow_ddl_event.rs | 48 +-- tests-integration/tests/gc_event.rs | 53 ++-- tests-integration/tests/http.rs | 25 ++ tests-integration/tests/region_migration.rs | 80 ++++- tests-integration/tests/repartition_event.rs | 43 +-- tests-integration/tests/sql.rs | 48 ++- tests-integration/tests/table_ddl_event.rs | 30 +- tests-integration/tests/view_ddl_event.rs | 51 ++-- tests-integration/tests/wal_prune_event.rs | 1 + 21 files changed, 504 insertions(+), 551 deletions(-) diff --git a/src/common/event-recorder/src/event_table.rs b/src/common/event-recorder/src/event_table.rs index 1ea4369a00a..25c16f433d1 100644 --- a/src/common/event-recorder/src/event_table.rs +++ b/src/common/event-recorder/src/event_table.rs @@ -310,286 +310,3 @@ pub fn nullable_json(value: Option<&serde_json::Value>) -> Value { pub fn jsonb_value(value: &serde_json::Value) -> Value { ValueData::BinaryValue(jsonb::Value::from(value).to_vec()).into() } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn base_schema_preserves_names_types_semantics_extensions_and_order() { - assert_eq!( - base_column_schemas(), - vec![ - ColumnSchema { - column_name: "type".to_string(), - datatype: ColumnDataType::String.into(), - semantic_type: SemanticType::Tag.into(), - ..Default::default() - }, - ColumnSchema { - column_name: "payload".to_string(), - datatype: ColumnDataType::Binary.into(), - semantic_type: SemanticType::Field.into(), - datatype_extension: Some(ColumnDataTypeExtension { - type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())), - }), - ..Default::default() - }, - ColumnSchema { - column_name: "timestamp".to_string(), - datatype: ColumnDataType::TimestampNanosecond.into(), - semantic_type: SemanticType::Timestamp.into(), - ..Default::default() - }, - ] - ); - } - - #[test] - fn procedure_envelope_schema_preserves_names_types_semantics_and_order() { - assert_eq!( - procedure_event_column_schemas(), - vec![ - ColumnSchema { - column_name: "procedure_id".to_string(), - datatype: ColumnDataType::String.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }, - ColumnSchema { - column_name: "procedure_state".to_string(), - datatype: ColumnDataType::String.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }, - ColumnSchema { - column_name: "procedure_error".to_string(), - datatype: ColumnDataType::String.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }, - ColumnSchema { - column_name: "procedure_trigger".to_string(), - datatype: ColumnDataType::Binary.into(), - semantic_type: SemanticType::Field.into(), - datatype_extension: Some(ColumnDataTypeExtension { - type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())), - }), - ..Default::default() - }, - ] - ); - } - - #[test] - fn shared_dimension_schema_preserves_names_types_semantics_and_order() { - assert_eq!( - column_schemas([&CATALOG_NAME_COLUMN, &SCHEMA_NAME_COLUMN]), - ["catalog_name", "schema_name"].map(|column_name| ColumnSchema { - column_name: column_name.to_string(), - datatype: ColumnDataType::String.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }) - ); - } - - #[test] - fn flow_dimension_schema_preserves_names_types_semantics_and_order() { - assert_eq!( - column_schemas([&FLOW_NAME_COLUMN, &FLOW_ID_COLUMN]), - [ - ("flow_name", ColumnDataType::String), - ("flow_id", ColumnDataType::Uint32), - ] - .map(|(column_name, datatype)| ColumnSchema { - column_name: column_name.to_string(), - datatype: datatype.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }) - ); - } - - #[test] - fn admin_function_schema_preserves_names_types_semantics_extensions_and_order() { - assert_eq!( - column_schemas([ - &ACTOR_COLUMN, - &ADMIN_FUNCTION_NAME_COLUMN, - &ADMIN_FUNCTION_STATUS_COLUMN, - &ADMIN_FUNCTION_OUTPUT_COLUMN, - ]), - vec![ - ColumnSchema { - column_name: "actor".to_string(), - datatype: ColumnDataType::String.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }, - ColumnSchema { - column_name: "admin_function_name".to_string(), - datatype: ColumnDataType::String.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }, - ColumnSchema { - column_name: "admin_function_status".to_string(), - datatype: ColumnDataType::String.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }, - ColumnSchema { - column_name: "admin_function_output".to_string(), - datatype: ColumnDataType::Binary.into(), - semantic_type: SemanticType::Field.into(), - datatype_extension: Some(ColumnDataTypeExtension { - type_ext: Some(TypeExt::JsonType(JsonTypeExtension::JsonBinary.into())), - }), - ..Default::default() - }, - ] - ); - } - - #[test] - fn table_dimension_schema_preserves_names_types_semantics_and_order() { - assert_eq!( - column_schemas([ - &TABLE_NAME_COLUMN, - &TABLE_ID_COLUMN, - &PHYSICAL_TABLE_ID_COLUMN, - ]), - [ - ("table_name", ColumnDataType::String), - ("table_id", ColumnDataType::Uint32), - ("physical_table_id", ColumnDataType::Uint32), - ] - .map(|(column_name, datatype)| ColumnSchema { - column_name: column_name.to_string(), - datatype: datatype.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }) - ); - } - - #[test] - fn view_dimension_schema_preserves_names_types_semantics_and_order() { - assert_eq!( - column_schemas([&VIEW_NAME_COLUMN, &VIEW_ID_COLUMN]), - [ - ("view_name", ColumnDataType::String), - ("view_id", ColumnDataType::Uint32), - ] - .map(|(column_name, datatype)| ColumnSchema { - column_name: column_name.to_string(), - datatype: datatype.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }) - ); - } - - #[test] - fn region_dimension_schema_preserves_names_types_semantics_and_order() { - assert_field_columns( - [ - ®ION_ID_COLUMN, - ®ION_NUMBER_COLUMN, - ®ION_MIGRATION_TRIGGER_REASON_COLUMN, - ®ION_MIGRATION_SRC_NODE_ID_COLUMN, - ®ION_MIGRATION_SRC_PEER_ADDR_COLUMN, - ®ION_MIGRATION_DST_NODE_ID_COLUMN, - ®ION_MIGRATION_DST_PEER_ADDR_COLUMN, - ], - [ - ("region_id", ColumnDataType::Uint64), - ("region_number", ColumnDataType::Uint32), - ("region_migration_trigger_reason", ColumnDataType::String), - ("region_migration_src_node_id", ColumnDataType::Uint64), - ("region_migration_src_peer_addr", ColumnDataType::String), - ("region_migration_dst_node_id", ColumnDataType::Uint64), - ("region_migration_dst_peer_addr", ColumnDataType::String), - ], - ); - } - - #[test] - fn repartition_dimension_schema_preserves_names_types_semantics_and_order() { - assert_field_columns( - [ - &TABLE_NAME_COLUMN, - &TABLE_ID_COLUMN, - &PARENT_PROCEDURE_ID_COLUMN, - &REPARTITION_GROUP_ID_COLUMN, - &SOURCE_REGION_ID_COLUMN, - &SOURCE_REGION_NUMBER_COLUMN, - &SOURCE_PARTITION_EXPR_COLUMN, - &TARGET_REGION_ID_COLUMN, - &TARGET_REGION_NUMBER_COLUMN, - &TARGET_PARTITION_EXPR_COLUMN, - ], - [ - ("table_name", ColumnDataType::String), - ("table_id", ColumnDataType::Uint32), - ("parent_procedure_id", ColumnDataType::String), - ("repartition_group_id", ColumnDataType::String), - ("source_region_id", ColumnDataType::Uint64), - ("source_region_number", ColumnDataType::Uint32), - ("source_partition_expr", ColumnDataType::String), - ("target_region_id", ColumnDataType::Uint64), - ("target_region_number", ColumnDataType::Uint32), - ("target_partition_expr", ColumnDataType::String), - ], - ); - } - - #[test] - fn wal_prune_schema_preserves_names_types_semantics_and_order() { - assert_field_columns( - [ - &TOPIC_NAME_COLUMN, - &PRUNABLE_ENTRY_ID_COLUMN, - &LATEST_OFFSET_COLUMN, - ], - [ - ("topic_name", ColumnDataType::String), - ("prunable_entry_id", ColumnDataType::Uint64), - ("latest_offset", ColumnDataType::Uint64), - ], - ); - } - - #[test] - fn nullable_values_preserve_types_and_nulls() { - assert_eq!( - nullable_string(Some("catalog")), - Value { - value_data: Some(ValueData::StringValue("catalog".to_string())) - } - ); - assert_eq!(nullable_value(None), Value { value_data: None }); - assert_eq!( - nullable_value(Some(ValueData::BoolValue(true))), - Value { - value_data: Some(ValueData::BoolValue(true)) - } - ); - } - - fn assert_field_columns( - columns: [&EventTableColumn; N], - expected: [(&str, ColumnDataType); N], - ) { - assert_eq!( - column_schemas(columns), - expected.map(|(column_name, datatype)| ColumnSchema { - column_name: column_name.to_string(), - datatype: datatype.into(), - semantic_type: SemanticType::Field.into(), - ..Default::default() - }) - ); - } -} diff --git a/src/common/meta/src/ddl/tests/event/database.rs b/src/common/meta/src/ddl/tests/event/database.rs index 8649982ed15..a617a1f5f6a 100644 --- a/src/common/meta/src/ddl/tests/event/database.rs +++ b/src/common/meta/src/ddl/tests/event/database.rs @@ -20,7 +20,7 @@ use api::v1::value::ValueData; use api::v1::{AlterDatabaseExpr, Row, SetDatabaseOptions as PbSetDatabaseOptions, Value}; use common_event_recorder::Event; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN as EVENT_TABLE_CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, + ACTOR_COLUMN, CATALOG_NAME_COLUMN as EVENT_TABLE_CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN as EVENT_TABLE_PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN as EVENT_TABLE_PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN as EVENT_TABLE_PROCEDURE_STATE_COLUMN, @@ -303,6 +303,7 @@ fn assert_procedure_event_contract( EVENT_TABLE_PROCEDURE_TRIGGER_COLUMN.column_schema(), EVENT_TABLE_CATALOG_NAME_COLUMN.column_schema(), EVENT_TABLE_SCHEMA_NAME_COLUMN.column_schema(), + ACTOR_COLUMN.column_schema(), EVENT_CONTEXT_COLUMN.column_schema(), ], &[Row { @@ -326,6 +327,7 @@ fn assert_procedure_event_contract( value_data: schema_name.map(|value| ValueData::StringValue(value.to_string())), }, Value { value_data: None }, + Value { value_data: None }, ], }], ); diff --git a/src/common/meta/src/ddl/tests/event/flow.rs b/src/common/meta/src/ddl/tests/event/flow.rs index cec0162adf8..297d6a6d9e2 100644 --- a/src/common/meta/src/ddl/tests/event/flow.rs +++ b/src/common/meta/src/ddl/tests/event/flow.rs @@ -19,8 +19,8 @@ use api::v1::{ColumnSchema, Row, Value}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_event_recorder::Event; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, jsonb_value, - procedure_event_column_schemas, + ACTOR_COLUMN, CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, FLOW_ID_COLUMN, FLOW_NAME_COLUMN, + jsonb_value, procedure_event_column_schemas, }; use common_event_recorder::testing::assert_event_contract; use common_procedure::{EventTrigger, ProcedureEvent, ProcedureId, ProcedureState}; @@ -241,6 +241,7 @@ fn assert_procedure_event_contract( ) { let mut schema = procedure_event_column_schemas(); schema.extend(flow_schema()); + schema.push(ACTOR_COLUMN.column_schema()); schema.push(EVENT_CONTEXT_COLUMN.column_schema()); assert_event_contract( event, @@ -260,6 +261,7 @@ fn assert_procedure_event_contract( .map(Into::into) .unwrap_or(Value { value_data: None }), Value { value_data: None }, + Value { value_data: None }, ], }], ); diff --git a/src/common/meta/src/ddl/tests/event/view.rs b/src/common/meta/src/ddl/tests/event/view.rs index 6ebb62cda82..85ecbae4ac8 100644 --- a/src/common/meta/src/ddl/tests/event/view.rs +++ b/src/common/meta/src/ddl/tests/event/view.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use api::v1::value::ValueData; use api::v1::{ColumnSchema, Row, Value}; use common_event_recorder::event_table::{ - CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, + ACTOR_COLUMN, CATALOG_NAME_COLUMN, EVENT_CONTEXT_COLUMN, SCHEMA_NAME_COLUMN, VIEW_ID_COLUMN, VIEW_NAME_COLUMN, jsonb_value, procedure_event_column_schemas, }; use common_event_recorder::testing::assert_event_contract; @@ -335,6 +335,7 @@ fn assert_procedure_event_contract( ) { let mut schema = procedure_event_column_schemas(); schema.extend(view_schema()); + schema.push(ACTOR_COLUMN.column_schema()); schema.push(EVENT_CONTEXT_COLUMN.column_schema()); let mut values = vec![ @@ -345,6 +346,7 @@ fn assert_procedure_event_contract( ]; values.extend(locator.values()); values.push(Value { value_data: None }); + values.push(Value { value_data: None }); assert_event_contract(event, event_type, &schema, &[Row { values }]); } diff --git a/src/common/procedure/src/event.rs b/src/common/procedure/src/event.rs index a4a6c290bf3..8ca3a6c67b3 100644 --- a/src/common/procedure/src/event.rs +++ b/src/common/procedure/src/event.rs @@ -19,8 +19,9 @@ use api::v1::{ColumnSchema, Row}; use common_event_recorder::Event; use common_event_recorder::error::{Result, SerializeEventSnafu}; use common_event_recorder::event_table::{ - EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, - PROCEDURE_TRIGGER_COLUMN, jsonb_value, nullable_json, procedure_event_column_schemas, + ACTOR_COLUMN, EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, + PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, jsonb_value, nullable_json, nullable_string, + procedure_event_column_schemas, }; use common_time::timestamp::{TimeUnit, Timestamp}; use snafu::ResultExt; @@ -100,6 +101,7 @@ impl Event for ProcedureEvent { let mut schema = procedure_event_column_schemas(); let mut internal_schema = self.internal_event.extra_schema(); schema.append(&mut internal_schema); + schema.push(ACTOR_COLUMN.column_schema()); schema.push(EVENT_CONTEXT_COLUMN.column_schema()); schema } @@ -124,9 +126,10 @@ impl Event for ProcedureEvent { .map(serde_json::to_value) .transpose() .context(SerializeEventSnafu)?; + let actor = self.context.actor.as_deref(); for internal_event_extra_row in internal_event_extra_rows.iter_mut() { - let mut values = Vec::with_capacity(5 + internal_event_extra_row.values.len()); + let mut values = Vec::with_capacity(6 + internal_event_extra_row.values.len()); values.extend([ ValueData::StringValue(procedure_id.clone()).into(), ValueData::StringValue(state.clone()).into(), @@ -134,6 +137,7 @@ impl Event for ProcedureEvent { jsonb_value(&trigger), ]); values.append(&mut internal_event_extra_row.values); + values.push(nullable_string(actor)); values.push(nullable_json(event_context.as_ref())); rows.push(Row { values }); } @@ -155,7 +159,7 @@ mod tests { use common_error::mock::MockError; use common_error::status_code::StatusCode; use common_event_recorder::event_table::{ - EVENT_CONTEXT_COLUMN, PROCEDURE_TRIGGER_COLUMN, jsonb_value, + ACTOR_COLUMN, EVENT_CONTEXT_COLUMN, PROCEDURE_TRIGGER_COLUMN, jsonb_value, }; use common_event_recorder::{Event, PersistentEventContext, TriggerReason}; use serde_json::json; @@ -223,6 +227,7 @@ mod tests { jsonb_value(&json!({"type": "Submitted"})), ValueData::StringValue("test_event1".to_string()).into(), Value { value_data: None }, + Value { value_data: None }, ], }, Row { @@ -233,6 +238,7 @@ mod tests { jsonb_value(&json!({"type": "Submitted"})), ValueData::StringValue("test_event2".to_string()).into(), Value { value_data: None }, + Value { value_data: None }, ], }, Row { @@ -243,6 +249,7 @@ mod tests { jsonb_value(&json!({"type": "Submitted"})), Value { value_data: None }, Value { value_data: None }, + Value { value_data: None }, ], }, ] @@ -313,16 +320,18 @@ mod tests { semantic_type: SemanticType::Field.into(), ..Default::default() }, + ACTOR_COLUMN.column_schema(), EVENT_CONTEXT_COLUMN.column_schema(), ] ); } #[test] - fn procedure_event_only_records_context_on_submission() { - let context = ProcedureContext::from_event_context(PersistentEventContext::new( + fn procedure_event_records_actor_for_all_lifecycle_triggers() { + let mut context = ProcedureContext::from_event_context(PersistentEventContext::new( TriggerReason::AutoRepartition, )); + context.actor = Some("alice".to_string()); let submitted = ProcedureEvent::new_with_context( ProcedureId::random(), Box::new(TestEvent {}), @@ -330,22 +339,46 @@ mod tests { EventTrigger::Submitted, context.clone(), ); - let succeeded = ProcedureEvent::new_with_context( - ProcedureId::random(), - Box::new(TestEvent {}), - ProcedureState::Done { output: None }, - EventTrigger::Succeeded, - context, - ); assert_eq!( submitted.extra_rows().unwrap()[0].values[5], - jsonb_value(&json!({"reason": "auto_repartition"})) + ValueData::StringValue("alice".to_string()).into() ); assert_eq!( - succeeded.extra_rows().unwrap()[0].values[5], - Value { value_data: None } + submitted.extra_rows().unwrap()[0].values[6], + jsonb_value(&json!({"reason": "auto_repartition"})) ); + + let child_id = ProcedureId::random(); + for trigger in [ + EventTrigger::Recovered, + EventTrigger::ChildSubmitted { + procedure_id: child_id, + outcome: ChildSubmissionOutcome::Accepted, + }, + EventTrigger::Retrying { + phase: RetryPhase::Execute, + attempt: 1, + }, + EventTrigger::RollingBack, + EventTrigger::Succeeded, + EventTrigger::Failed, + EventTrigger::Poisoned, + ] { + let event = ProcedureEvent::new_with_context( + ProcedureId::random(), + Box::new(TestEvent {}), + ProcedureState::Running, + trigger, + context.clone(), + ); + let values = &event.extra_rows().unwrap()[0].values; + assert_eq!( + values[5], + ValueData::StringValue("alice".to_string()).into() + ); + assert_eq!(values[6], Value { value_data: None }); + } } #[test] diff --git a/src/meta-client/src/client/procedure.rs b/src/meta-client/src/client/procedure.rs index ccf2c28188f..9c32b4fc880 100644 --- a/src/meta-client/src/client/procedure.rs +++ b/src/meta-client/src/client/procedure.rs @@ -63,13 +63,9 @@ pub(crate) fn procedure_event_context(context: &ExecutorContext) -> Option Option { - context - .actor - .as_deref() - .filter(|username| !username.is_empty()) - .map(|username| ProcedureActor { - username: username.to_string(), - }) + context.actor.as_ref().map(|username| ProcedureActor { + username: username.clone(), + }) } #[derive(Clone, Debug)] @@ -531,7 +527,7 @@ mod tests { use tonic::{Request, Response, Status}; use crate::client::MetaClientBuilder; - use crate::client::procedure::{gc_timeout_secs, procedure_actor, procedure_event_context}; + use crate::client::procedure::{gc_timeout_secs, procedure_event_context}; #[test] fn test_gc_timeout_secs() { @@ -578,21 +574,6 @@ mod tests { ); } - #[test] - fn test_procedure_actor() { - let context = ExecutorContext { - actor: Some(String::new()), - ..Default::default() - }; - assert!(procedure_actor(&context).is_none()); - - let context = ExecutorContext { - actor: Some("alice".to_string()), - ..Default::default() - }; - assert_eq!(procedure_actor(&context).unwrap().username, "alice"); - } - #[derive(Clone)] struct MockHeartbeat { leader_addr: String, @@ -728,18 +709,18 @@ mod tests { username: "alice".to_string(), created_at_ns: 42, }; - let executor_context = |actor: String| ExecutorContext { + let executor_context = ExecutorContext { query_context: Some(QueryContext { channel: Channel::Postgres as u8, ..Default::default() }), - actor: Some(actor), + actor: Some("effective-user".to_string()), event_input: Some(ProcedureEventInput::new(TriggerReason::Manual)), ..Default::default() }; ProcedureExecutor::submit_ddl_task( &client, - executor_context("effective-user".to_string()), + executor_context, SubmitDdlTaskRequest::new(DdlTask::new_create_database( "greptime".to_string(), "metrics".to_string(), @@ -772,27 +753,6 @@ mod tests { let extensions = &request.query_context.unwrap().extensions; assert_eq!(extensions[CREATE_DATABASE_CREATOR_EXTENSION_KEY], encoded); - ProcedureExecutor::submit_ddl_task( - &client, - executor_context(String::new()), - SubmitDdlTaskRequest::new(DdlTask::new_drop_database( - "greptime".to_string(), - "metrics".to_string(), - false, - )), - ) - .await - .unwrap(); - assert!( - request_rx - .recv() - .await - .unwrap() - .into_inner() - .actor - .is_none() - ); - server_handle.abort(); } diff --git a/src/meta-srv/src/event/gc.rs b/src/meta-srv/src/event/gc.rs index 3575ac20c2f..7f679f1785b 100644 --- a/src/meta-srv/src/event/gc.rs +++ b/src/meta-srv/src/event/gc.rs @@ -212,8 +212,8 @@ mod tests { use api::v1::ColumnSchema; use common_event_recorder::event_table::{ - EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, - PROCEDURE_TRIGGER_COLUMN, jsonb_value, + ACTOR_COLUMN, EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, + PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, jsonb_value, }; use common_event_recorder::testing::assert_event_contract; use common_event_recorder::{EventTypeFilter, PersistentEventContext, TriggerReason}; @@ -536,6 +536,7 @@ mod tests { ); let mut event_schema = procedure_schema(); event_schema.extend(schema()); + event_schema.push(ACTOR_COLUMN.column_schema()); event_schema.push(EVENT_CONTEXT_COLUMN.column_schema()); let mut values = vec![ ValueData::StringValue(procedure_id.to_string()).into(), @@ -554,6 +555,7 @@ mod tests { .values, ); values.push(Value { value_data: None }); + values.push(Value { value_data: None }); assert_event_contract( &event, diff --git a/src/meta-srv/src/event/repartition.rs b/src/meta-srv/src/event/repartition.rs index 1faf3f18212..c1e18fd565e 100644 --- a/src/meta-srv/src/event/repartition.rs +++ b/src/meta-srv/src/event/repartition.rs @@ -364,8 +364,8 @@ mod tests { use api::v1::{ColumnSchema, Row, Value}; use common_event_recorder::Event; use common_event_recorder::event_table::{ - EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, PROCEDURE_STATE_COLUMN, - PROCEDURE_TRIGGER_COLUMN, jsonb_value, + ACTOR_COLUMN, EVENT_CONTEXT_COLUMN, PROCEDURE_ERROR_COLUMN, PROCEDURE_ID_COLUMN, + PROCEDURE_STATE_COLUMN, PROCEDURE_TRIGGER_COLUMN, jsonb_value, }; use common_event_recorder::testing::assert_event_contract; use common_procedure::{EventTrigger, ProcedureEvent, ProcedureId, ProcedureState}; @@ -668,6 +668,7 @@ mod tests { ); let mut schema = procedure_schema(); schema.extend(parent_schema()); + schema.push(ACTOR_COLUMN.column_schema()); schema.push(EVENT_CONTEXT_COLUMN.column_schema()); assert_event_contract( @@ -685,6 +686,7 @@ mod tests { ValueData::StringValue("repartition_events".to_string()).into(), ValueData::U32Value(1024).into(), Value { value_data: None }, + Value { value_data: None }, ], }], ); diff --git a/src/meta-srv/src/service/procedure.rs b/src/meta-srv/src/service/procedure.rs index 9e226bc18fe..df27e2cb32b 100644 --- a/src/meta-srv/src/service/procedure.rs +++ b/src/meta-srv/src/service/procedure.rs @@ -55,28 +55,14 @@ struct ProcedureSubmission { event_context: Option, } -impl TryFrom<(Option, Option)> for ProcedureSubmission { - type Error = Status; - - fn try_from( +impl From<(Option, Option)> for ProcedureSubmission { + fn from( (actor, event_context): (Option, Option), - ) -> Result { - let actor = actor - .map(|actor| { - if actor.username.is_empty() { - Err(Status::invalid_argument( - "Procedure actor must not be empty", - )) - } else { - Ok(actor.username) - } - }) - .transpose()?; - - Ok(Self { - actor, + ) -> Self { + Self { + actor: actor.map(|actor| actor.username), event_context, - }) + } } } @@ -139,7 +125,7 @@ impl procedure_service_server::ProcedureService for Metasrv { let ProcedureSubmission { actor, event_context, - } = ProcedureSubmission::try_from((actor, event_context))?; + } = ProcedureSubmission::from((actor, event_context)); let mut query_context = query_context .context(error::MissingRequiredParameterSnafu { param: "query_context", @@ -191,7 +177,7 @@ impl procedure_service_server::ProcedureService for Metasrv { let _header = header.context(error::MissingRequestHeaderSnafu)?; let procedure_context = - ProcedureContext::from(ProcedureSubmission::try_from((actor, event_context))?); + ProcedureContext::from(ProcedureSubmission::from((actor, event_context))); let from_peer = self .lookup_datanode_peer(from_peer) .await? @@ -329,7 +315,7 @@ impl procedure_service_server::ProcedureService for Metasrv { let _header = header.context(error::MissingRequestHeaderSnafu)?; let procedure_context = - ProcedureContext::from(ProcedureSubmission::try_from((actor, event_context))?); + ProcedureContext::from(ProcedureSubmission::from((actor, event_context))); let response = self .handle_gc_regions( @@ -361,7 +347,7 @@ impl procedure_service_server::ProcedureService for Metasrv { let _header = header.context(error::MissingRequestHeaderSnafu)?; let procedure_context = - ProcedureContext::from(ProcedureSubmission::try_from((actor, event_context))?); + ProcedureContext::from(ProcedureSubmission::from((actor, event_context))); let response = self .handle_gc_table( @@ -593,15 +579,14 @@ fn gc_response_to_table_pb(resp: GcResponse) -> GcTableResponse { mod tests { use std::time::Duration; - use api::v1::meta::{ProcedureActor, ProcedureEventContext, Role}; + use api::v1::meta::Role; use common_meta::rpc::ddl::{ CREATE_DATABASE_CREATOR_EXTENSION_KEY, CREATE_DATABASE_CREATOR_METADATA_KEY, CreatorGrantIntent, DdlTask, QueryContext, }; - use common_procedure::ProcedureContext; use tonic::metadata::{MetadataMap, MetadataValue}; - use super::{Metasrv, ProcedureSubmission, restore_create_database_creator}; + use super::{Metasrv, restore_create_database_creator}; fn create_database_task() -> DdlTask { DdlTask::new_create_database( @@ -738,41 +723,4 @@ mod tests { assert_eq!(task.creator, Some(creator)); assert!(trusted.extensions.is_empty()); } - - #[test] - fn test_procedure_submission_rejects_empty_actor() { - let submission = ProcedureSubmission::try_from(( - Some(ProcedureActor { - username: "alice".to_string(), - }), - Some(ProcedureEventContext { - reason: "manual".to_string(), - protocol: "postgres".to_string(), - extensions: Default::default(), - }), - )) - .unwrap(); - assert_eq!(submission.actor.as_deref(), Some("alice")); - assert_eq!( - submission - .event_context - .as_ref() - .map(|context| context.protocol.as_str()), - Some("postgres") - ); - - assert!( - ProcedureSubmission::try_from(( - Some(ProcedureActor { - username: String::new(), - }), - None, - )) - .is_err() - ); - - let procedure_context = - ProcedureContext::from(ProcedureSubmission::try_from((None, None)).unwrap()); - assert_eq!(procedure_context.event_context, None); - } } diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index d124c5c6cdb..30b9a07348d 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -18,9 +18,13 @@ use std::net::SocketAddr; use std::sync::Arc; use std::time::Duration; -use auth::{DefaultPermissionChecker, PermissionCheckerRef, UserProviderRef}; +use api::v1::Basic; +use auth::{ + DefaultPermissionChecker, PermissionCheckerRef, UserProviderRef, user_provider_from_option, +}; use axum::Router; use catalog::kvbackend::KvBackendCatalogManager; +use client::{Client, Database}; use common_base::Plugins; use common_config::Configurable; use common_meta::key::TableMetadataManager; @@ -732,14 +736,57 @@ async fn setup_grpc_server_for_instance( grpc_config: Option, memory_limiter: Option, ) -> (GreptimeDbStandalone, Arc) { + let grpc_server = setup_grpc_server_for_frontend_instance_with( + instance.fe_instance().clone(), + user_provider, + grpc_config, + memory_limiter, + ) + .await; + (instance, grpc_server) +} + +/// Builds and starts an authenticated gRPC server on an existing frontend instance. +pub async fn setup_grpc_server_for_frontend_instance( + instance: Arc, + user_provider: Option, +) -> Arc { + setup_grpc_server_for_frontend_instance_with(instance, user_provider, None, None).await +} + +/// Starts a gRPC server backed by `instance` and returns a client authenticated +/// as the supplied test user. The server handle keeps the listener alive. +pub async fn setup_authenticated_grpc_database( + instance: Arc, + username: &str, + password: &str, +) -> (Database, Arc) { + let user_provider = + user_provider_from_option(&format!("static_user_provider:cmd:{username}={password}")) + .unwrap(); + let grpc_server = setup_grpc_server_for_frontend_instance(instance, Some(user_provider)).await; + let grpc_addr = grpc_server.bind_addr().unwrap().to_string(); + let mut database = + Database::new_with_dbname("greptime-public", Client::with_urls(vec![grpc_addr])); + database.set_auth(api::v1::auth_header::AuthScheme::Basic(Basic { + username: username.to_string(), + password: password.to_string(), + })); + (database, grpc_server) +} + +async fn setup_grpc_server_for_frontend_instance_with( + fe_instance_ref: Arc, + user_provider: Option, + grpc_config: Option, + memory_limiter: Option, +) -> Arc { let runtime: Runtime = RuntimeBuilder::default() .worker_threads(2) .thread_name("grpc-handlers") .build() .unwrap(); - let fe_instance_ref = instance.fe_instance().clone(); - let greptime_request_handler = GreptimeRequestHandler::new( fe_instance_ref.clone(), user_provider.clone(), @@ -769,7 +816,7 @@ async fn setup_grpc_server_for_instance( let fe_grpc_addr = "127.0.0.1:0".parse::().unwrap(); grpc_server.start(fe_grpc_addr).await.unwrap(); - (instance, Arc::new(grpc_server)) + Arc::new(grpc_server) } pub async fn setup_mysql_server( diff --git a/tests-integration/tests/database_ddl_event.rs b/tests-integration/tests/database_ddl_event.rs index 6cbc65a2673..4dfda4f32b9 100644 --- a/tests-integration/tests/database_ddl_event.rs +++ b/tests-integration/tests/database_ddl_event.rs @@ -14,17 +14,21 @@ use std::sync::Arc; -use client::OutputData; +use client::{Database, OutputData}; use common_test_util::temp_dir::create_temp_dir; -use servers::query_handler::sql::SqlQueryHandler; -use session::context::QueryContext; use tests_integration::cluster::GreptimeDbClusterBuilder; use tests_integration::standalone::GreptimeDbStandaloneBuilder; -use tests_integration::test_util::{StorageType, get_test_store_config}; +use tests_integration::test_util::{ + StorageType, get_test_store_config, setup_authenticated_grpc_database, +}; -use crate::event_recorder_test_util::{assert_single_event, find_eventually_string}; +use crate::event_recorder_test_util::{ + assert_procedure_actor, assert_single_event, find_eventually_string, +}; const DATABASE_NAME: &str = "database_ddl_events"; +const PROCEDURE_ACTOR: &str = "procedure_actor"; +const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; #[tokio::test(flavor = "multi_thread")] async fn test_database_ddl_events() { @@ -44,8 +48,14 @@ async fn test_database_ddl_events() { .build(true) .await; let instance = cluster.fe_instance(); + let (database, _grpc_server) = setup_authenticated_grpc_database( + instance.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, + ) + .await; - execute_database_ddl(instance).await; + execute_database_ddl(&database).await; assert_database_events(instance).await; } @@ -57,23 +67,25 @@ async fn test_standalone_database_ddl_events() { .build() .await; let instance = standalone.fe_instance(); + let (database, _grpc_server) = setup_authenticated_grpc_database( + instance.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, + ) + .await; - execute_database_ddl(instance).await; + execute_database_ddl(&database).await; assert_database_events(instance).await; } -async fn execute_database_ddl(instance: &Arc) { +async fn execute_database_ddl(database: &Database) { for sql in [ format!("CREATE DATABASE IF NOT EXISTS {DATABASE_NAME} WITH (ttl = '1h')"), format!("ALTER DATABASE {DATABASE_NAME} SET 'ttl' = '2h'"), format!("DROP DATABASE IF EXISTS {DATABASE_NAME}"), ] { - let output = instance - .do_query(&sql, QueryContext::arc()) - .await - .remove(0) - .unwrap(); + let output = database.sql(sql).await.unwrap(); assert!(matches!(output.data, OutputData::AffectedRows(_))); } } @@ -111,10 +123,27 @@ async fn assert_database_event( event_type: &str, submitted_payload_predicate: &str, ) { + let procedure_id = find_eventually_string( + instance, + &format!( + r#"SELECT procedure_id +FROM greptime_private.events +WHERE type = '{event_type}' + AND schema_name = '{DATABASE_NAME}' + AND json_path_match(procedure_trigger, '$.type == "Submitted"') +ORDER BY timestamp DESC +LIMIT 1"# + ), + "procedure_id", + ) + .await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; + let submitted = format!( r#"SELECT count(*) AS event_count FROM greptime_private.events WHERE type = '{event_type}' + AND procedure_id = '{procedure_id}' AND procedure_state = 'Running' AND json_path_match(procedure_trigger, '$.type == "Submitted"') AND catalog_name = 'greptime' @@ -125,17 +154,18 @@ WHERE type = '{event_type}' let actual = find_eventually_string( instance, &format!( - "SELECT json_to_string(event_context) AS event_context FROM greptime_private.events WHERE type = '{event_type}' AND procedure_state = 'Running' AND json_path_match(procedure_trigger, '$.type == \"Submitted\"') ORDER BY timestamp DESC LIMIT 1" + "SELECT json_to_string(event_context) AS event_context FROM greptime_private.events WHERE procedure_id = '{procedure_id}' AND json_path_match(procedure_trigger, '$.type == \"Submitted\"')" ), "event_context", ) .await; - assert_eq!(r#"{"reason":"manual"}"#, actual); + assert_eq!(r#"{"protocol":"grpc","reason":"manual"}"#, actual); let lifecycle = format!( r#"SELECT count(*) AS event_count FROM greptime_private.events WHERE type = '{event_type}' + AND procedure_id = '{procedure_id}' AND procedure_state = 'Done' AND json_path_match(procedure_trigger, '$.type == "Succeeded"') AND catalog_name = 'greptime' diff --git a/tests-integration/tests/event_recorder_test_util.rs b/tests-integration/tests/event_recorder_test_util.rs index 0441cc08d3b..38bb18861b3 100644 --- a/tests-integration/tests/event_recorder_test_util.rs +++ b/tests-integration/tests/event_recorder_test_util.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::future::Future; use std::sync::Arc; use std::time::Duration; @@ -26,6 +27,62 @@ use session::context::QueryContext; const MAX_ATTEMPTS: usize = 60; const POLL_INTERVAL: Duration = Duration::from_millis(250); +/// Asserts through the supplied SQL executor that a table procedure records +/// the expected actor for both submission and completion events. +pub(crate) async fn assert_procedure_actor_by_table( + event_type: &str, + table_name: &str, + actor: &str, + mut execute: F, +) where + F: FnMut(String) -> Fut, + Fut: Future, +{ + let query = format!( + "SELECT count(DISTINCT json_get_string(procedure_trigger, 'type')) = 2 \ + AND count(*) = count(CASE WHEN actor = '{actor}' THEN 1 END) AS actor_matches \ + FROM greptime_private.events \ + WHERE type = '{event_type}' AND table_name = '{table_name}' \ + AND json_get_string(procedure_trigger, 'type') IN ('Submitted', 'Succeeded')" + ); + + for _ in 0..MAX_ATTEMPTS { + if execute(query.clone()).await { + return; + } + tokio::time::sleep(POLL_INTERVAL).await; + } + + panic!("timed out waiting for procedure actor: {query}"); +} + +/// Asserts that submission and completion events record the expected actor. +pub(crate) async fn assert_procedure_actor( + instance: &Arc, + procedure_id: &str, + actor: Option<&str>, +) { + let actor_predicate = actor.map_or_else( + || "count(actor) = 0".to_string(), + |actor| format!("count(*) = count(CASE WHEN actor = '{actor}' THEN 1 END)"), + ); + for trigger in ["Submitted", "Succeeded"] { + assert_eventually_eq( + instance, + &format!( + "SELECT count(*) > 0 AND {actor_predicate} AS actor_matches FROM greptime_private.events WHERE procedure_id = '{procedure_id}' AND json_get_string(procedure_trigger, 'type') = '{trigger}'" + ), + "\ ++---------------+ +| actor_matches | ++---------------+ +| true | ++---------------+", + ) + .await; + } +} + /// Asserts that the query eventually returns exactly one event. pub(crate) async fn assert_single_event(instance: &Arc, query: &str) { for _ in 0..MAX_ATTEMPTS { diff --git a/tests-integration/tests/flow_ddl_event.rs b/tests-integration/tests/flow_ddl_event.rs index c5fc9be94c9..122c36c84f4 100644 --- a/tests-integration/tests/flow_ddl_event.rs +++ b/tests-integration/tests/flow_ddl_event.rs @@ -14,15 +14,18 @@ use std::sync::Arc; -use servers::query_handler::sql::SqlQueryHandler; -use session::context::QueryContext; use tests_integration::standalone::GreptimeDbStandaloneBuilder; +use tests_integration::test_util::setup_authenticated_grpc_database; use uuid::Uuid; -use crate::event_recorder_test_util::{assert_single_event, find_eventually_string}; +use crate::event_recorder_test_util::{ + assert_procedure_actor, assert_single_event, find_eventually_string, +}; const CREATE_FLOW_EVENT_TYPE: &str = "create_flow"; const DROP_FLOW_EVENT_TYPE: &str = "drop_flow"; +const PROCEDURE_ACTOR: &str = "procedure_actor"; +const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; #[tokio::test(flavor = "multi_thread")] async fn test_standalone_flow_ddl_events() { @@ -35,39 +38,39 @@ async fn test_standalone_flow_ddl_events() { let missing_source = format!("flow_ddl_event_missing_source_{suffix}"); let sink = format!("flow_ddl_event_sink_{suffix}"); let flow = format!("flow_ddl_event_{suffix}"); + let (database, _grpc_server) = setup_authenticated_grpc_database( + instance.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, + ) + .await; - instance - .do_query( - &format!("CREATE TABLE {sink} (val STRING, ts TIMESTAMP TIME INDEX)"), - QueryContext::arc(), - ) + database + .sql(format!( + "CREATE TABLE {sink} (val STRING, ts TIMESTAMP TIME INDEX)" + )) .await - .remove(0) .unwrap(); - instance - .do_query( - &format!( - "CREATE FLOW IF NOT EXISTS {flow} SINK TO {sink} EVAL INTERVAL '10s' \ - WITH (defer_on_missing_source = true) AS SELECT val, ts FROM {missing_source}" - ), - QueryContext::arc(), - ) + database + .sql(format!( + "CREATE FLOW IF NOT EXISTS {flow} SINK TO {sink} EVAL INTERVAL '10s' \ + WITH (defer_on_missing_source = true) AS SELECT val, ts FROM {missing_source}" + )) .await - .remove(0) .unwrap(); assert_create_events(instance, &flow).await; - instance - .do_query(&format!("DROP FLOW IF EXISTS {flow}"), QueryContext::arc()) + database + .sql(format!("DROP FLOW IF EXISTS {flow}")) .await - .remove(0) .unwrap(); assert_drop_events(instance, &flow).await; } async fn assert_create_events(instance: &Arc, flow: &str) { let procedure_id = find_submitted_procedure_id(instance, CREATE_FLOW_EVENT_TYPE, flow).await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; assert_single_event( instance, &format!( @@ -120,11 +123,12 @@ async fn assert_event_context( "event_context", ) .await; - assert_eq!(r#"{"reason":"manual"}"#, actual); + assert_eq!(r#"{"protocol":"grpc","reason":"manual"}"#, actual); } async fn assert_drop_events(instance: &Arc, flow: &str) { let procedure_id = find_submitted_procedure_id(instance, DROP_FLOW_EVENT_TYPE, flow).await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; assert_single_event( instance, &format!( diff --git a/tests-integration/tests/gc_event.rs b/tests-integration/tests/gc_event.rs index 54a0787d12a..4721a442a80 100644 --- a/tests-integration/tests/gc_event.rs +++ b/tests-integration/tests/gc_event.rs @@ -15,19 +15,21 @@ use std::sync::Arc; use std::time::Duration; -use common_event_recorder::{PersistentEventContext, TriggerReason}; -use common_procedure::{ProcedureContext, ProcedureId, ProcedureWithId, watcher}; use common_test_util::temp_dir::create_temp_dir; -use meta_srv::gc::{BatchGcProcedure, GcSchedulerOptions}; +use meta_srv::gc::GcSchedulerOptions; use mito2::gc::GcConfig; use servers::query_handler::sql::SqlQueryHandler; use session::context::QueryContext; use tests_integration::cluster::{GreptimeDbCluster, GreptimeDbClusterBuilder}; -use tests_integration::test_util::{StorageType, get_test_store_config}; +use tests_integration::test_util::{ + StorageType, get_test_store_config, setup_authenticated_grpc_database, +}; -use crate::event_recorder_test_util::find_eventually_string; +use crate::event_recorder_test_util::{assert_procedure_actor, find_eventually_string}; const TABLE_NAME: &str = "batch_gc_events"; +const PROCEDURE_ACTOR: &str = "procedure_actor"; +const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; const MAX_ATTEMPTS: usize = 60; const POLL_INTERVAL: Duration = Duration::from_millis(250); @@ -121,31 +123,28 @@ async fn test_batch_gc_event() { .collect::>(); assert_eq!(regions.len(), 1); - let procedure = BatchGcProcedure::new( - cluster.metasrv.mailbox().clone(), - cluster.metasrv.table_metadata_manager().clone(), - cluster.metasrv.options().grpc.server_addr.clone(), - regions, - false, - Duration::from_secs(10), - Default::default(), - ); - let procedure_id = ProcedureId::parse_str("00000000-0000-0000-0000-00000000bac0").unwrap(); - let mut watcher = cluster - .metasrv - .procedure_manager() - .submit(ProcedureWithId { - id: procedure_id, - procedure: Box::new(procedure), - context: ProcedureContext::from_event_context(PersistentEventContext::new( - TriggerReason::Manual, - )), - }) + let (actor_db, _actor_grpc_server) = setup_authenticated_grpc_database( + instance.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, + ) + .await; + actor_db + .sql(format!("ADMIN GC_REGIONS({})", regions[0].as_u64())) .await .unwrap(); - watcher::wait(&mut watcher).await.unwrap(); assert_sst_count(&cluster, 1).await; + let procedure_id = find_eventually_string( + instance, + &format!( + "SELECT procedure_id FROM greptime_private.events WHERE type = 'batch_gc' AND actor = '{PROCEDURE_ACTOR}' AND json_get_string(procedure_trigger, 'type') = 'Submitted'" + ), + "procedure_id", + ) + .await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; + let submitted = format!( r#"SELECT json_to_string(event_context) AS event_context FROM greptime_private.events @@ -155,7 +154,7 @@ WHERE type = 'batch_gc' AND json_get_string(procedure_trigger, 'type') = 'Submitted'"#, ); assert_eq!( - r#"{"reason":"manual"}"#, + r#"{"protocol":"grpc","reason":"manual"}"#, find_eventually_string(instance, &submitted, "event_context").await ); diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 530b581a2ca..079439bf78c 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -84,6 +84,8 @@ use tests_integration::test_util::{ use urlencoding::encode; use yaml_rust::YamlLoader; +use crate::event_recorder_test_util::assert_procedure_actor_by_table; + #[macro_export] macro_rules! http_test { ($service:ident, $($(#[$meta:meta])* $test:ident),*,) => { @@ -258,6 +260,29 @@ pub async fn test_http_auth(store_type: StorageType) { .send() .await; assert_eq!(res.status(), StatusCode::OK); + let actor_query_client = &client; + assert_procedure_actor_by_table( + "create_table", + "auth_test", + "writeonly_user", + |query| async move { + let res = actor_query_client + .get(&format!("/v1/sql?db=public&sql={}", encode(&query))) + .header("Authorization", basic_auth("greptime_user", "greptime_pwd")) + .send() + .await; + let Ok(body) = serde_json::from_str::(&res.text().await) else { + return false; + }; + matches!( + body.output().first(), + Some(GreptimeQueryOutput::Records(records)) + if records.rows().first().and_then(|row| row.first()).and_then(Value::as_bool) + == Some(true) + ) + }, + ) + .await; let res = client .get("/v1/sql?db=public&sql=insert into auth_test values(1);") .header( diff --git a/tests-integration/tests/region_migration.rs b/tests-integration/tests/region_migration.rs index d67180249c9..17a58de2eed 100644 --- a/tests-integration/tests/region_migration.rs +++ b/tests-integration/tests/region_migration.rs @@ -15,11 +15,11 @@ use std::sync::Arc; use std::time::Duration; -use client::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, OutputData}; +use client::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, Database, OutputData}; use common_catalog::consts::DEFAULT_PRIVATE_SCHEMA_NAME; use common_event_recorder::event_table::{ - REGION_ID_COLUMN, REGION_MIGRATION_DST_NODE_ID_COLUMN, REGION_MIGRATION_SRC_NODE_ID_COLUMN, - REGION_MIGRATION_TRIGGER_REASON_COLUMN, + ACTOR_COLUMN, REGION_ID_COLUMN, REGION_MIGRATION_DST_NODE_ID_COLUMN, + REGION_MIGRATION_SRC_NODE_ID_COLUMN, REGION_MIGRATION_TRIGGER_REASON_COLUMN, }; use common_event_recorder::{ DEFAULT_EVENTS_TABLE_NAME, DEFAULT_FLUSH_INTERVAL_SECONDS, EVENTS_TABLE_TIMESTAMP_COLUMN_NAME, @@ -58,10 +58,14 @@ use session::context::{QueryContext, QueryContextRef}; use store_api::storage::RegionId; use table::metadata::TableId; use tests_integration::cluster::{GreptimeDbCluster, GreptimeDbClusterBuilder}; -use tests_integration::test_util::{PEER_PLACEHOLDER_ADDR, StorageType, get_test_store_config}; +use tests_integration::test_util::{ + PEER_PLACEHOLDER_ADDR, StorageType, get_test_store_config, setup_authenticated_grpc_database, +}; use uuid::Uuid; const TEST_TABLE_NAME: &str = "migration_target"; +const PROCEDURE_ACTOR: &str = "procedure_actor"; +const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; #[macro_export] macro_rules! region_migration_test { @@ -263,6 +267,7 @@ pub async fn test_region_migration(store_type: StorageType, endpoints: Vec String { + let output = database + .sql(format!( + "admin migrate_region({region_id}, {from_peer_id}, {to_peer_id})" + )) + .await + .unwrap(); + let recordbatches = match output.data { + OutputData::RecordBatches(recordbatches) => recordbatches, + OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(), + OutputData::AffectedRows(_) => unreachable!(), + }; + + let record_batch = &recordbatches.take()[0]; + record_batch + .column(0) + .as_string::() + .value(0) + .to_string() +} + /// Query procedure state by SQL. async fn query_procedure_by_sql(instance: &Arc, pid: &str) -> String { let OutputData::RecordBatches(recordbatches) = run_sql( @@ -1331,6 +1373,7 @@ async fn run_sql( } enum RegionMigrationEvents { + Actor, ProcedureId, Timestamp, ProcedureState, @@ -1349,6 +1392,7 @@ impl Iden for RegionMigrationEvents { s, "{}", match self { + Self::Actor => ACTOR_COLUMN.name(), Self::ProcedureId => EVENTS_TABLE_PROCEDURE_ID_COLUMN_NAME, Self::Timestamp => EVENTS_TABLE_TIMESTAMP_COLUMN_NAME, Self::ProcedureState => EVENTS_TABLE_PROCEDURE_STATE_COLUMN_NAME, @@ -1371,6 +1415,7 @@ async fn check_region_migration_events_system_table( region_id: u64, from_peer_id: u64, to_peer_id: u64, + actor: Option<&str>, ) { // Sleep enough time to ensure the event is recorded. tokio::time::sleep(DEFAULT_FLUSH_INTERVAL_SECONDS * 2).await; @@ -1386,7 +1431,8 @@ async fn check_region_migration_events_system_table( // region_migration_src_node_id = ${from_peer_id} AND // region_migration_dst_node_id = ${to_peer_id} // ORDER BY timestamp ASC - let query = Query::select() + let mut query = Query::select(); + query .column(RegionMigrationEvents::RegionMigrationTriggerReason) .column(RegionMigrationEvents::ProcedureState) .expr_as( @@ -1402,7 +1448,13 @@ async fn check_region_migration_events_system_table( .and_where(Expr::col(RegionMigrationEvents::ProcedureId).eq(procedure_id)) .and_where(Expr::col(RegionMigrationEvents::RegionId).eq(region_id)) .and_where(Expr::col(RegionMigrationEvents::SrcNodeId).eq(from_peer_id)) - .and_where(Expr::col(RegionMigrationEvents::DstNodeId).eq(to_peer_id)) + .and_where(Expr::col(RegionMigrationEvents::DstNodeId).eq(to_peer_id)); + if let Some(actor) = actor { + query.and_where(Expr::col(RegionMigrationEvents::Actor).eq(actor)); + } else { + query.and_where(Expr::col(RegionMigrationEvents::Actor).is_null()); + } + let query = query .order_by(RegionMigrationEvents::Timestamp, Order::Asc) .to_string(PostgresQueryBuilder); @@ -1411,12 +1463,22 @@ async fn check_region_migration_events_system_table( .await .remove(0); - let expected = "\ + let expected = if actor == Some(PROCEDURE_ACTOR) { + "\ ++---------------------------------+-----------------+-------------------+---------------------------------------+ +| region_migration_trigger_reason | procedure_state | procedure_trigger | event_context | ++---------------------------------+-----------------+-------------------+---------------------------------------+ +| Manual | Running | Submitted | {\"protocol\":\"grpc\",\"reason\":\"manual\"} | +| Manual | Done | Succeeded | | ++---------------------------------+-----------------+-------------------+---------------------------------------+" + } else { + "\ +---------------------------------+-----------------+-------------------+---------------------+ | region_migration_trigger_reason | procedure_state | procedure_trigger | event_context | +---------------------------------+-----------------+-------------------+---------------------+ | Manual | Running | Submitted | {\"reason\":\"manual\"} | | Manual | Done | Succeeded | | -+---------------------------------+-----------------+-------------------+---------------------+"; ++---------------------------------+-----------------+-------------------+---------------------+" + }; check_output_stream(result.unwrap().data, expected).await; } diff --git a/tests-integration/tests/repartition_event.rs b/tests-integration/tests/repartition_event.rs index 65503a4a8ee..906eb3cfee3 100644 --- a/tests-integration/tests/repartition_event.rs +++ b/tests-integration/tests/repartition_event.rs @@ -15,17 +15,22 @@ use std::sync::Arc; use std::time::Duration; +use client::Database; use common_test_util::temp_dir::create_temp_dir; use meta_srv::gc::GcSchedulerOptions; use mito2::gc::GcConfig; -use servers::query_handler::sql::SqlQueryHandler; -use session::context::QueryContext; use tests_integration::cluster::GreptimeDbClusterBuilder; -use tests_integration::test_util::{StorageType, get_test_store_config}; +use tests_integration::test_util::{ + StorageType, get_test_store_config, setup_authenticated_grpc_database, +}; -use crate::event_recorder_test_util::{assert_eventually_eq, find_eventually_string}; +use crate::event_recorder_test_util::{ + assert_eventually_eq, assert_procedure_actor, find_eventually_string, +}; const TABLE_NAME: &str = "repartition_events"; +const PROCEDURE_ACTOR: &str = "procedure_actor"; +const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; #[tokio::test(flavor = "multi_thread")] async fn test_repartition_event() { @@ -56,11 +61,17 @@ async fn test_repartition_event() { .build(true) .await; let instance = cluster.fe_instance(); + let (database, _grpc_server) = setup_authenticated_grpc_database( + instance.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, + ) + .await; - execute_partition(instance).await; + execute_partition(&database).await; let partition_procedure_id = assert_repartition_event(instance).await; assert_repartition_group_event(instance, &partition_procedure_id).await; - execute_merge(instance).await; + execute_merge(&database).await; let merge_procedure_id = find_submitted_procedure_id( instance, "repartition", @@ -70,10 +81,11 @@ async fn test_repartition_event() { ), ) .await; + assert_procedure_actor(instance, &merge_procedure_id, Some(PROCEDURE_ACTOR)).await; assert_merge_repartition_group_event(instance, &merge_procedure_id).await; } -async fn execute_partition(instance: &Arc) { +async fn execute_partition(database: &Database) { for sql in [ format!( "CREATE TABLE {TABLE_NAME} (host INT, ts TIMESTAMP TIME INDEX, PRIMARY KEY(host)) ENGINE = mito" @@ -83,21 +95,13 @@ async fn execute_partition(instance: &Arc) { ), format!("ALTER TABLE {TABLE_NAME} PARTITION ON COLUMNS (host) (host < 10, host >= 10)"), ] { - instance - .do_query(&sql, QueryContext::arc()) - .await - .remove(0) - .unwrap(); + database.sql(sql).await.unwrap(); } } -async fn execute_merge(instance: &Arc) { +async fn execute_merge(database: &Database) { let sql = format!("ALTER TABLE {TABLE_NAME} MERGE PARTITION (host < 10, host >= 10)"); - instance - .do_query(&sql, QueryContext::arc()) - .await - .remove(0) - .unwrap(); + database.sql(sql).await.unwrap(); } async fn assert_repartition_event(instance: &Arc) -> String { @@ -110,6 +114,7 @@ async fn assert_repartition_event(instance: &Arc) ), ) .await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; let query = format!( r#"SELECT type, catalog_name, schema_name, table_name, table_id, json_path_match(payload, '$.version == 2') AS payload_version, @@ -168,6 +173,7 @@ async fn assert_repartition_group_event( ), ) .await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; let query = format!( r#"SELECT type, catalog_name, schema_name, table_name, table_id, parent_procedure_id = '{parent_procedure_id}' AS matches_parent_procedure_id, @@ -235,6 +241,7 @@ async fn assert_merge_repartition_group_event( ), ) .await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; let query = format!( r#"SELECT type, catalog_name, schema_name, table_name, table_id, parent_procedure_id = '{parent_procedure_id}' AS matches_parent_procedure_id, diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index 58df5424d39..e46a8f58d1b 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -35,6 +35,8 @@ use tests_integration::test_util::{ }; use tokio_postgres::{Client, NoTls, SimpleQueryMessage}; +use crate::event_recorder_test_util::assert_procedure_actor_by_table; + #[macro_export] macro_rules! sql_test { ($service:ident, $($(#[$meta:meta])* $test:ident),*,) => { @@ -148,6 +150,7 @@ pub async fn test_mysql_auth(store_type: StorageType) { .await; assert!(conn_re.is_ok()); + let read_pool = conn_re.unwrap(); // 4. readonly user let conn_re = MySqlPoolOptions::new() @@ -181,6 +184,19 @@ pub async fn test_mysql_auth(store_type: StorageType) { .execute("CREATE TABLE test (ts timestamp time index)") .await .unwrap(); + let actor_query_pool = &read_pool; + assert_procedure_actor_by_table( + "create_table", + "test", + "writeonly_user", + |query| async move { + sqlx::query_scalar::<_, bool>(&query) + .fetch_one(actor_query_pool) + .await + .unwrap_or(false) + }, + ) + .await; let err = pool.execute("SHOW TABLES").await.unwrap_err(); assert!( err.to_string() @@ -466,8 +482,10 @@ pub async fn test_mysql_timezone(store_type: StorageType) { } pub async fn test_postgres_auth(store_type: StorageType) { - let user_provider = - user_provider_from_option("static_user_provider:cmd:greptime_user=greptime_pwd").unwrap(); + let user_provider = user_provider_from_option( + "static_user_provider:cmd:greptime_user=greptime_pwd,writeonly_user:wo=writeonly_pwd", + ) + .unwrap(); let (mut guard, fe_pg_server) = setup_pg_server_with_user_provider(store_type, "sql_crud", Some(user_provider)).await; @@ -518,6 +536,32 @@ pub async fn test_postgres_auth(store_type: StorageType) { .await; assert!(conn_re.is_ok()); + let read_pool = conn_re.unwrap(); + + let write_pool = PgPoolOptions::new() + .max_connections(2) + .connect(&format!( + "postgres://writeonly_user:writeonly_pwd@{addr}/public" + )) + .await + .unwrap(); + write_pool + .execute("CREATE TABLE postgres_auth_actor (ts timestamp time index)") + .await + .unwrap(); + let actor_query_pool = &read_pool; + assert_procedure_actor_by_table( + "create_table", + "postgres_auth_actor", + "writeonly_user", + |query| async move { + sqlx::query_scalar::<_, bool>(&query) + .fetch_one(actor_query_pool) + .await + .unwrap_or(false) + }, + ) + .await; let _ = fe_pg_server.shutdown().await; guard.remove_all().await; diff --git a/tests-integration/tests/table_ddl_event.rs b/tests-integration/tests/table_ddl_event.rs index afd41b844ac..8bad66dae3c 100644 --- a/tests-integration/tests/table_ddl_event.rs +++ b/tests-integration/tests/table_ddl_event.rs @@ -41,9 +41,11 @@ use servers::query_handler::grpc::GrpcQueryHandler; use servers::query_handler::sql::SqlQueryHandler; use session::context::{Channel, QueryContext}; use tests_integration::cluster::GreptimeDbClusterBuilder; +use tests_integration::test_util::setup_authenticated_grpc_database; use crate::event_recorder_test_util::{ - assert_eventually_eq, assert_single_event, find_eventually_string, find_eventually_u32, + assert_eventually_eq, assert_procedure_actor, assert_single_event, find_eventually_string, + find_eventually_u32, }; const EVENTS_TABLE: &str = "greptime_private.events"; @@ -53,6 +55,8 @@ const LOGICAL_TABLE: &str = "table_ddl_events_logical"; const AUTO_TABLE: &str = "table_ddl_events_auto"; const AUTO_INFLUX_TABLE: &str = "table_ddl_events_auto_influx"; const GRPC_TABLE: &str = "table_ddl_events_grpc"; +const PROCEDURE_ACTOR: &str = "procedure_actor"; +const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; #[tokio::test(flavor = "multi_thread")] async fn test_table_ddl_procedure_events() { @@ -150,6 +154,7 @@ async fn test_table_ddl_procedure_events() { .unwrap(); let grpc_create_procedure_id = submitted_procedure_id(&frontend, "create_table", GRPC_TABLE).await; + assert_procedure_actor(&frontend, &grpc_create_procedure_id, Some("greptime")).await; assert_event_context( &frontend, "create_table", @@ -200,26 +205,27 @@ async fn test_table_ddl_procedure_events() { // Act / Assert: Create Table retains its rich submitted row and records the // created table ID when it completes. - run_sql_with_context( - &frontend, - &format!( - "CREATE TABLE {TABLE} (host STRING PRIMARY KEY, ts TIMESTAMP TIME INDEX, val DOUBLE)" - ), - Arc::new(QueryContext::with_channel( - "greptime", - "public", - Channel::HttpSql, - )), + let (actor_db, _actor_grpc_server) = setup_authenticated_grpc_database( + frontend.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, ) .await; + actor_db + .sql(format!( + "CREATE TABLE {TABLE} (host STRING PRIMARY KEY, ts TIMESTAMP TIME INDEX, val DOUBLE)" + )) + .await + .unwrap(); let table_id = find_table_id(&frontend, TABLE).await; let create_table_procedure_id = submitted_procedure_id(&frontend, "create_table", TABLE).await; + assert_procedure_actor(&frontend, &create_table_procedure_id, Some(PROCEDURE_ACTOR)).await; assert_event_context( &frontend, "create_table", &create_table_procedure_id, "manual", - Some("httpsql"), + Some("grpc"), ) .await; assert_named_submitted_event( diff --git a/tests-integration/tests/view_ddl_event.rs b/tests-integration/tests/view_ddl_event.rs index fce33d09bc8..b21838995c3 100644 --- a/tests-integration/tests/view_ddl_event.rs +++ b/tests-integration/tests/view_ddl_event.rs @@ -16,17 +16,21 @@ use std::sync::Arc; use common_test_util::temp_dir::create_temp_dir; use common_wal::config::DatanodeWalConfig; -use servers::query_handler::sql::SqlQueryHandler; -use session::context::QueryContext; use tests_integration::cluster::GreptimeDbClusterBuilder; use tests_integration::standalone::GreptimeDbStandaloneBuilder; -use tests_integration::test_util::{StorageType, get_test_store_config}; +use tests_integration::test_util::{ + StorageType, get_test_store_config, setup_authenticated_grpc_database, +}; use uuid::Uuid; -use crate::event_recorder_test_util::{assert_single_event, find_eventually_string}; +use crate::event_recorder_test_util::{ + assert_procedure_actor, assert_single_event, find_eventually_string, +}; const CREATE_VIEW_EVENT_TYPE: &str = "create_view"; const DROP_VIEW_EVENT_TYPE: &str = "drop_view"; +const PROCEDURE_ACTOR: &str = "procedure_actor"; +const PROCEDURE_ACTOR_PASSWORD: &str = "procedure_actor_pwd"; #[tokio::test(flavor = "multi_thread")] async fn test_view_ddl_events() { @@ -72,37 +76,35 @@ async fn execute_view_ddl( source_table: &str, view: &str, ) { - instance - .do_query( - &format!( - "CREATE TABLE {source_table} (host STRING PRIMARY KEY, amount DOUBLE, ts TIMESTAMP TIME INDEX)" - ), - QueryContext::arc(), - ) + let (database, _grpc_server) = setup_authenticated_grpc_database( + instance.clone(), + PROCEDURE_ACTOR, + PROCEDURE_ACTOR_PASSWORD, + ) + .await; + + database + .sql(format!( + "CREATE TABLE {source_table} (host STRING PRIMARY KEY, amount DOUBLE, ts TIMESTAMP TIME INDEX)" + )) .await - .remove(0) .unwrap(); - instance - .do_query( - &format!("CREATE VIEW {view} AS SELECT amount FROM {source_table}"), - QueryContext::arc(), - ) + database + .sql(format!( + "CREATE VIEW {view} AS SELECT amount FROM {source_table}" + )) .await - .remove(0) .unwrap(); assert_create_events(instance, view).await; - instance - .do_query(&format!("DROP VIEW {view}"), QueryContext::arc()) - .await - .remove(0) - .unwrap(); + database.sql(format!("DROP VIEW {view}")).await.unwrap(); assert_drop_events(instance, view).await; } async fn assert_create_events(instance: &Arc, view: &str) { let procedure_id = find_submitted_procedure_id(instance, CREATE_VIEW_EVENT_TYPE, view).await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; assert_single_event( instance, &format!( @@ -157,11 +159,12 @@ async fn assert_event_context( "event_context", ) .await; - assert_eq!(r#"{"reason":"manual"}"#, actual); + assert_eq!(r#"{"protocol":"grpc","reason":"manual"}"#, actual); } async fn assert_drop_events(instance: &Arc, view: &str) { let procedure_id = find_submitted_procedure_id(instance, DROP_VIEW_EVENT_TYPE, view).await; + assert_procedure_actor(instance, &procedure_id, Some(PROCEDURE_ACTOR)).await; assert_single_event( instance, &format!( diff --git a/tests-integration/tests/wal_prune_event.rs b/tests-integration/tests/wal_prune_event.rs index 05ceb1b92d9..db86ade4d9d 100644 --- a/tests-integration/tests/wal_prune_event.rs +++ b/tests-integration/tests/wal_prune_event.rs @@ -100,6 +100,7 @@ WHERE type = 'wal_prune' AND procedure_id = '{procedure_id}' AND procedure_state = 'Done' AND json_path_match(procedure_trigger, '$.type == "Succeeded"') + AND actor IS NULL AND topic_name = '{topic_name}' AND prunable_entry_id = {pruned_entry_id} AND latest_offset = 3 From 072810159a6c19b153bcd59f9abbbc3fe1f39462 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Fri, 14 Aug 2026 09:40:50 +0000 Subject: [PATCH 52/92] chore: add v2 version label to prom metrics (#8885) Signed-off-by: shuiyisong --- src/servers/src/http/prom_store.rs | 28 ++++++++++++------------ src/servers/src/metrics.rs | 23 ++++++++----------- src/servers/src/prom_remote_write/mod.rs | 9 +++++++- src/servers/src/prom_remote_write/v2.rs | 10 ++++++--- src/servers/src/prom_store.rs | 5 ++++- 5 files changed, 42 insertions(+), 33 deletions(-) diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index ed960264b0d..1f0d26789eb 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -47,9 +47,11 @@ use crate::http::header::{ }; use crate::pending_rows_batcher::PendingRowsBatcher; use crate::prom_remote_write::decode::PromSeriesProcessor; -use crate::prom_remote_write::decode_remote_write_request; use crate::prom_remote_write::v2::decode_remote_write_v2; use crate::prom_remote_write::validation::PromValidationMode; +use crate::prom_remote_write::{ + REMOTE_WRITE_V1_VERSION, REMOTE_WRITE_V2_VERSION, decode_remote_write_request, +}; use crate::prom_store::snappy_decompress; use crate::query_handler::{PipelineHandlerRef, PromStoreProtocolHandlerRef, PromStoreResponse}; @@ -57,8 +59,6 @@ pub const PHYSICAL_TABLE_PARAM: &str = "physical_table"; pub const DEFAULT_ENCODING: &str = "snappy"; pub const VM_ENCODING: &str = "zstd"; pub const VM_PROTO_VERSION: &str = "1"; -const REMOTE_WRITE_V1_VERSION: &str = "1.0"; -const REMOTE_WRITE_V2_VERSION: &str = "2.0"; const REMOTE_WRITE_V1_PROTO: &str = "prometheus.WriteRequest"; const REMOTE_WRITE_V2_PROTO: &str = "io.prometheus.write.v2.Request"; const CONTENT_TYPE_PROTO_PARAM: &str = "proto"; @@ -194,11 +194,11 @@ async fn remote_write_v1( { Ok(outcome) => outcome, Err(error) => { - record_remote_write_samples(&db, error.rows_written); + record_remote_write_samples(&db, REMOTE_WRITE_V1_VERSION, error.rows_written); return Err(error.error); } }; - record_remote_write_samples(&db, outcome.rows_written); + record_remote_write_samples(&db, REMOTE_WRITE_V1_VERSION, outcome.rows_written); Ok(( StatusCode::NO_CONTENT, @@ -264,8 +264,8 @@ async fn remote_write_v2( { Ok(outcome) => outcome, Err(error) => { - record_remote_write_samples(&db, error.samples_written); - record_remote_write_histograms(&db, error.histograms_written); + record_remote_write_samples(&db, REMOTE_WRITE_V2_VERSION, error.samples_written); + record_remote_write_histograms(&db, REMOTE_WRITE_V2_VERSION, error.histograms_written); return Ok(remote_write_v2_error_response( error.error, error.samples_written, @@ -276,8 +276,8 @@ async fn remote_write_v2( }; debug_assert_eq!(outcome.samples_written, sample_count); debug_assert_eq!(outcome.histograms_written, histogram_count); - record_remote_write_samples(&db, outcome.samples_written); - record_remote_write_histograms(&db, outcome.histograms_written); + record_remote_write_samples(&db, REMOTE_WRITE_V2_VERSION, outcome.samples_written); + record_remote_write_histograms(&db, REMOTE_WRITE_V2_VERSION, outcome.histograms_written); let mut headers = write_cost_header_map(outcome.write_cost); append_remote_write_v2_written_headers( @@ -322,7 +322,7 @@ fn prepare_remote_write_context( query_ctx.set_extension(SEMANTIC_SOURCE_VERSION, remote_write_version); query_ctx.set_extension(SEMANTIC_METRIC_METADATA_QUALITY, METADATA_QUALITY_INFERRED); let timer = crate::metrics::METRIC_HTTP_PROM_STORE_WRITE_ELAPSED - .with_label_values(&[db.as_str()]) + .with_label_values(&[db.as_str(), remote_write_version]) .start_timer(); (db, query_ctx, timer) @@ -608,21 +608,21 @@ fn incomplete_prom_write_error() -> error::Error { .build() } -fn record_remote_write_samples(db: &str, rows: u64) { +fn record_remote_write_samples(db: &str, version: &str, rows: u64) { if rows == 0 { return; } crate::metrics::PROM_STORE_REMOTE_WRITE_SAMPLES - .with_label_values(&[db]) + .with_label_values(&[db, version]) .inc_by(rows); } -fn record_remote_write_histograms(db: &str, rows: u64) { +fn record_remote_write_histograms(db: &str, version: &str, rows: u64) { if rows == 0 { return; } crate::metrics::PROM_STORE_REMOTE_WRITE_HISTOGRAMS - .with_label_values(&[db]) + .with_label_values(&[db, version]) .inc_by(rows); } diff --git a/src/servers/src/metrics.rs b/src/servers/src/metrics.rs index 41a16bfea23..0d70540be31 100644 --- a/src/servers/src/metrics.rs +++ b/src/servers/src/metrics.rs @@ -44,6 +44,7 @@ pub(crate) const METRIC_POSTGRES_EXTENDED_QUERY: &str = "extended"; pub(crate) const METRIC_METHOD_LABEL: &str = "method"; pub(crate) const METRIC_PATH_LABEL: &str = "path"; pub(crate) const METRIC_RESULT_LABEL: &str = "result"; +pub(crate) const METRIC_VERSION_LABEL: &str = "version"; pub(crate) const METRIC_SUCCESS_VALUE: &str = "success"; pub(crate) const METRIC_FAILURE_VALUE: &str = "failure"; @@ -100,39 +101,33 @@ lazy_static! { vec![0.005, 0.01, 0.05, 0.1, 0.5, 1.0, 5.0, 10.0, 60.0, 300.0] ) .unwrap(); - /// Http prometheus write duration per database. + /// Http prometheus write duration per database and remote write protocol version. pub static ref METRIC_HTTP_PROM_STORE_WRITE_ELAPSED: HistogramVec = register_histogram_vec!( "greptime_servers_http_prometheus_write_elapsed", "servers http prometheus write elapsed", - &[METRIC_DB_LABEL], + &[METRIC_DB_LABEL, METRIC_VERSION_LABEL], vec![0.005, 0.01, 0.05, 0.1, 0.5, 1.0, 5.0, 10.0, 60.0, 300.0] ) .unwrap(); - /// Prometheus remote write codec duration. + /// Prometheus remote write codec duration per protocol version. pub static ref METRIC_HTTP_PROM_STORE_CODEC_ELAPSED: HistogramVec = register_histogram_vec!( "greptime_servers_http_prometheus_codec_elapsed", "servers http prometheus request codec duration", - &["type"], + &["type", METRIC_VERSION_LABEL], ) .unwrap(); - /// Decode duration of prometheus write request. - pub static ref METRIC_HTTP_PROM_STORE_DECODE_ELAPSED: Histogram = METRIC_HTTP_PROM_STORE_CODEC_ELAPSED - .with_label_values(&["decode"]); - /// Duration to convert prometheus write request to gRPC request. - pub static ref METRIC_HTTP_PROM_STORE_CONVERT_ELAPSED: Histogram = METRIC_HTTP_PROM_STORE_CODEC_ELAPSED - .with_label_values(&["convert"]); - /// The samples count of Prometheus remote write. + /// The samples count of Prometheus remote write per protocol version. pub static ref PROM_STORE_REMOTE_WRITE_SAMPLES: IntCounterVec = register_int_counter_vec!( "greptime_servers_prometheus_remote_write_samples", "frontend prometheus remote write samples", - &[METRIC_DB_LABEL] + &[METRIC_DB_LABEL, METRIC_VERSION_LABEL] ) .unwrap(); - /// The native histograms count of Prometheus remote write. + /// The native histograms count of Prometheus remote write per protocol version. pub static ref PROM_STORE_REMOTE_WRITE_HISTOGRAMS: IntCounterVec = register_int_counter_vec!( "greptime_servers_prometheus_remote_write_histograms", "frontend prometheus remote write native histograms", - &[METRIC_DB_LABEL] + &[METRIC_DB_LABEL, METRIC_VERSION_LABEL] ) .unwrap(); pub static ref PENDING_BATCHES: IntGauge = register_int_gauge!( diff --git a/src/servers/src/prom_remote_write/mod.rs b/src/servers/src/prom_remote_write/mod.rs index 1c6b972e3db..73f968b6084 100644 --- a/src/servers/src/prom_remote_write/mod.rs +++ b/src/servers/src/prom_remote_write/mod.rs @@ -37,6 +37,11 @@ use crate::prom_remote_write::row_builder::TablesBuilder; use crate::prom_remote_write::validation::PromValidationMode; use crate::prom_store::{snappy_decompress, zstd_decompress}; +/// Prometheus remote write protocol versions, also used as the `version` label +/// of the remote write metrics. +pub const REMOTE_WRITE_V1_VERSION: &str = "1.0"; +pub const REMOTE_WRITE_V2_VERSION: &str = "2.0"; + lazy_static! { static ref PROM_WRITE_REQUEST_POOL: Pool> = Pool::new(256, PromWriteRequest::default); @@ -56,7 +61,9 @@ pub fn decode_remote_write_request( prom_validation_mode: PromValidationMode, processor: &mut PromSeriesProcessor, ) -> crate::error::Result> { - let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_DECODE_ELAPSED.start_timer(); + let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED + .with_label_values(&["decode", REMOTE_WRITE_V1_VERSION]) + .start_timer(); // due to vmagent's limitation, there is a chance that vmagent is // sending content type wrong so we have to apply a fallback with decoding diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 31da7ac1f35..1bcec98d5f5 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -43,8 +43,8 @@ use table::requests::{ use crate::error::{self, Result}; use crate::prom_remote_write::row_builder::PromCtx; -use crate::prom_remote_write::try_decompress; use crate::prom_remote_write::validation::validate_label_name; +use crate::prom_remote_write::{REMOTE_WRITE_V2_VERSION, try_decompress}; #[allow(deprecated)] use crate::prom_store::{ DATABASE_LABEL, DATABASE_LABEL_ALT, METRIC_NAME_LABEL, PHYSICAL_TABLE_LABEL, @@ -128,7 +128,9 @@ pub(crate) fn decode_remote_write_v2( body: Bytes, native_histograms_enabled: bool, ) -> Result { - let decode_timer = crate::metrics::METRIC_HTTP_PROM_STORE_DECODE_ELAPSED.start_timer(); + let decode_timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED + .with_label_values(&["decode", REMOTE_WRITE_V2_VERSION]) + .start_timer(); // Match the v1 decoder's VictoriaMetrics fallback: some clients may send a // mismatched content-encoding header, so try the other compression on failure. @@ -140,7 +142,9 @@ pub(crate) fn decode_remote_write_v2( let request = BorrowedRequest::decode(&buf).context(error::DecodePromRemoteRequestSnafu)?; drop(decode_timer); - let _convert_timer = crate::metrics::METRIC_HTTP_PROM_STORE_CONVERT_ELAPSED.start_timer(); + let _convert_timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED + .with_label_values(&["convert", REMOTE_WRITE_V2_VERSION]) + .start_timer(); convert_remote_write_v2(request, native_histograms_enabled) } diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index 33707b29358..c4f2f929cae 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -39,6 +39,7 @@ use snafu::{OptionExt, ResultExt, ensure}; use snap::raw::{Decoder, Encoder}; use crate::error::{self, Result}; +use crate::prom_remote_write::REMOTE_WRITE_V1_VERSION; use crate::row_writer::{self, MultiTableData}; pub const METRIC_NAME_LABEL: &str = "__name__"; @@ -458,7 +459,9 @@ fn recordbatch_to_timeseries( } pub fn to_grpc_row_insert_requests(request: &WriteRequest) -> Result<(RowInsertRequests, usize)> { - let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_CONVERT_ELAPSED.start_timer(); + let _timer = crate::metrics::METRIC_HTTP_PROM_STORE_CODEC_ELAPSED + .with_label_values(&["convert", REMOTE_WRITE_V1_VERSION]) + .start_timer(); let mut multi_table_data = MultiTableData::new(); From 0ef521d28736644b15e0eaaff80379bac49d56c7 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Fri, 14 Aug 2026 10:17:43 +0000 Subject: [PATCH 53/92] docs: refresh README product copy (#8886) * docs: refresh README product copy Align the README with the current product messaging on the website: - Drop the Observability 2.0 / wide events framing and the unqualified cost multiples from the headline and feature table. - Remove the claim that the open-source distribution scales reads with horizontal replicas; read replicas are Enterprise. - Replace the competitor comparison table with a per-protocol compatibility matrix, since query-side coverage is narrower than ingestion (no LogQL, partial QueryDSL in Enterprise only). - Add a section stating what is in the Apache-2.0 build and what Enterprise adds. - Fix the canary badge caption, add gRPC to the ingest list. Signed-off-by: Dennis Zhuang * docs: link the Enterprise overview instead of a full feature list Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- README.md | 85 +++++++++++++++++++++++++++++++------------------------ 1 file changed, 48 insertions(+), 37 deletions(-) diff --git a/README.md b/README.md index 21980b1a22b..e95f7849e4b 100644 --- a/README.md +++ b/README.md @@ -6,10 +6,11 @@

-

One database for metrics, logs, and traces
-replacing Prometheus, Loki, and Elasticsearch

+

Metrics, logs, and traces.
+One engine, on your infrastructure.

-> The unified OpenTelemetry backend — with SQL + PromQL on object storage. +> A columnar database for metrics, logs, and traces on object storage. +> Apache-2.0 licensed core.

@@ -28,7 +29,7 @@ replacing Prometheus, Loki, and Elasticsearch

Nightly -stable for production  ·  latest includes pre-releases  ·  nightly is a weekly snapshot of main +stable for production  ·  canary includes pre-releases  ·  nightly is a weekly snapshot of main Docker Pulls @@ -57,9 +58,11 @@ replacing Prometheus, Loki, and Elasticsearch
- [Introduction](#introduction) +- [Why You Might Use It](#why-you-might-use-it) - [Overview](#overview) -- [Features](#features) -- [How GreptimeDB Compares](#how-greptimedb-compares) +- [What's Supported](#whats-supported) +- [Compatibility and Migration](#compatibility-and-migration) +- [Limitations and Edition Boundary](#limitations-and-edition-boundary) - [Architecture](#architecture) - [Try GreptimeDB](#try-greptimedb) - [Getting Started](#getting-started) @@ -74,9 +77,20 @@ replacing Prometheus, Loki, and Elasticsearch ## Introduction -**GreptimeDB** is an open-source observability database built for [Observability 2.0](https://docs.greptime.com/user-guide/concepts/observability-2/) — treating metrics, logs, and traces as one unified data model (wide events) instead of three separate pillars. +**GreptimeDB** is an open-source observability database. Metrics, logs, and traces run on one columnar engine over object storage and share one [table model](https://docs.greptime.com/user-guide/concepts/data-model/): tags, timestamp, and fields. When signals carry common identifiers such as service, host, or trace ID, you can correlate them in SQL without moving data between databases. -Use it as the single OpenTelemetry backend — replacing Prometheus, Loki, and Elasticsearch with one database built on object storage. Query with SQL and PromQL, scale without pain, cut costs up to 50×. +Ingest through OpenTelemetry, Prometheus Remote Write, Loki Push, or Elasticsearch Bulk. Use SQL across observability data and PromQL for metrics. Migrate ingestion one signal at a time without rebuilding your collectors. + +## Why You Might Use It + +- You run Prometheus plus Loki or Elasticsearch and want one backend instead of three +- You have outgrown Prometheus on cardinality or retention and don't want the Thanos/Mimir operational surface +- You need long retention on object storage without a separate analytics stack +- You want to query telemetry with SQL, not only a domain query language +- You are storing GenAI or agent telemetry ([OTel GenAI conventions](https://opentelemetry.io/docs/specs/semconv/gen-ai/)) alongside infrastructure signals +- You need the same engine and semantics on resource-constrained devices + +Learn more in [Why GreptimeDB](https://docs.greptime.com/user-guide/concepts/why-greptimedb). ## Overview @@ -88,42 +102,38 @@ A quick overview of what GreptimeDB ingests, how it connects to other systems, a

-## Features +## What's Supported -| Feature | Description | -|---------|-------------| -| **Observability 2.0 native** | Logs, metrics, and traces in one engine with [SQL + PromQL](https://docs.greptime.com/user-guide/query-data/overview/). Native [OpenTelemetry](https://docs.greptime.com/user-guide/ingest-data/for-observability/opentelemetry/), [Prometheus remote write](https://docs.greptime.com/user-guide/ingest-data/for-observability/prometheus/), and [Jaeger](https://docs.greptime.com/user-guide/query-data/jaeger/). Migrate one signal at a time, or use as a single backend. | -| **Elastic compute-storage separation** | Scale reads independently with horizontal replicas. Serve high-concurrency workloads from dashboards, alerting, and AI agents — without resharding or data migration. | -| **Sub-second on PB–EB-scale data** | Columnar engine with [fulltext, inverted, and skipping indexes](https://docs.greptime.com/user-guide/manage-data/data-index). Written in Rust. Designed for high-concurrency point queries, not just analytical scans. | -| **50× lower cost** | Object storage (S3, GCS, Azure Blob) as [primary storage](https://docs.greptime.com/user-guide/deployments-administration/configuration/#storage-options), with a tiered cache (memory + local disk) to keep writes and queries fast. | +| | | +|---|---| +| **Ingest** | [OpenTelemetry (OTLP)](https://docs.greptime.com/user-guide/ingest-data/for-observability/opentelemetry/), [Prometheus Remote Write](https://docs.greptime.com/user-guide/ingest-data/for-observability/prometheus/), [Loki Push](https://docs.greptime.com/user-guide/ingest-data/for-observability/loki), [Elasticsearch Bulk](https://docs.greptime.com/user-guide/ingest-data/for-observability/elasticsearch), [InfluxDB line protocol](https://docs.greptime.com/user-guide/protocols/influxdb-line-protocol), [gRPC](https://docs.greptime.com/user-guide/protocols/grpc/) | +| **Query** | [SQL](https://docs.greptime.com/user-guide/query-data/overview/), [PromQL](https://docs.greptime.com/user-guide/query-data/promql/), [Jaeger-compatible trace queries](https://docs.greptime.com/user-guide/query-data/jaeger/), MySQL and PostgreSQL wire protocols | +| **Storage** | S3, GCS, Azure Blob and S3-compatible endpoints as [primary storage](https://docs.greptime.com/user-guide/deployments-administration/configuration/#storage-options), with memory and local-disk caches | +| **Built in** | Retention policies, downsampling, [continuous aggregation](https://docs.greptime.com/user-guide/flow-computation/overview), explicit table partitioning, and [inverted / skipping / fulltext indexes](https://docs.greptime.com/user-guide/manage-data/data-index) | -**Perfect for:** - * Replacing Prometheus + Loki + Elasticsearch with a single observability backend - * Scaling past Prometheus — high cardinality, long-term storage, no Thanos/Mimir overhead - * AI/agent workloads — store GenAI telemetry ([OTel GenAI conventions](https://opentelemetry.io/docs/specs/semconv/gen-ai/)), and serve high-concurrency reads from SRE/developer agents via horizontal read replicas - * Cutting observability costs with object storage (up to 50× savings on traces, 30% on logs) - * Edge-to-cloud observability with unified APIs on resource-constrained devices +Compute and storage are disaggregated: object storage holds the data, while memory and local-disk caches keep recent and frequently queried data close to compute. -> **Why Observability 2.0?** Three separate databases for metrics, logs, and traces means three storage layers, three query languages, and three sets of dashboards. GreptimeDB stores all three as timestamped wide events in one columnar engine — JOIN across signals in SQL, run one stack instead of three, and ingest AI agent telemetry the same way. Read more: [Observability 2.0 and the Database for It](https://greptime.com/blogs/2025-04-25-greptimedb-observability2-new-database). +## Compatibility and Migration -Learn more in [Why GreptimeDB](https://docs.greptime.com/user-guide/concepts/why-greptimedb). +Compatibility is per protocol, and query-side coverage is narrower than ingestion. -## How GreptimeDB Compares - -| Capability | GreptimeDB | Prometheus / Thanos / Mimir | Grafana Loki | Elasticsearch | -|---|---|---|---|---| -| Data types | Metrics, logs, traces | Metrics only | Logs only | Logs, traces | -| Query language | SQL + PromQL | PromQL | LogQL | Query DSL | -| Storage | Native object storage (S3, etc.) | Local disk + object storage (Thanos/Mimir) | Object storage (chunks) | Local disk | -| Scaling | Compute-storage separation, stateless nodes | Federation / Thanos / Mimir — multi-component, ops heavy | Stateless + object storage | Shard-based, ops heavy | -| Cost efficiency | Up to 50× lower storage cost | High at scale | Moderate | High (inverted index overhead) | -| OpenTelemetry | Native (metrics + logs + traces) | Partial (metrics only) | Partial (logs only) | Via instrumentation | +| | Compatible | Not compatible | +|---|---|---| +| **Prometheus** | Remote Write ingestion; PromQL queries | Gaps are listed in [PromQL compatibility](https://docs.greptime.com/user-guide/query-data/promql/) | +| **Loki** | Push ingestion; dual-write through Grafana Alloy makes the [cutover gradual](https://greptime.com/blogs/2026-07-23-from-loki-to-greptimedb-dual-write-migration) | LogQL and the rest of the Loki query API | +| **Elasticsearch** | `_bulk` ingestion in the open-source core; [QueryDSL](https://docs.greptime.com/enterprise/elasticsearch-compatible/query/) partially, in Enterprise | Most other Elasticsearch APIs | **Benchmarks:** * [GreptimeDB tops JSONBench's billion-record cold run test](https://greptime.com/blogs/2025-03-18-jsonbench-greptimedb-performance) * [TSBS Benchmark](https://github.com/GreptimeTeam/greptimedb/tree/main/docs/benchmarks/tsbs) * [More benchmark reports](https://docs.greptime.com/user-guide/concepts/features-that-you-concern#how-is-greptimedbs-performance-compared-to-other-solutions) +## Limitations and Edition Boundary + +Cluster deployment, object storage, the Flow engine, and every ingestion protocol listed above are in the Apache-2.0 build. Repartitioning, region migration, and index creation are manual operations there. + +Read replicas, workload isolation, and automated repartitioning are **GreptimeDB Enterprise** features, along with enterprise security and governance. The [Enterprise overview](https://docs.greptime.com/enterprise/overview/) has the current list, and [pricing](https://www.greptime.com/pricing#differences) has the edition comparison. + ## Architecture GreptimeDB can run in two modes: @@ -131,7 +141,7 @@ GreptimeDB can run in two modes: * **Distributed** — four components, each independently scalable: - **Frontend** — protocol entry (OTel, Prometheus, MySQL/PostgreSQL, gRPC, ingestion APIs for Elasticsearch/InfluxDB/Loki) and the distributed query engine. Stateless, scales horizontally. - **Datanode** — region engine with WAL, memtable, SST, cache, compaction, and indexes. Persists data to object storage. Elastic. - - **Metasrv** — metadata, routing, repartitioning, autopilot, and security. Backed by a pluggable KV layer (etcd or RDS). + - **Metasrv** — metadata, routing, repartitioning, and security. Backed by a pluggable KV layer (etcd or RDS). - **Flownode** (optional) — continuous flow computation (streaming and materialized views). For deeper coverage, see the [architecture doc](https://docs.greptime.com/contributor-guide/overview/#architecture) or [DeepWiki](https://deepwiki.com/GreptimeTeam/greptimedb/1-overview). @@ -209,7 +219,7 @@ See the [Contribution Guidelines](CONTRIBUTING.md) for the full developer workfl ## Project Status -GreptimeDB is generally available, with stable APIs and regular releases. It runs in production at scale — [OceanBase Cloud](https://greptime.com/blogs/2025-07-22-user-case-obcloud-log-management-greptimedb) operates 80+ GreptimeDB clusters managing 300 TB of logs, cutting log storage cost by 60% after migrating from Grafana Loki. See more in [case studies](https://greptime.com/blogs/?category=Use%20Case). +GreptimeDB is generally available, with stable APIs and regular releases. It runs in production at scale — [OceanBase Cloud](https://greptime.com/blogs/2025-07-22-user-case-obcloud-log-management-greptimedb) operates 80+ GreptimeDB clusters managing 300 TB of logs, cutting log storage cost by 60%+ after migrating from Grafana Loki. See more in [case studies](https://greptime.com/blogs/?category=Use%20Case). Read the [v1.0 highlights](https://greptime.com/blogs/2025-11-05-greptimedb-v1-highlights) and [2026 roadmap](https://greptime.com/blogs/2026-02-11-greptimedb-roadmap-2026), or browse the [version reference](https://docs.greptime.com/nightly/reference/about-greptimedb-version). @@ -241,8 +251,9 @@ that license carry an explicit Enterprise License header. ## Commercial Support -Running GreptimeDB in your organization? -We offer enterprise add-ons, services, training, and consulting. +Scaling observability on your infrastructure? +[GreptimeDB Enterprise](https://docs.greptime.com/enterprise/overview/) adds the operational, +security, and support layer for production deployments. [Contact us](https://greptime.com/contactus) for details. ## Contributing From cb30837cd56c8e6abbdbfba098728de3ea61990d Mon Sep 17 00:00:00 2001 From: Yingwen Date: Fri, 14 Aug 2026 12:06:27 +0000 Subject: [PATCH 54/92] feat: add incremental primary key index writer (#8788) * feat: initial implementation of the pk index writer Signed-off-by: evenyag * perf: optimize primary key index writer Signed-off-by: evenyag * chore: add todo Signed-off-by: evenyag * feat(mito2): track series count in pk index metrics Signed-off-by: evenyag * refactor(mito2): simplify pk index writer cleanup Signed-off-by: evenyag * fix(mito2): clean up aborted pk index writers Signed-off-by: evenyag --------- Signed-off-by: evenyag --- src/mito2/AGENTS.md | 1 + src/mito2/src/lib.rs | 1 + src/mito2/src/pk_index.rs | 19 + src/mito2/src/pk_index/writer.rs | 1172 ++++++++++++++++++++++++++++++ 4 files changed, 1193 insertions(+) create mode 100644 src/mito2/src/pk_index.rs create mode 100644 src/mito2/src/pk_index/writer.rs diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index 73fc517fa7e..af60d784f1c 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -26,6 +26,7 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`. | `compaction` | `src/mito2/src/compaction/` | Compaction scheduler (`scheduler.rs` + `scheduler/`), TWCS picker, strict-window manual picker, compactor, memory control | | `access_layer` | `src/mito2/src/access_layer.rs` | SST read/write over the object store | | `sst` | `src/mito2/src/sst/` | Parquet format, file metadata, index layout | +| `pk_index` | `src/mito2/src/pk_index/` | Incremental writer for aggregate primary-key index files | | `read` | `src/mito2/src/read/` | `ScanRegion`, merge, dedup, projection, streaming | | `manifest` | `src/mito2/src/manifest/` | `RegionManifestManager`, manifest actions/edits | | `cache` | `src/mito2/src/cache.rs` | Write/file/page caches | diff --git a/src/mito2/src/lib.rs b/src/mito2/src/lib.rs index a212cf69533..9dfa478060b 100644 --- a/src/mito2/src/lib.rs +++ b/src/mito2/src/lib.rs @@ -38,6 +38,7 @@ pub mod gc; pub mod manifest; pub mod memtable; mod metrics; +pub mod pk_index; pub mod read; pub mod region; mod region_write_ctx; diff --git a/src/mito2/src/pk_index.rs b/src/mito2/src/pk_index.rs new file mode 100644 index 00000000000..61318026d6c --- /dev/null +++ b/src/mito2/src/pk_index.rs @@ -0,0 +1,19 @@ +// Copyright 2023 Greptime Team +// +// 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. + +//! Primary-key index writer. + +mod writer; + +pub use writer::{PkIndexWriter, PkIndexWriterMetrics, PkIndexWriterOptions, pk_columns_schema}; diff --git a/src/mito2/src/pk_index/writer.rs b/src/mito2/src/pk_index/writer.rs new file mode 100644 index 00000000000..f2407d97423 --- /dev/null +++ b/src/mito2/src/pk_index/writer.rs @@ -0,0 +1,1172 @@ +// Copyright 2023 Greptime Team +// +// 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. + +use std::cmp::Ordering; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use bytes::Bytes; +use datatypes::arrow::array::{ + Array, ArrayRef, BinaryArray, DictionaryArray, Int64Array, StringArray, UInt32Array, + UInt64Array, +}; +use datatypes::arrow::datatypes::{DataType, Field, Schema, SchemaRef, UInt32Type}; +use datatypes::arrow::record_batch::RecordBatch; +use datatypes::prelude::ConcreteDataType; +use datatypes::timestamp::timestamp_array_to_primitive; +use datatypes::value::Value; +use futures::future::BoxFuture; +use mito_codec::row_converter::{CompositeValues, PrimaryKeyCodec, build_primary_key_codec}; +use object_store::{ObjectStore, Writer}; +use parquet::arrow::AsyncArrowWriter; +use parquet::arrow::async_writer::AsyncFileWriter; +use parquet::basic::{Compression, Encoding, ZstdLevel}; +use parquet::errors::ParquetError; +use parquet::file::properties::WriterProperties; +use snafu::{OptionExt, ResultExt, ensure}; +use store_api::codec::PrimaryKeyEncoding; +use store_api::metadata::RegionMetadataRef; +use store_api::storage::ColumnId; +use store_api::storage::consts::ReservedColumnId; + +use crate::access_layer::TempFileCleaner; +use crate::error::{ + DecodeSnafu, InvalidMetaSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, OpenDalSnafu, + Result, UnexpectedSnafu, WriteParquetSnafu, +}; +use crate::sst::parquet::DEFAULT_ROW_GROUP_SIZE; +use crate::sst::parquet::flat_format::{primary_key_column_index, time_index_column_index}; +use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY}; + +const MIN_TS_COLUMN: &str = "__pk_min_ts"; +const MAX_TS_COLUMN: &str = "__pk_max_ts"; +const ROW_COUNT_COLUMN: &str = "__pk_row_count"; +const TABLE_ID_COLUMN: &str = "__table_id"; +const TSID_COLUMN: &str = "__tsid"; +const WRITE_BATCH_SIZE: usize = 1024; + +type ParquetWriter = AsyncArrowWriter; + +/// Bridges an OpenDAL [`Writer`] with Parquet's [`AsyncFileWriter`] and tracks +/// the number of bytes successfully submitted to the object store. +struct AsyncWriter { + inner: Writer, + output_bytes: u64, +} + +impl AsyncWriter { + fn new(inner: Writer) -> Self { + Self { + inner, + output_bytes: 0, + } + } + + fn output_bytes(&self) -> u64 { + self.output_bytes + } + + fn into_inner(self) -> Writer { + self.inner + } +} + +impl AsyncFileWriter for AsyncWriter { + fn write(&mut self, bytes: Bytes) -> BoxFuture<'_, parquet::errors::Result<()>> { + Box::pin(async move { + let len = bytes.len() as u64; + self.inner + .write(bytes) + .await + .map_err(|error| ParquetError::External(Box::new(error)))?; + self.output_bytes += len; + Ok(()) + }) + } + + fn complete(&mut self) -> BoxFuture<'_, parquet::errors::Result<()>> { + Box::pin(async move { + self.inner + .close() + .await + .map(|_| ()) + .map_err(|error| ParquetError::External(Box::new(error))) + }) + } +} + +/// Options for writing a primary-key index. +#[derive(Debug, Clone)] +pub struct PkIndexWriterOptions { + /// Maximum number of rows in a Parquet row group. + pub row_group_size: usize, +} + +impl Default for PkIndexWriterOptions { + fn default() -> Self { + Self { + row_group_size: DEFAULT_ROW_GROUP_SIZE, + } + } +} + +/// Metrics collected by a [`PkIndexWriter`]. +#[derive(Debug, Clone, Default)] +pub struct PkIndexWriterMetrics { + /// Number of input record batches passed to the writer. + pub input_batches: usize, + /// Number of logical rows passed to the writer. + pub input_rows: usize, + /// Number of unique time series (primary keys) aggregated by the writer. + pub num_series: usize, + /// Size of the committed index file. This remains zero for an aborted writer. + pub output_bytes: u64, + /// Time spent opening the object-store and Parquet writers. + pub open_elapsed: Duration, + /// Time spent validating, aggregating, and decoding primary keys. + pub aggregate_elapsed: Duration, + /// Time spent encoding and writing Parquet batches. + pub write_elapsed: Duration, + /// Time spent closing a completed file. + pub finish_elapsed: Duration, + /// Time spent removing an incomplete output. + pub cleanup_elapsed: Duration, + /// Whether this writer was explicitly aborted. + pub aborted: bool, +} + +impl PkIndexWriterMetrics { + /// Returns time spent by writer-owned work. + pub fn total_elapsed(&self) -> Duration { + self.open_elapsed + + self.aggregate_elapsed + + self.write_elapsed + + self.finish_elapsed + + self.cleanup_elapsed + } +} + +#[derive(Debug)] +struct PkIndexRow { + min_ts: i64, + max_ts: i64, + row_count: u64, + table_id: u32, + tsid: u64, + tags: Vec>, +} + +/// Incrementally aggregates sorted flat record batches and writes `pk_columns` Parquet. +pub struct PkIndexWriter { + codec: Arc, + tag_columns: Vec<(ColumnId, String)>, + schema: SchemaRef, + object_store: ObjectStore, + file_name: String, + writer: Option, + current_primary_key: Option>, + current_row: Option, + buffered_rows: Vec, + metrics: PkIndexWriterMetrics, + failed: bool, +} + +impl PkIndexWriter { + /// Creates a writer for `path` in `object_store`. + /// + /// The file name in `path` must be unique in the object store so aborting + /// the writer only removes temporary files that belong to this writer. + pub async fn try_new( + metadata: RegionMetadataRef, + object_store: ObjectStore, + path: &str, + options: PkIndexWriterOptions, + ) -> Result { + let open_start = Instant::now(); + ensure!( + options.row_group_size > 0, + InvalidMetaSnafu { + reason: "primary-key index row group size must be greater than zero", + } + ); + let schema = pk_columns_schema(&metadata)?; + let tag_columns = tag_columns(&metadata); + let file_name = path.rsplit('/').next().unwrap_or(path).to_string(); + let output = object_store + .writer_with(path) + .chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize) + .concurrent(DEFAULT_WRITE_CONCURRENCY) + .await + .context(OpenDalSnafu)?; + let properties = WriterProperties::builder() + .set_compression(Compression::ZSTD(ZstdLevel::default())) + .set_encoding(Encoding::PLAIN) + .set_max_row_group_row_count(Some(options.row_group_size)) + .set_column_index_truncate_length(None) + .set_statistics_truncate_length(None) + .build(); + let writer = + AsyncArrowWriter::try_new(AsyncWriter::new(output), schema.clone(), Some(properties)) + .context(WriteParquetSnafu)?; + let codec = build_primary_key_codec(&metadata); + + Ok(Self { + codec, + tag_columns, + schema, + object_store, + file_name, + writer: Some(writer), + current_primary_key: None, + current_row: None, + buffered_rows: Vec::with_capacity(WRITE_BATCH_SIZE), + metrics: PkIndexWriterMetrics { + open_elapsed: open_start.elapsed(), + ..Default::default() + }, + failed: false, + }) + } + + /// Returns the metrics collected so far. + pub fn metrics(&self) -> &PkIndexWriterMetrics { + &self.metrics + } + + /// Aggregates one sorted flat record batch. + /// + /// After this method returns an error, callers must call [`Self::abort`]. + pub async fn write(&mut self, batch: &RecordBatch) -> Result<()> { + ensure!( + !self.failed, + InvalidRecordBatchSnafu { + reason: "cannot write to a failed primary-key index writer", + } + ); + + self.metrics.input_batches += 1; + self.metrics.input_rows += batch.num_rows(); + let aggregate_start = Instant::now(); + let write_before = self.metrics.write_elapsed; + let result = self.write_inner(batch).await; + let write_cost = self.metrics.write_elapsed.saturating_sub(write_before); + self.metrics.aggregate_elapsed += aggregate_start.elapsed().saturating_sub(write_cost); + if result.is_err() { + self.failed = true; + } + result + } + + /// Finishes and commits the index file. + pub async fn finish(mut self) -> Result { + if self.failed { + let error = InvalidRecordBatchSnafu { + reason: "cannot finish a failed primary-key index writer", + } + .build(); + self.cleanup().await; + return Err(error); + } + + let result = self.finish_inner().await; + if result.is_err() { + self.cleanup().await; + } + result.map(|_| self.metrics) + } + + /// Aborts the writer and removes incomplete output files. + pub async fn abort(mut self) -> Result { + self.metrics.aborted = true; + self.cleanup().await; + Ok(self.metrics) + } + + async fn write_inner(&mut self, batch: &RecordBatch) -> Result<()> { + if batch.num_rows() == 0 { + return Ok(()); + } + ensure!( + batch.num_columns() >= 4, + InvalidRecordBatchSnafu { + reason: format!( + "primary-key index input has too few columns: {}", + batch.num_columns() + ), + } + ); + + let pk_idx = primary_key_column_index(batch.num_columns()); + let ts_idx = time_index_column_index(batch.num_columns()); + let primary_keys = batch.column(pk_idx); + let timestamps = timestamp_values(batch.column(ts_idx))?; + ensure!( + primary_keys.len() == batch.num_rows() && timestamps.len() == batch.num_rows(), + InvalidRecordBatchSnafu { + reason: "primary-key or timestamp array length does not match the batch", + } + ); + + if let Some(array) = primary_keys.as_any().downcast_ref::() { + ensure!( + array.null_count() == 0, + InvalidRecordBatchSnafu { + reason: "primary-key index input contains null primary keys", + } + ); + self.write_binary_primary_keys(array, timestamps.values()) + .await + } else if let Some(array) = primary_keys + .as_any() + .downcast_ref::>() + { + ensure!( + array.null_count() == 0, + InvalidRecordBatchSnafu { + reason: "primary-key index input contains null primary keys", + } + ); + self.write_dictionary_primary_keys(array, timestamps.values()) + .await + } else { + InvalidRecordBatchSnafu { + reason: format!( + "primary-key index requires Binary or Dictionary(UInt32, Binary) primary keys, got {:?}", + primary_keys.data_type() + ), + } + .fail() + } + } + + async fn write_binary_primary_keys( + &mut self, + primary_keys: &BinaryArray, + timestamps: &[i64], + ) -> Result<()> { + let mut start = 0; + while start < primary_keys.len() { + let primary_key = primary_keys.value(start); + let mut end = start + 1; + while end < primary_keys.len() && primary_keys.value(end) == primary_key { + end += 1; + } + + self.update_primary_key( + primary_key, + timestamps[start], + timestamps[end - 1], + (end - start) as u64, + ) + .await?; + start = end; + } + Ok(()) + } + + async fn write_dictionary_primary_keys( + &mut self, + primary_keys: &DictionaryArray, + timestamps: &[i64], + ) -> Result<()> { + let values = primary_keys + .values() + .as_any() + .downcast_ref::() + .context(InvalidRecordBatchSnafu { + reason: "primary-key dictionary values are not binary", + })?; + let keys = primary_keys.keys().values(); + let mut start = 0; + while start < keys.len() { + let key = keys[start]; + let mut end = start + 1; + while end < keys.len() && keys[end] == key { + end += 1; + } + + self.update_primary_key( + values.value(key as usize), + timestamps[start], + timestamps[end - 1], + (end - start) as u64, + ) + .await?; + start = end; + } + Ok(()) + } + + async fn update_primary_key( + &mut self, + primary_key: &[u8], + min_ts: i64, + max_ts: i64, + row_count: u64, + ) -> Result<()> { + if let Some(current) = self.current_primary_key.as_deref() { + match primary_key.cmp(current) { + Ordering::Less => { + return InvalidRecordBatchSnafu { + reason: "primary-key index input is not sorted by primary key", + } + .fail(); + } + Ordering::Equal => { + let row = self.current_row.as_mut().context(InvalidRecordBatchSnafu { + reason: "primary-key index aggregation state is incomplete", + })?; + row.min_ts = row.min_ts.min(min_ts); + row.max_ts = row.max_ts.max(max_ts); + row.row_count += row_count; + return Ok(()); + } + Ordering::Greater => self.finish_current_row().await?, + } + } + + let row = decode_primary_key( + self.codec.as_ref(), + primary_key, + min_ts, + max_ts, + row_count, + &self.tag_columns, + )?; + self.current_primary_key = Some(primary_key.to_vec()); + self.current_row = Some(row); + Ok(()) + } + + async fn finish_current_row(&mut self) -> Result<()> { + self.current_primary_key = None; + if let Some(row) = self.current_row.take() { + self.buffered_rows.push(row); + self.metrics.num_series += 1; + } + if self.buffered_rows.len() >= WRITE_BATCH_SIZE { + self.flush_rows().await?; + } + Ok(()) + } + + async fn flush_rows(&mut self) -> Result<()> { + if self.buffered_rows.is_empty() { + return Ok(()); + } + let batch = rows_to_batch(&self.schema, &self.buffered_rows)?; + let start = Instant::now(); + let result = self + .writer + .as_mut() + .context(UnexpectedSnafu { + reason: "primary-key index Parquet writer is closed", + })? + .write(&batch) + .await + .context(WriteParquetSnafu); + self.metrics.write_elapsed += start.elapsed(); + result?; + self.buffered_rows.clear(); + Ok(()) + } + + async fn finish_inner(&mut self) -> Result<()> { + let aggregate_start = Instant::now(); + let write_before = self.metrics.write_elapsed; + self.finish_current_row().await?; + self.flush_rows().await?; + let write_cost = self.metrics.write_elapsed.saturating_sub(write_before); + self.metrics.aggregate_elapsed += aggregate_start.elapsed().saturating_sub(write_cost); + + let finish_start = Instant::now(); + self.writer + .as_mut() + .context(UnexpectedSnafu { + reason: "primary-key index Parquet writer is closed", + })? + .finish() + .await + .context(WriteParquetSnafu)?; + let writer = self.writer.take().context(UnexpectedSnafu { + reason: "primary-key index Parquet writer is closed", + })?; + self.metrics.output_bytes = writer.into_inner().output_bytes(); + self.metrics.finish_elapsed += finish_start.elapsed(); + Ok(()) + } + + async fn cleanup(&mut self) { + let start = Instant::now(); + if let Some(writer) = self.writer.take() { + let mut writer = writer.into_inner().into_inner(); + if let Err(error) = writer.abort().await { + common_telemetry::warn!(error; "Failed to abort primary-key index writer"); + } + } + self.current_primary_key = None; + self.current_row = None; + self.buffered_rows.clear(); + + TempFileCleaner::clean_atomic_dir_files(&self.object_store, &[&self.file_name]).await; + self.metrics.output_bytes = 0; + self.metrics.cleanup_elapsed += start.elapsed(); + } +} + +/// Returns the Arrow schema of a `pk_columns` primary-key index. +pub fn pk_columns_schema(metadata: &RegionMetadataRef) -> Result { + validate_metadata(metadata)?; + let mut fields = vec![ + Field::new(MIN_TS_COLUMN, DataType::Int64, false), + Field::new(MAX_TS_COLUMN, DataType::Int64, false), + Field::new(ROW_COUNT_COLUMN, DataType::UInt64, false), + Field::new(TABLE_ID_COLUMN, DataType::UInt32, false), + Field::new(TSID_COLUMN, DataType::UInt64, false), + ]; + fields.extend( + tag_columns(metadata) + .into_iter() + .map(|(_, name)| Field::new(name, DataType::Utf8, true)), + ); + Ok(Arc::new(Schema::new(fields))) +} + +fn validate_metadata(metadata: &RegionMetadataRef) -> Result<()> { + for column in &metadata.column_metadatas { + ensure!( + !matches!( + column.column_schema.name.as_str(), + MIN_TS_COLUMN | MAX_TS_COLUMN | ROW_COUNT_COLUMN + ), + InvalidMetaSnafu { + reason: format!( + "primary-key index internal column name {} is already in use", + column.column_schema.name + ), + } + ); + } + ensure!( + metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse, + InvalidMetaSnafu { + reason: "primary-key index only supports sparse primary-key encoding", + } + ); + ensure!( + metadata + .primary_key + .starts_with(&[ReservedColumnId::table_id(), ReservedColumnId::tsid()]), + InvalidMetaSnafu { + reason: "primary-key index requires (__table_id, __tsid) as the primary-key prefix", + } + ); + let table_id = metadata + .column_by_id(ReservedColumnId::table_id()) + .context(InvalidMetaSnafu { + reason: "primary-key index metadata is missing __table_id", + })?; + let tsid = metadata + .column_by_id(ReservedColumnId::tsid()) + .context(InvalidMetaSnafu { + reason: "primary-key index metadata is missing __tsid", + })?; + ensure!( + table_id.column_schema.data_type == ConcreteDataType::uint32_datatype() + && tsid.column_schema.data_type == ConcreteDataType::uint64_datatype(), + InvalidMetaSnafu { + reason: "primary-key index requires UInt32 __table_id and UInt64 __tsid", + } + ); + for column in metadata.primary_key_columns() { + if is_reserved_column(column.column_id) { + continue; + } + ensure!( + column.column_schema.data_type == ConcreteDataType::string_datatype(), + InvalidMetaSnafu { + reason: format!( + "primary-key index requires string tag column {}, got {}", + column.column_schema.name, column.column_schema.data_type + ), + } + ); + } + Ok(()) +} + +fn tag_columns(metadata: &RegionMetadataRef) -> Vec<(ColumnId, String)> { + metadata + .primary_key_columns() + .filter(|column| !is_reserved_column(column.column_id)) + .map(|column| (column.column_id, column.column_schema.name.clone())) + .collect() +} + +fn is_reserved_column(column_id: ColumnId) -> bool { + column_id == ReservedColumnId::table_id() || column_id == ReservedColumnId::tsid() +} + +fn timestamp_values(array: &ArrayRef) -> Result { + let timestamps = if let Some(array) = array.as_any().downcast_ref::() { + array.clone() + } else { + timestamp_array_to_primitive(array) + .map(|(array, _)| array) + .with_context(|| InvalidRecordBatchSnafu { + reason: format!( + "primary-key index requires an Int64 or timestamp time index, got {:?}", + array.data_type() + ), + })? + }; + ensure!( + timestamps.null_count() == 0, + InvalidRecordBatchSnafu { + reason: "primary-key index input contains null timestamps", + } + ); + Ok(timestamps) +} + +// TODO(yingwen): Bench and optimize the performance if this is costly. +fn decode_primary_key( + codec: &dyn PrimaryKeyCodec, + primary_key: &[u8], + min_ts: i64, + max_ts: i64, + row_count: u64, + tag_columns: &[(ColumnId, String)], +) -> Result { + let CompositeValues::Sparse(values) = codec.decode(primary_key).context(DecodeSnafu)? else { + return InvalidRecordBatchSnafu { + reason: "decoded primary key is not sparse", + } + .fail(); + }; + let table_id = match values.get(&ReservedColumnId::table_id()) { + Some(Value::UInt32(value)) => *value, + value => { + return InvalidRecordBatchSnafu { + reason: format!("missing or invalid sparse __table_id: {value:?}"), + } + .fail(); + } + }; + let tsid = match values.get(&ReservedColumnId::tsid()) { + Some(Value::UInt64(value)) => *value, + value => { + return InvalidRecordBatchSnafu { + reason: format!("missing or invalid sparse __tsid: {value:?}"), + } + .fail(); + } + }; + let tags = tag_columns + .iter() + .map(|(column_id, _)| match values.get(column_id) { + None | Some(Value::Null) => Ok(None), + Some(Value::String(value)) => Ok(Some(value.as_utf8().to_string())), + value => InvalidRecordBatchSnafu { + reason: format!( + "invalid sparse string tag value for column {column_id}: {value:?}" + ), + } + .fail(), + }) + .collect::>>()?; + Ok(PkIndexRow { + min_ts, + max_ts, + row_count, + table_id, + tsid, + tags, + }) +} + +fn rows_to_batch(schema: &SchemaRef, rows: &[PkIndexRow]) -> Result { + let mut arrays: Vec = vec![ + Arc::new(Int64Array::from_iter_values( + rows.iter().map(|row| row.min_ts), + )), + Arc::new(Int64Array::from_iter_values( + rows.iter().map(|row| row.max_ts), + )), + Arc::new(UInt64Array::from_iter_values( + rows.iter().map(|row| row.row_count), + )), + Arc::new(UInt32Array::from_iter_values( + rows.iter().map(|row| row.table_id), + )), + Arc::new(UInt64Array::from_iter_values( + rows.iter().map(|row| row.tsid), + )), + ]; + for tag_idx in 0..schema.fields().len() - 5 { + arrays.push(Arc::new(StringArray::from_iter( + rows.iter().map(|row| row.tags[tag_idx].as_deref()), + ))); + } + RecordBatch::try_new(schema.clone(), arrays).context(NewRecordBatchSnafu) +} + +#[cfg(test)] +mod tests { + use api::v1::SemanticType; + use datatypes::arrow::array::{ + BinaryDictionaryBuilder, TimestampMicrosecondArray, TimestampMillisecondArray, + TimestampNanosecondArray, TimestampSecondArray, UInt8Array, + }; + use datatypes::arrow::datatypes::{TimeUnit, UInt32Type}; + use datatypes::schema::ColumnSchema; + use mito_codec::row_converter::{PrimaryKeyCodec, SparsePrimaryKeyCodec}; + use object_store::ErrorKind; + use object_store::services::Memory; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use store_api::codec::PrimaryKeyEncoding; + use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; + + use super::*; + use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding}; + + fn object_store() -> ObjectStore { + ObjectStore::new(Memory::default()).unwrap().finish() + } + + fn flat_schema(primary_key_type: DataType) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("__primary_key", primary_key_type, false), + Field::new("__sequence", DataType::UInt64, false), + Field::new("__op_type", DataType::UInt8, false), + ])) + } + + fn binary_batch(primary_keys: &[&[u8]], timestamps: &[i64]) -> RecordBatch { + RecordBatch::try_new( + flat_schema(DataType::Binary), + vec![ + Arc::new(TimestampMillisecondArray::from(timestamps.to_vec())), + Arc::new(BinaryArray::from_iter_values(primary_keys.iter().copied())), + Arc::new(UInt64Array::from(vec![1; timestamps.len()])), + Arc::new(UInt8Array::from(vec![0; timestamps.len()])), + ], + ) + .unwrap() + } + + fn dictionary_batch(primary_keys: &[&[u8]], timestamps: &[i64]) -> RecordBatch { + let mut builder = BinaryDictionaryBuilder::::new(); + for primary_key in primary_keys { + builder.append(*primary_key).unwrap(); + } + RecordBatch::try_new( + flat_schema(DataType::Dictionary( + Box::new(DataType::UInt32), + Box::new(DataType::Binary), + )), + vec![ + Arc::new(TimestampMillisecondArray::from(timestamps.to_vec())), + Arc::new(builder.finish()), + Arc::new(UInt64Array::from(vec![1; timestamps.len()])), + Arc::new(UInt8Array::from(vec![0; timestamps.len()])), + ], + ) + .unwrap() + } + + async fn read_index(store: &ObjectStore, path: &str) -> (u64, usize, Vec) { + let bytes = store.read(path).await.unwrap().to_bytes(); + let output_bytes = bytes.len() as u64; + let builder = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap(); + let row_groups = builder.metadata().num_row_groups(); + let batches = builder + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + (output_bytes, row_groups, batches) + } + + #[test] + fn test_pk_columns_schema() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let schema = pk_columns_schema(&metadata).unwrap(); + assert_eq!( + schema + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(), + [ + "__pk_min_ts", + "__pk_max_ts", + "__pk_row_count", + "__table_id", + "__tsid", + "tag_0", + "tag_1", + ] + ); + assert!(!schema.field(4).is_nullable()); + assert!(schema.field(5).is_nullable()); + + let dense = Arc::new(sst_region_metadata_with_encoding(PrimaryKeyEncoding::Dense)); + assert!(pk_columns_schema(&dense).is_err()); + } + + #[tokio::test] + async fn test_reject_pk_index_internal_column_names_before_opening_writer() { + for (index, name) in [MIN_TS_COLUMN, MAX_TS_COLUMN, ROW_COUNT_COLUMN] + .into_iter() + .enumerate() + { + let mut builder = RegionMetadataBuilder::from_existing( + sst_region_metadata_with_encoding(PrimaryKeyEncoding::Sparse), + ); + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new(name, ConcreteDataType::string_datatype(), true), + semantic_type: SemanticType::Field, + column_id: 100 + index as u32, + }); + let metadata = Arc::new(builder.build().unwrap()); + let store = object_store(); + let path = format!("collision-{index}.parquet"); + let error = PkIndexWriter::try_new( + metadata, + store.clone(), + &path, + PkIndexWriterOptions::default(), + ) + .await + .err() + .unwrap(); + + assert!(error.to_string().contains(name), "{error}"); + assert_eq!( + store.stat(&path).await.unwrap_err().kind(), + ErrorKind::NotFound + ); + } + } + + #[test] + fn test_timestamp_values() { + let arrays: Vec = vec![ + Arc::new(Int64Array::from(vec![1, 2])), + Arc::new(TimestampSecondArray::from(vec![1, 2])), + Arc::new(TimestampMillisecondArray::from(vec![1, 2])), + Arc::new(TimestampMicrosecondArray::from(vec![1, 2])), + Arc::new(TimestampNanosecondArray::from(vec![1, 2])), + ]; + for array in arrays { + assert_eq!(timestamp_values(&array).unwrap().values().as_ref(), &[1, 2]); + } + + let nulls: ArrayRef = Arc::new(TimestampMillisecondArray::from(vec![Some(1), None])); + let error = timestamp_values(&nulls).unwrap_err(); + assert!(error.to_string().contains("null timestamps"), "{error}"); + + let unsupported: ArrayRef = Arc::new(UInt8Array::from(vec![1, 2])); + let error = timestamp_values(&unsupported).unwrap_err(); + assert!( + error.to_string().contains("requires an Int64 or timestamp"), + "{error}" + ); + } + + #[tokio::test] + async fn test_write_batches_and_metrics() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let primary_key_1 = new_sparse_primary_key(&["a", "x"], &metadata, 1, 10); + let primary_key_2 = new_sparse_primary_key(&["b", "y"], &metadata, 1, 20); + let store = object_store(); + let mut writer = PkIndexWriter::try_new( + metadata, + store.clone(), + "pk.parquet", + PkIndexWriterOptions { row_group_size: 2 }, + ) + .await + .unwrap(); + + writer + .write(&dictionary_batch( + &[ + primary_key_1.as_slice(), + primary_key_1.as_slice(), + primary_key_1.as_slice(), + primary_key_1.as_slice(), + ], + &[70, 80, 90, 100], + )) + .await + .unwrap(); + writer + .write(&binary_batch( + &[ + primary_key_1.as_slice(), + primary_key_1.as_slice(), + primary_key_1.as_slice(), + primary_key_2.as_slice(), + primary_key_2.as_slice(), + primary_key_2.as_slice(), + primary_key_2.as_slice(), + ], + &[110, 120, 130, 200, 210, 220, 230], + )) + .await + .unwrap(); + assert_eq!(writer.metrics().input_batches, 2); + assert_eq!(writer.metrics().input_rows, 11); + + let metrics = writer.finish().await.unwrap(); + assert_eq!(metrics.input_batches, 2); + assert_eq!(metrics.input_rows, 11); + assert_eq!(metrics.num_series, 2); + assert!(!metrics.aborted); + + let (output_bytes, row_groups, batches) = read_index(&store, "pk.parquet").await; + assert_eq!(metrics.output_bytes, output_bytes); + assert_eq!(row_groups, 1); + assert_eq!(batches.len(), 1); + let batch = &batches[0]; + assert_eq!(batch.num_rows(), 2); + assert_eq!( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(), + &Int64Array::from(vec![70, 200]) + ); + assert_eq!( + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(), + &Int64Array::from(vec![130, 230]) + ); + assert_eq!( + batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(), + &UInt64Array::from(vec![7, 4]) + ); + assert_eq!( + batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap(), + &UInt32Array::from(vec![1, 1]) + ); + assert_eq!( + batch + .column(4) + .as_any() + .downcast_ref::() + .unwrap(), + &UInt64Array::from(vec![10, 20]) + ); + assert_eq!( + batch + .column(5) + .as_any() + .downcast_ref::() + .unwrap(), + &StringArray::from(vec![Some("a"), Some("b")]) + ); + } + + #[tokio::test] + async fn test_row_group_size_and_empty_file() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let store = object_store(); + let mut writer = PkIndexWriter::try_new( + metadata.clone(), + store.clone(), + "groups.parquet", + PkIndexWriterOptions { row_group_size: 2 }, + ) + .await + .unwrap(); + let keys = (0..5) + .map(|tsid| new_sparse_primary_key(&["a", "x"], &metadata, 1, tsid)) + .collect::>(); + let key_refs = keys.iter().map(Vec::as_slice).collect::>(); + writer + .write(&binary_batch(&key_refs, &[1, 2, 3, 4, 5])) + .await + .unwrap(); + writer.finish().await.unwrap(); + let (_, row_groups, _) = read_index(&store, "groups.parquet").await; + assert_eq!(row_groups, 3); + + let empty = PkIndexWriter::try_new( + metadata, + store.clone(), + "empty.parquet", + PkIndexWriterOptions::default(), + ) + .await + .unwrap() + .finish() + .await + .unwrap(); + assert_eq!(empty.input_rows, 0); + assert_eq!(empty.num_series, 0); + let (output_bytes, _, batches) = read_index(&store, "empty.parquet").await; + assert_eq!(empty.output_bytes, output_bytes); + assert!(batches.is_empty()); + } + + #[tokio::test] + async fn test_abort_and_out_of_order_input() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let primary_key_1 = new_sparse_primary_key(&["a", "x"], &metadata, 1, 10); + let primary_key_2 = new_sparse_primary_key(&["b", "y"], &metadata, 1, 20); + let store = object_store(); + let mut writer = PkIndexWriter::try_new( + metadata.clone(), + store.clone(), + "abort.parquet", + PkIndexWriterOptions { row_group_size: 1 }, + ) + .await + .unwrap(); + let error = writer + .write(&binary_batch( + &[primary_key_2.as_slice(), primary_key_1.as_slice()], + &[1, 2], + )) + .await + .unwrap_err(); + assert!(error.to_string().contains("not sorted"), "{error}"); + store + .write("abort.parquet", Bytes::from_static(b"existing")) + .await + .unwrap(); + let metrics = writer.abort().await.unwrap(); + assert!(metrics.aborted); + assert_eq!(metrics.output_bytes, 0); + assert_eq!( + store.read("abort.parquet").await.unwrap().to_bytes(), + Bytes::from_static(b"existing") + ); + + let mut writer = PkIndexWriter::try_new( + metadata, + store.clone(), + "dictionary-abort.parquet", + PkIndexWriterOptions { row_group_size: 1 }, + ) + .await + .unwrap(); + let error = writer + .write(&dictionary_batch( + &[ + primary_key_2.as_slice(), + primary_key_2.as_slice(), + primary_key_1.as_slice(), + ], + &[1, 2, 3], + )) + .await + .unwrap_err(); + assert!(error.to_string().contains("not sorted"), "{error}"); + writer.abort().await.unwrap(); + assert_eq!( + store + .stat("dictionary-abort.parquet") + .await + .unwrap_err() + .kind(), + ErrorKind::NotFound + ); + } + + #[tokio::test] + async fn test_nullable_tag_and_invalid_options() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let codec = SparsePrimaryKeyCodec::new(&metadata); + let mut primary_key = Vec::new(); + codec + .encode_value_refs( + &[ + ( + ReservedColumnId::table_id(), + datatypes::value::ValueRef::UInt32(1), + ), + ( + ReservedColumnId::tsid(), + datatypes::value::ValueRef::UInt64(10), + ), + (0, datatypes::value::ValueRef::String("a")), + ], + &mut primary_key, + ) + .unwrap(); + let store = object_store(); + let mut writer = PkIndexWriter::try_new( + metadata.clone(), + store.clone(), + "nullable.parquet", + PkIndexWriterOptions::default(), + ) + .await + .unwrap(); + writer + .write(&binary_batch(&[primary_key.as_slice()], &[1])) + .await + .unwrap(); + writer.finish().await.unwrap(); + let (_, _, batches) = read_index(&store, "nullable.parquet").await; + let tag_1 = batches[0] + .column(6) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(tag_1.is_null(0)); + + assert!( + PkIndexWriter::try_new( + metadata, + store, + "invalid.parquet", + PkIndexWriterOptions { row_group_size: 0 }, + ) + .await + .is_err() + ); + } +} From a8924bb95c43b562e4312af309d2a5e80c103185 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Fri, 14 Aug 2026 12:20:04 +0000 Subject: [PATCH 55/92] refactor(udaf): replace uddsketch implementation (#8867) * refactor(function): replace uddsketch implementation Signed-off-by: Lei, HUANG * bench(function): compare uddsketch batch ingestion Signed-off-by: Lei, HUANG * perf(function): avoid copying non-null uddsketch batches Signed-off-by: Lei, HUANG * fix(function): decode legacy uddsketch states Signed-off-by: Lei, HUANG * fix(function): harden legacy uddsketch validation Signed-off-by: Lei, HUANG * format: taplo Signed-off-by: Lei, HUANG * test: add compatibility tests for uddsketch functions Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- Cargo.lock | 18 +- Cargo.toml | 2 + src/common/function/Cargo.toml | 8 +- src/common/function/benches/uddsketch.rs | 201 ++++++ .../src/aggrs/approximate/uddsketch.rs | 299 +++++---- src/common/function/src/lib.rs | 1 + .../function/src/scalars/uddsketch_calc.rs | 99 ++- src/common/function/src/uddsketch_compat.rs | 619 ++++++++++++++++++ .../cases/uddsketch_legacy_state/case.toml | 9 + .../cases/uddsketch_legacy_state/setup.sql | 31 + .../uddsketch_legacy_state/verify.result | 63 ++ .../cases/uddsketch_legacy_state/verify.sql | 33 + 12 files changed, 1235 insertions(+), 148 deletions(-) create mode 100644 src/common/function/benches/uddsketch.rs create mode 100644 src/common/function/src/uddsketch_compat.rs create mode 100644 tests/compatibility/cases/uddsketch_legacy_state/case.toml create mode 100644 tests/compatibility/cases/uddsketch_legacy_state/setup.sql create mode 100644 tests/compatibility/cases/uddsketch_legacy_state/verify.result create mode 100644 tests/compatibility/cases/uddsketch_legacy_state/verify.sql diff --git a/Cargo.lock b/Cargo.lock index ece381b7b84..0f473d4e030 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2477,6 +2477,7 @@ dependencies = [ "common-telemetry", "common-time", "common-version", + "criterion 0.7.0", "datafusion", "datafusion-common", "datafusion-expr", @@ -2511,7 +2512,8 @@ dependencies = [ "store-api", "table", "tokio", - "uddsketch", + "uddsketch-rs", + "vu128", "wkt", ] @@ -15823,11 +15825,13 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971" [[package]] -name = "uddsketch" +name = "uddsketch-rs" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/timescaledb-toolkit.git?rev=84828fe8fb494a6a61412a3da96517fc80f7bb20#84828fe8fb494a6a61412a3da96517fc80f7bb20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f0495d3e02c0e9b074e5631de7af419b9f4fdbfec4b1ca1b3180c8b3a926cb5" dependencies = [ - "serde", + "thiserror 2.0.17", + "vu128", ] [[package]] @@ -16216,6 +16220,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "vu128" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b18da3bd753c6f4373511e5f025423986560dfe4a5e7d642cc9a0266847f9fdd" + [[package]] name = "walkdir" version = "2.5.0" diff --git a/Cargo.toml b/Cargo.toml index bfd6153bf5b..daee1bd4090 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -266,8 +266,10 @@ tracing-appender = "0.2" tracing-opentelemetry = "0.33.0" tracing-subscriber = { version = "0.3", features = ["env-filter", "json", "fmt"] } typetag = "0.2" +uddsketch-rs = "0.1.0" uuid = { version = "1.17", features = ["serde", "v4", "v7", "fast-rng"] } vrl = "0.33" +vu128 = "1.1" zstd = "0.13" # DO_NOT_REMOVE_THIS: END_OF_EXTERNAL_DEPENDENCIES diff --git a/src/common/function/Cargo.toml b/src/common/function/Cargo.toml index 622f3ddfd23..a1ea31bc32d 100644 --- a/src/common/function/Cargo.toml +++ b/src/common/function/Cargo.toml @@ -65,12 +65,18 @@ snafu.workspace = true sql.workspace = true store-api.workspace = true table.workspace = true -uddsketch = { git = "https://github.com/GreptimeTeam/timescaledb-toolkit.git", rev = "84828fe8fb494a6a61412a3da96517fc80f7bb20" } +uddsketch-rs.workspace = true +vu128.workspace = true wkt = { version = "0.11", optional = true } [dev-dependencies] approx = "0.5" +criterion.workspace = true futures.workspace = true pretty_assertions.workspace = true serde = { version = "1.0", features = ["derive"] } tokio.workspace = true + +[[bench]] +name = "uddsketch" +harness = false diff --git a/src/common/function/benches/uddsketch.rs b/src/common/function/benches/uddsketch.rs new file mode 100644 index 00000000000..6c8ad8779c2 --- /dev/null +++ b/src/common/function/benches/uddsketch.rs @@ -0,0 +1,201 @@ +// Copyright 2023 Greptime Team +// +// 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. + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Float64Array, Int64Array}; +use arrow::datatypes::{DataType, Field, Schema}; +use common_function::aggrs::approximate::uddsketch::UddSketchState; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datafusion::common::ScalarValue; +use datafusion::logical_expr::Accumulator; +use datafusion::logical_expr::function::AccumulatorArgs; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr::expressions::{Column, Literal}; + +const BATCH_SIZES: [usize; 4] = [128, 256, 1024, 2048]; +const BUCKET_SIZE: i64 = 128; +const ERROR_RATE: f64 = 0.01; + +struct AccumulatorFactory { + udf: datafusion::logical_expr::AggregateUDF, + schema: Schema, + exprs: Vec>, + expr_fields: Vec>, + return_field: Arc, +} + +impl AccumulatorFactory { + fn new() -> Self { + let udf = UddSketchState::state_udf_impl(); + let schema = Schema::new(vec![ + Field::new("bucket_size", DataType::Int64, false), + Field::new("error", DataType::Float64, false), + Field::new("value", DataType::Float64, false), + ]); + let exprs: Vec> = vec![ + Arc::new(Literal::new(ScalarValue::Int64(Some(BUCKET_SIZE)))), + Arc::new(Literal::new(ScalarValue::Float64(Some(ERROR_RATE)))), + Arc::new(Column::new("value", 2)), + ]; + let expr_fields = exprs + .iter() + .map(|expr| expr.return_field(&schema).unwrap()) + .collect::>(); + let return_type = udf + .return_type(&[DataType::Int64, DataType::Float64, DataType::Float64]) + .unwrap(); + + Self { + udf, + schema, + exprs, + expr_fields, + return_field: Arc::new(Field::new("uddsketch_state", return_type, true)), + } + } + + fn create(&self) -> Box { + self.udf + .accumulator(AccumulatorArgs { + return_field: Arc::clone(&self.return_field), + schema: &self.schema, + ignore_nulls: false, + order_bys: &[], + is_reversed: false, + name: "uddsketch_state", + is_distinct: false, + exprs: &self.exprs, + expr_fields: &self.expr_fields, + }) + .unwrap() + } +} + +fn input_arrays(batch_size: usize) -> Vec { + let bucket_sizes = Arc::new(Int64Array::from_value(BUCKET_SIZE, batch_size)) as ArrayRef; + let errors = Arc::new(Float64Array::from_value(ERROR_RATE, batch_size)) as ArrayRef; + let mut state = 0x9e37_79b9_7f4a_7c15_u64; + let values = (0..batch_size) + .map(|index| { + state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407) + .wrapping_add(index as u64); + let unit = (state >> 11) as f64 * (1.0 / (1_u64 << 53) as f64); + let magnitude = 10_f64.powf(-9.0 + 18.0 * unit); + if state & 1 == 0 { + magnitude + } else { + -magnitude + } + }) + .collect::>(); + let values = Arc::new(Float64Array::from(values)) as ArrayRef; + + vec![bucket_sizes, errors, values] +} + +fn validate(factory: &AccumulatorFactory, values: &[ArrayRef]) { + let mut accumulator = factory.create(); + accumulator.update_batch(values).unwrap(); + match accumulator.evaluate().unwrap() { + ScalarValue::Binary(Some(encoded)) => assert!(!encoded.is_empty()), + encoded => panic!("expected non-empty Binary, got {encoded:?}"), + } +} + +fn bench_uddsketch(c: &mut Criterion) { + let factory = AccumulatorFactory::new(); + let inputs = BATCH_SIZES + .into_iter() + .map(|batch_size| { + let values = input_arrays(batch_size); + validate(&factory, &values); + (batch_size, values) + }) + .collect::>(); + + let mut group = c.benchmark_group("uddsketch/ingest/fresh"); + for (batch_size, values) in &inputs { + group.throughput(Throughput::Elements(*batch_size as u64)); + group.bench_with_input( + BenchmarkId::new("batch_size", batch_size), + values, + |b, values| { + b.iter(|| { + let mut accumulator = factory.create(); + accumulator.update_batch(black_box(values)).unwrap(); + black_box(accumulator); + }); + }, + ); + } + group.finish(); + + let mut group = c.benchmark_group("uddsketch/ingest/reused"); + for (batch_size, values) in &inputs { + group.throughput(Throughput::Elements(*batch_size as u64)); + group.bench_with_input( + BenchmarkId::new("batch_size", batch_size), + values, + |b, values| { + let mut accumulator = factory.create(); + b.iter(|| { + accumulator.update_batch(black_box(values)).unwrap(); + black_box(&mut accumulator); + }); + }, + ); + } + group.finish(); + + let mut group = c.benchmark_group("uddsketch/ingest_evaluate/fresh"); + for (batch_size, values) in &inputs { + group.throughput(Throughput::Elements(*batch_size as u64)); + group.bench_with_input( + BenchmarkId::new("batch_size", batch_size), + values, + |b, values| { + b.iter(|| { + let mut accumulator = factory.create(); + accumulator.update_batch(black_box(values)).unwrap(); + black_box(accumulator.evaluate().unwrap()); + }); + }, + ); + } + group.finish(); + + let mut group = c.benchmark_group("uddsketch/ingest_evaluate/reused"); + for (batch_size, values) in &inputs { + group.throughput(Throughput::Elements(*batch_size as u64)); + group.bench_with_input( + BenchmarkId::new("batch_size", batch_size), + values, + |b, values| { + let mut accumulator = factory.create(); + b.iter(|| { + accumulator.update_batch(black_box(values)).unwrap(); + black_box(accumulator.evaluate().unwrap()); + }); + }, + ); + } + group.finish(); +} + +criterion_group!(benches, bench_uddsketch); +criterion_main!(benches); diff --git a/src/common/function/src/aggrs/approximate/uddsketch.rs b/src/common/function/src/aggrs/approximate/uddsketch.rs index 24b5c5455c0..f7d1558d13c 100644 --- a/src/common/function/src/aggrs/approximate/uddsketch.rs +++ b/src/common/function/src/aggrs/approximate/uddsketch.rs @@ -21,7 +21,6 @@ use std::sync::Arc; use common_query::prelude::*; -use common_telemetry::trace; use datafusion::common::cast::{as_binary_array, as_primitive_array}; use datafusion::common::not_impl_err; use datafusion::error::{DataFusionError, Result as DfResult}; @@ -29,27 +28,40 @@ use datafusion::logical_expr::function::AccumulatorArgs; use datafusion::logical_expr::{Accumulator as DfAccumulator, AggregateUDF, Volatility}; use datafusion::physical_plan::expressions::Literal; use datafusion::prelude::create_udaf; -use datatypes::arrow::array::ArrayRef; +use datatypes::arrow::array::{Array, ArrayRef}; use datatypes::arrow::datatypes::{DataType, Float64Type}; -use serde::{Deserialize, Serialize}; -use uddsketch::{SketchHashKey, UDDSketch}; +use uddsketch::{BatchWorkspace, UddSketch}; + +use crate::uddsketch_compat; pub const UDDSKETCH_STATE_NAME: &str = "uddsketch_state"; pub const UDDSKETCH_MERGE_NAME: &str = "uddsketch_merge"; -#[derive(Debug, Serialize, Deserialize)] +const MAX_BUCKETS: u32 = 1_000_000; + +#[derive(Debug)] pub struct UddSketchState { - uddsketch: UDDSketch, - error_rate: f64, + uddsketch: UddSketch, + workspace: BatchWorkspace, + values: Vec, } impl UddSketchState { - pub fn new(bucket_size: u64, error_rate: f64) -> Self { - Self { - uddsketch: UDDSketch::new(bucket_size, error_rate), - error_rate, + pub fn new(bucket_size: u32, error_rate: f64) -> DfResult { + if bucket_size > MAX_BUCKETS { + return Err(DataFusionError::Plan(format!( + "UDDSketch bucket size exceeds the maximum of {}", + MAX_BUCKETS + ))); } + let uddsketch = UddSketch::new(bucket_size, error_rate) + .map_err(|e| DataFusionError::Plan(e.to_string()))?; + Ok(Self { + uddsketch, + workspace: BatchWorkspace::default(), + values: Vec::new(), + }) } pub fn state_udf_impl() -> AggregateUDF { @@ -60,7 +72,7 @@ impl UddSketchState { Volatility::Immutable, Arc::new(|args| { let (bucket_size, error_rate) = downcast_accumulator_args(args)?; - Ok(Box::new(UddSketchState::new(bucket_size, error_rate))) + Ok(Box::new(UddSketchState::new(bucket_size, error_rate)?)) }), Arc::new(vec![DataType::Binary]), ) @@ -80,51 +92,44 @@ impl UddSketchState { Volatility::Immutable, Arc::new(|args| { let (bucket_size, error_rate) = downcast_accumulator_args(args)?; - Ok(Box::new(UddSketchState::new(bucket_size, error_rate))) + Ok(Box::new(UddSketchState::new(bucket_size, error_rate)?)) }), Arc::new(vec![DataType::Binary]), ) } - fn update(&mut self, value: f64) { - self.uddsketch.add_value(value); - } - fn merge(&mut self, raw: &[u8]) -> DfResult<()> { - if let Ok(uddsketch) = bincode::deserialize::(raw) { - if uddsketch.uddsketch.count() != 0 { - if self.uddsketch.max_allowed_buckets() != uddsketch.uddsketch.max_allowed_buckets() - || (self.error_rate - uddsketch.error_rate).abs() >= 1e-9 - { - return Err(DataFusionError::Plan(format!( - "Merging UDDSketch with different parameters: arguments={:?} vs actual input={:?}", - (self.uddsketch.max_allowed_buckets(), self.error_rate), - ( - uddsketch.uddsketch.max_allowed_buckets(), - uddsketch.error_rate - ) - ))); - } - self.uddsketch.merge_sketch(&uddsketch.uddsketch); - } - } else { - trace!("Warning: Failed to deserialize UDDSketch from {:?}", raw); - return Err(DataFusionError::Plan( - "Failed to deserialize UDDSketch from binary".to_string(), - )); + let uddsketch = uddsketch_compat::decode(raw).map_err(|e| { + common_telemetry::trace!("Failed to deserialize UDDSketch: {}", e); + DataFusionError::Plan("Failed to deserialize UDDSketch from binary".to_string()) + })?; + if uddsketch.count() == 0 { + return Ok(()); } - - Ok(()) + if self.uddsketch.max_buckets() != uddsketch.max_buckets() + || self.uddsketch.initial_error().to_bits() != uddsketch.initial_error().to_bits() + { + return Err(DataFusionError::Plan(format!( + "Merging UDDSketch with different parameters: arguments={:?} vs actual input={:?}", + (self.uddsketch.max_buckets(), self.uddsketch.initial_error()), + (uddsketch.max_buckets(), uddsketch.initial_error()) + ))); + } + self.uddsketch + .merge(&uddsketch) + .map_err(|e| DataFusionError::Plan(e.to_string())) } } -fn downcast_accumulator_args(args: AccumulatorArgs) -> DfResult<(u64, f64)> { +fn downcast_accumulator_args(args: AccumulatorArgs) -> DfResult<(u32, f64)> { let bucket_size = match args.exprs[0] .as_any() .downcast_ref::() .map(|lit| lit.value()) { - Some(ScalarValue::Int64(Some(value))) => *value as u64, + Some(ScalarValue::Int64(Some(value))) => u32::try_from(*value).map_err(|_| { + DataFusionError::Plan(format!("Invalid UDDSketch bucket size: {}", value)) + })?, _ => { return not_impl_err!( "{} not supported for bucket size: {}", @@ -158,9 +163,16 @@ impl DfAccumulator for UddSketchState { match array.data_type() { DataType::Float64 => { let f64_array = as_primitive_array::(array)?; - for v in f64_array.iter().flatten() { - self.update(v); - } + let values: &[f64] = if f64_array.null_count() == 0 { + f64_array.values().as_ref() + } else { + self.values.clear(); + self.values.extend(f64_array.iter().flatten()); + self.values.as_slice() + }; + self.uddsketch + .add_batch_with_workspace(values, &mut self.workspace) + .map_err(|e| DataFusionError::Execution(e.to_string()))?; } // meaning instantiate as `uddsketch_merge` DataType::Binary => self.merge_batch(std::slice::from_ref(array))?, @@ -176,35 +188,21 @@ impl DfAccumulator for UddSketchState { } fn evaluate(&mut self) -> DfResult { - Ok(ScalarValue::Binary(Some( - bincode::serialize(&self).map_err(|e| { - DataFusionError::Internal(format!("Failed to serialize UDDSketch: {}", e)) - })?, - ))) + Ok(ScalarValue::Binary(Some(self.uddsketch.encode().map_err( + |e| DataFusionError::Internal(format!("Failed to serialize UDDSketch: {}", e)), + )?))) } fn size(&self) -> usize { - // Base size of UDDSketch struct fields - let mut total_size = std::mem::size_of::() * 3 + // alpha, gamma, values_sum - std::mem::size_of::() + // compactions - std::mem::size_of::() * 2; // max_buckets, num_values - - // Size of buckets (SketchHashMap) - // Each bucket entry contains: - // - SketchHashKey (enum with i64/Zero/Invalid variants) - // - SketchHashEntry (count: u64, next: SketchHashKey) - let bucket_entry_size = std::mem::size_of::() + // key - std::mem::size_of::() + // count - std::mem::size_of::(); // next - - total_size += self.uddsketch.current_buckets_count() * bucket_entry_size; - - total_size + std::mem::size_of::() - std::mem::size_of::() + + self.uddsketch.allocated_size() + + self.workspace.allocated_size() + + self.values.capacity() * std::mem::size_of::() } fn state(&mut self) -> DfResult> { Ok(vec![ScalarValue::Binary(Some( - bincode::serialize(&self).map_err(|e| { + self.uddsketch.encode().map_err(|e| { DataFusionError::Internal(format!("Failed to serialize UDDSketch: {}", e)) })?, ))]) @@ -224,20 +222,21 @@ impl DfAccumulator for UddSketchState { #[cfg(test)] mod tests { use datafusion::arrow::array::{BinaryArray, Float64Array}; + use uddsketch::UddSketchRef; use super::*; #[test] fn test_uddsketch_state_basic() { - let mut state = UddSketchState::new(10, 0.01); - state.update(1.0); - state.update(2.0); - state.update(3.0); + let mut state = UddSketchState::new(10, 0.01).unwrap(); + state.uddsketch.add(1.0).unwrap(); + state.uddsketch.add(2.0).unwrap(); + state.uddsketch.add(3.0).unwrap(); let result = state.evaluate().unwrap(); if let ScalarValue::Binary(Some(bytes)) = result { - let deserialized: UddSketchState = bincode::deserialize(&bytes).unwrap(); - assert_eq!(deserialized.uddsketch.count(), 3); + let encoded = UddSketchRef::parse(&bytes).unwrap(); + assert_eq!(encoded.count(), 3); } else { panic!("Expected binary scalar value"); } @@ -245,39 +244,38 @@ mod tests { #[test] fn test_uddsketch_state_roundtrip() { - let mut state = UddSketchState::new(10, 0.01); - state.update(1.0); - state.update(2.0); + let mut state = UddSketchState::new(10, 0.01).unwrap(); + state.uddsketch.add(1.0).unwrap(); + state.uddsketch.add(2.0).unwrap(); // Serialize let serialized = state.evaluate().unwrap(); // Create new state and merge the serialized data - let mut new_state = UddSketchState::new(10, 0.01); + let mut new_state = UddSketchState::new(10, 0.01).unwrap(); if let ScalarValue::Binary(Some(bytes)) = &serialized { new_state.merge(bytes).unwrap(); - // Verify the merged state matches original by comparing deserialized values - let original_sketch: UddSketchState = bincode::deserialize(bytes).unwrap(); - let original_sketch = original_sketch.uddsketch; + let original_sketch = UddSketchRef::parse(bytes).unwrap(); let new_result = new_state.evaluate().unwrap(); if let ScalarValue::Binary(Some(new_bytes)) = new_result { - let new_sketch: UddSketchState = bincode::deserialize(&new_bytes).unwrap(); - let new_sketch = new_sketch.uddsketch; + let new_sketch = UddSketchRef::parse(&new_bytes).unwrap(); assert_eq!(original_sketch.count(), new_sketch.count()); assert_eq!(original_sketch.sum(), new_sketch.sum()); - assert_eq!(original_sketch.mean(), new_sketch.mean()); - assert_eq!(original_sketch.max_error(), new_sketch.max_error()); + assert_eq!( + original_sketch.max_error().unwrap(), + new_sketch.max_error().unwrap() + ); // Compare a few quantiles to ensure statistical equivalence for q in [0.1, 0.5, 0.9].iter() { + let original = original_sketch.quantile(*q).unwrap().unwrap(); + let merged = new_sketch.quantile(*q).unwrap().unwrap(); assert!( - (original_sketch.estimate_quantile(*q) - new_sketch.estimate_quantile(*q)) - .abs() - < 1e-10, + (original - merged).abs() < 1e-10, "Quantile {} mismatch: original={}, new={}", q, - original_sketch.estimate_quantile(*q), - new_sketch.estimate_quantile(*q) + original, + merged ); } } else { @@ -288,10 +286,41 @@ mod tests { } } + #[test] + fn test_uddsketch_state_merges_legacy_state() { + let mut state = UddSketchState::new(128, 0.01).unwrap(); + + state.merge(uddsketch_compat::LEGACY_STATE).unwrap(); + + let ScalarValue::Binary(Some(encoded)) = state.evaluate().unwrap() else { + panic!("Expected binary scalar value"); + }; + let sketch = UddSketchRef::parse(&encoded).unwrap(); + assert_eq!(sketch.count(), 4); + assert_eq!(sketch.sum(), 1.0); + assert_eq!(sketch.quantile(0.5).unwrap(), Some(0.9900000000000001)); + } + + #[test] + fn test_uddsketch_state_merges_compacted_legacy_state() { + let mut legacy_state = uddsketch_compat::COMPACTED_LEGACY_SKETCH.to_vec(); + legacy_state.extend_from_slice(&0.01_f64.to_le_bytes()); + let mut state = UddSketchState::new(7, 0.01).unwrap(); + + state.merge(&legacy_state).unwrap(); + + let ScalarValue::Binary(Some(encoded)) = state.evaluate().unwrap() else { + panic!("Expected binary scalar value"); + }; + let sketch = UddSketchRef::parse(&encoded).unwrap(); + assert_eq!(sketch.count(), 201); + assert_eq!(sketch.times_compacted(), 12); + } + #[test] fn test_uddsketch_state_batch_update() { - let mut state = UddSketchState::new(10, 0.01); - let values = vec![1.0f64, 2.0, 3.0]; + let mut state = UddSketchState::new(10, 0.01).unwrap(); + let values = vec![Some(1.0f64), None, Some(2.0), Some(3.0)]; let array = Arc::new(Float64Array::from(values)) as ArrayRef; state @@ -300,25 +329,52 @@ mod tests { let result = state.evaluate().unwrap(); if let ScalarValue::Binary(Some(bytes)) = result { - let deserialized: UddSketchState = bincode::deserialize(&bytes).unwrap(); - let deserialized = deserialized.uddsketch; - assert_eq!(deserialized.count(), 3); + let encoded = UddSketchRef::parse(&bytes).unwrap(); + assert_eq!(encoded.count(), 3); } else { panic!("Expected binary scalar value"); } } + #[test] + fn test_uddsketch_state_non_null_batch_avoids_values_buffer() { + let mut state = UddSketchState::new(10, 0.01).unwrap(); + let array = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])) as ArrayRef; + + state + .update_batch(&[array.clone(), array.clone(), array]) + .unwrap(); + + assert_eq!(state.uddsketch.count(), 3); + assert_eq!(state.values.capacity(), 0); + } + + #[test] + fn test_uddsketch_state_non_null_sliced_batch() { + let mut state = UddSketchState::new(10, 0.01).unwrap(); + let array = Float64Array::from(vec![1.0, 2.0, 3.0, 4.0]); + let array = array.slice(1, 2); + let array = Arc::new(array) as ArrayRef; + + state + .update_batch(&[array.clone(), array.clone(), array]) + .unwrap(); + + assert_eq!(state.uddsketch.count(), 2); + assert_eq!(state.uddsketch.sum(), 5.0); + } + #[test] fn test_uddsketch_state_merge_batch() { - let mut state1 = UddSketchState::new(10, 0.01); - state1.update(1.0); + let mut state1 = UddSketchState::new(10, 0.01).unwrap(); + state1.uddsketch.add(1.0).unwrap(); let state1_binary = state1.evaluate().unwrap(); - let mut state2 = UddSketchState::new(10, 0.01); - state2.update(2.0); + let mut state2 = UddSketchState::new(10, 0.01).unwrap(); + state2.uddsketch.add(2.0).unwrap(); let state2_binary = state2.evaluate().unwrap(); - let mut merged_state = UddSketchState::new(10, 0.01); + let mut merged_state = UddSketchState::new(10, 0.01).unwrap(); if let (ScalarValue::Binary(Some(bytes1)), ScalarValue::Binary(Some(bytes2))) = (&state1_binary, &state2_binary) { @@ -330,9 +386,8 @@ mod tests { let result = merged_state.evaluate().unwrap(); if let ScalarValue::Binary(Some(bytes)) = result { - let deserialized: UddSketchState = bincode::deserialize(&bytes).unwrap(); - let deserialized = deserialized.uddsketch; - assert_eq!(deserialized.count(), 2); + let encoded = UddSketchRef::parse(&bytes).unwrap(); + assert_eq!(encoded.count(), 2); } else { panic!("Expected binary scalar value"); } @@ -343,13 +398,14 @@ mod tests { #[test] fn test_uddsketch_state_size() { - let mut state = UddSketchState::new(10, 0.01); + let mut state = UddSketchState::new(10, 0.01).unwrap(); let initial_size = state.size(); // Add some values to create buckets - state.update(1.0); - state.update(2.0); - state.update(3.0); + let array = Arc::new(Float64Array::from_iter_values((0..64).map(f64::from))) as ArrayRef; + state + .update_batch(&[array.clone(), array.clone(), array]) + .unwrap(); let size_with_values = state.size(); assert!( @@ -358,14 +414,29 @@ mod tests { initial_size, size_with_values ); + } - // Verify size increases with more buckets - state.update(10.0); // This should create a new bucket - assert!( - state.size() > size_with_values, - "Size should increase after adding new bucket: prev={}, new={}", - size_with_values, - state.size() - ); + #[test] + fn test_uddsketch_state_rejects_invalid_config() { + assert!(UddSketchState::new(6, 0.01).is_err()); + assert!(UddSketchState::new(10, 1.0).is_err()); + + let mut maximum = UddSketchState::new(1_000_000, 0.01).unwrap(); + let ScalarValue::Binary(Some(encoded)) = maximum.evaluate().unwrap() else { + panic!("Expected binary scalar value"); + }; + UddSketchRef::parse(&encoded).unwrap(); + assert!(UddSketchState::new(1_000_001, 0.01).is_err()); + } + + #[test] + fn test_uddsketch_state_rejects_nan_batch() { + let mut state = UddSketchState::new(10, 0.01).unwrap(); + let array = Arc::new(Float64Array::from(vec![1.0, f64::NAN])) as ArrayRef; + + let error = state + .update_batch(&[array.clone(), array.clone(), array]) + .unwrap_err(); + assert!(error.to_string().contains("NaN values are not supported")); } } diff --git a/src/common/function/src/lib.rs b/src/common/function/src/lib.rs index 7abd5953671..336f8d489ca 100644 --- a/src/common/function/src/lib.rs +++ b/src/common/function/src/lib.rs @@ -18,6 +18,7 @@ mod admin; mod flush_flow; mod macros; mod system; +mod uddsketch_compat; pub mod aggrs; pub mod function; diff --git a/src/common/function/src/scalars/uddsketch_calc.rs b/src/common/function/src/scalars/uddsketch_calc.rs index d5d91e7edb6..d28bb78d1be 100644 --- a/src/common/function/src/scalars/uddsketch_calc.rs +++ b/src/common/function/src/scalars/uddsketch_calc.rs @@ -22,10 +22,10 @@ use datafusion_common::DataFusionError; use datafusion_common::arrow::array::{Array, AsArray, Float64Builder}; use datafusion_common::arrow::datatypes::{DataType, Float64Type}; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility}; -use uddsketch::UDDSketch; use crate::function::{Function, extract_args}; use crate::function_registry::FunctionRegistry; +use crate::uddsketch_compat; const NAME: &str = "uddsketch_calc"; @@ -114,26 +114,19 @@ impl Function for UddSketchCalcFunction { let sketch_bytes = sketch_opt.unwrap(); let perc = perc_opt.unwrap(); - // Deserialize the UDDSketch from its bincode representation - let sketch: UDDSketch = match bincode::deserialize(sketch_bytes) { - Ok(s) => s, + let value = match uddsketch_compat::quantile(sketch_bytes, perc) { + Ok(value) => value, Err(e) => { - common_telemetry::trace!("Failed to deserialize UDDSketch: {}", e); + common_telemetry::trace!("Failed to parse UDDSketch: {}", e); builder.append_null(); continue; } }; - // Check if the sketch is empty, if so, return null - // This is important to avoid panics when calling estimate_quantile on an empty sketch - // In practice, this will happen if input is all null - if sketch.bucket_iter().count() == 0 { - builder.append_null(); - continue; + match value { + Some(value) => builder.append_value(value), + None => builder.append_null(), } - // Compute the estimated quantile from the sketch - let result = sketch.estimate_quantile(perc); - builder.append_value(result); } Ok(ColumnarValue::Array(Arc::new(builder.finish()))) @@ -146,6 +139,7 @@ mod tests { use arrow_schema::Field; use datafusion_common::arrow::array::{BinaryArray, Float64Array}; + use uddsketch::UddSketch; use super::*; @@ -159,24 +153,17 @@ mod tests { ); // Create a test sketch - let mut sketch = UDDSketch::new(128, 0.01); - sketch.add_value(10.0); - sketch.add_value(20.0); - sketch.add_value(30.0); - sketch.add_value(40.0); - sketch.add_value(50.0); - sketch.add_value(60.0); - sketch.add_value(70.0); - sketch.add_value(80.0); - sketch.add_value(90.0); - sketch.add_value(100.0); + let mut sketch = UddSketch::new(128, 0.01).unwrap(); + sketch + .add_batch(&[10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0]) + .unwrap(); // Get expected values directly from the sketch - let expected_p50 = sketch.estimate_quantile(0.5); - let expected_p90 = sketch.estimate_quantile(0.9); - let expected_p95 = sketch.estimate_quantile(0.95); + let expected_p50 = sketch.quantile(0.5).unwrap().unwrap(); + let expected_p90 = sketch.quantile(0.9).unwrap().unwrap(); + let expected_p95 = sketch.quantile(0.95).unwrap().unwrap(); - let serialized = bincode::serialize(&sketch).unwrap(); + let serialized = sketch.encode().unwrap(); let percentiles = vec![0.5, 0.9, 0.95]; let args = vec![ @@ -207,6 +194,34 @@ mod tests { assert!((result.value(2) - expected_p95).abs() < 1e-10); } + #[test] + fn test_uddsketch_calc_function_reads_legacy_state() { + let function = UddSketchCalcFunction::default(); + let args = vec![ + ColumnarValue::Array(Arc::new(Float64Array::from(vec![0.5]))), + ColumnarValue::Array(Arc::new(BinaryArray::from_iter_values(vec![ + uddsketch_compat::LEGACY_STATE, + ]))), + ]; + + let result = function + .invoke_with_args(ScalarFunctionArgs { + args, + arg_fields: vec![], + number_rows: 1, + return_field: Arc::new(Field::new("x", DataType::Float64, false)), + config_options: Arc::new(Default::default()), + }) + .unwrap(); + let ColumnarValue::Array(result) = result else { + unreachable!() + }; + let result = result.as_primitive::(); + assert_eq!(result.len(), 1); + assert!(!result.is_null(0)); + assert_eq!(result.value(0), 0.9900000000000001); + } + #[test] fn test_uddsketch_calc_function_errors() { let function = UddSketchCalcFunction::default(); @@ -249,5 +264,31 @@ mod tests { let result = result.as_primitive::(); assert_eq!(result.len(), 1); assert!(result.is_null(0)); + + let empty = UddSketch::new(128, 0.01).unwrap().encode().unwrap(); + let mut populated = UddSketch::new(128, 0.01).unwrap(); + populated.add(1.0).unwrap(); + let populated = populated.encode().unwrap(); + let args = vec![ + ColumnarValue::Array(Arc::new(Float64Array::from(vec![0.5, -0.1, f64::NAN]))), + ColumnarValue::Array(Arc::new(BinaryArray::from_iter_values(vec![ + empty, + populated.clone(), + populated, + ]))), + ]; + let result = function + .invoke_with_args(ScalarFunctionArgs { + args, + arg_fields: vec![], + number_rows: 3, + return_field: Arc::new(Field::new("x", DataType::Float64, false)), + config_options: Arc::new(Default::default()), + }) + .unwrap(); + let ColumnarValue::Array(result) = result else { + unreachable!() + }; + assert_eq!(result.as_primitive::().null_count(), 3); } } diff --git a/src/common/function/src/uddsketch_compat.rs b/src/common/function/src/uddsketch_compat.rs new file mode 100644 index 00000000000..733b4ab043a --- /dev/null +++ b/src/common/function/src/uddsketch_compat.rs @@ -0,0 +1,619 @@ +// Copyright 2023 Greptime Team +// +// 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. + +//! Compatibility decoder for UDDSketch states written before the canonical v1 format. + +use std::collections::{HashMap, HashSet}; + +use bincode::Options; +use serde::Deserialize; +use uddsketch::{UddSketch, UddSketchRef}; + +const MAX_BYTES: usize = 64 * 1024 * 1024; +const MAX_BUCKETS: usize = 1_000_000; +const HEADER_LEN: usize = 48; + +#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, PartialEq)] +enum LegacyBucketKey { + Negative(i64), + Zero, + Positive(i64), + Invalid, +} + +#[derive(Debug, Deserialize)] +struct LegacyBucket { + count: u64, + next: LegacyBucketKey, +} + +#[derive(Debug, Deserialize)] +struct LegacyBucketStore { + map: HashMap, + head: LegacyBucketKey, +} + +#[derive(Debug, Deserialize)] +struct LegacyUddSketch { + buckets: LegacyBucketStore, + alpha: f64, + gamma: f64, + compactions: u32, + max_buckets: u64, + count: u64, + sum: f64, +} + +#[derive(Debug, Deserialize)] +struct LegacyUddSketchState { + uddsketch: LegacyUddSketch, + initial_error: f64, +} + +pub(crate) fn decode(raw: &[u8]) -> Result { + match UddSketch::decode(raw) { + Ok(sketch) => Ok(sketch), + Err(current_error) => decode_legacy_state(raw).map_err(|legacy_error| { + format!( + "canonical decode failed: {current_error}; legacy decode failed: {legacy_error}" + ) + }), + } +} + +pub(crate) fn quantile(raw: &[u8], quantile: f64) -> Result, String> { + match UddSketchRef::parse(raw) { + Ok(sketch) => sketch.quantile(quantile).map_err(|error| error.to_string()), + Err(current_error) => decode_legacy_sketch(raw) + .and_then(|sketch| sketch.quantile(quantile)) + .map_err(|legacy_error| { + format!( + "canonical decode failed: {current_error}; legacy decode failed: {legacy_error}" + ) + }), + } +} + +fn validate_legacy_input(raw: &[u8]) -> Result<(), String> { + if raw.len() > MAX_BYTES { + return Err("input exceeds the legacy decode byte limit".to_string()); + } + let map_len = raw + .get(..8) + .and_then(|bytes| bytes.try_into().ok()) + .map(u64::from_le_bytes) + .ok_or_else(|| "legacy input is truncated before the bucket count".to_string())?; + if map_len > MAX_BUCKETS as u64 { + return Err("legacy populated bucket count exceeds decode limit".to_string()); + } + Ok(()) +} + +fn legacy_options() -> impl Options { + bincode::DefaultOptions::new() + .with_fixint_encoding() + .with_limit(MAX_BYTES as u64) + .reject_trailing_bytes() +} + +fn decode_legacy_state(raw: &[u8]) -> Result { + validate_legacy_input(raw)?; + let state = legacy_options() + .deserialize::(raw) + .map_err(|error| error.to_string())?; + let encoded = state.into_canonical()?; + UddSketch::decode(&encoded).map_err(|error| error.to_string()) +} + +fn decode_legacy_sketch(raw: &[u8]) -> Result { + validate_legacy_input(raw)?; + legacy_options() + .deserialize::(raw) + .map(|state| state.uddsketch) + .or_else(|state_error| { + legacy_options() + .deserialize::(raw) + .map_err(|sketch_error| { + format!( + "legacy state decode failed: {state_error}; legacy sketch decode failed: {sketch_error}" + ) + }) + }) +} + +impl LegacyUddSketchState { + fn into_canonical(self) -> Result, String> { + let sketch = self.uddsketch; + let max_buckets = u32::try_from(sketch.max_buckets) + .map_err(|_| "legacy maximum bucket count exceeds u32".to_string())?; + if !(7..=MAX_BUCKETS as u32).contains(&max_buckets) { + return Err("legacy maximum bucket count is outside supported limits".to_string()); + } + let compactions = u8::try_from(sketch.compactions) + .map_err(|_| "legacy compaction count exceeds u8".to_string())?; + if compactions > 63 { + return Err("legacy compaction count exceeds 63".to_string()); + } + + let (expected_alpha, expected_gamma) = mapping(self.initial_error, compactions)?; + if sketch.alpha.to_bits() != expected_alpha.to_bits() + || sketch.gamma.to_bits() != expected_gamma.to_bits() + { + return Err("legacy mapping metadata is inconsistent".to_string()); + } + + let buckets = sketch.buckets.ordered()?; + if buckets.len() > max_buckets as usize { + return Err("legacy populated bucket count exceeds its configured limit".to_string()); + } + let decoded_count = buckets.iter().try_fold(0_u64, |total, (_, count)| { + total + .checked_add(*count) + .ok_or_else(|| "legacy bucket count sum overflows u64".to_string()) + })?; + if decoded_count != sketch.count { + return Err("legacy bucket counts do not match the value count".to_string()); + } + if sketch.count == 0 && sketch.sum.to_bits() != 0.0_f64.to_bits() { + return Err("legacy empty sketch sum is not positive zero".to_string()); + } + + encode_canonical( + max_buckets, + self.initial_error, + compactions, + sketch.count, + sketch.sum, + &buckets, + ) + } +} + +impl LegacyUddSketch { + fn quantile(self, quantile: f64) -> Result, String> { + if !quantile.is_finite() || !(0.0..=1.0).contains(&quantile) { + return Err("invalid quantile".to_string()); + } + if self.compactions >= 64 { + return Err("legacy compaction count must be below 64".to_string()); + } + validate_current_mapping(self.alpha, self.gamma, self.compactions)?; + if self.max_buckets == 0 || self.max_buckets > MAX_BUCKETS as u64 { + return Err("legacy maximum bucket count is outside supported limits".to_string()); + } + + let count = self.count; + let alpha = self.alpha; + let gamma = self.gamma; + let buckets = self.buckets.ordered()?; + if buckets.len() > self.max_buckets as usize { + return Err("legacy populated bucket count exceeds its configured limit".to_string()); + } + for (key, _) in &buckets { + if !legacy_bucket_key_is_attainable(self.gamma, *key) { + return Err("legacy bucket index cannot represent an f64".to_string()); + } + } + let decoded_count = buckets.iter().try_fold(0_u64, |total, (_, count)| { + total + .checked_add(*count) + .ok_or_else(|| "legacy bucket count sum overflows u64".to_string()) + })?; + if decoded_count != count { + return Err("legacy bucket counts do not match the value count".to_string()); + } + if count == 0 { + return Ok(None); + } + + let target = if quantile == 1.0 { + count + } else { + ((count as f64 * quantile) as u64) + .saturating_add(1) + .min(count) + }; + let mut seen = 0_u64; + for (key, bucket_count) in buckets { + seen += bucket_count; + if seen >= target { + return Ok(Some(legacy_bucket_value(alpha, gamma, key)?)); + } + } + Err("legacy bucket counts do not cover the quantile rank".to_string()) + } +} + +fn validate_current_mapping( + mut alpha: f64, + mut gamma: f64, + compactions: u32, +) -> Result<(), String> { + if !alpha.is_finite() || !(0.0..=1.0).contains(&alpha) { + return Err("legacy current error is outside [0, 1]".to_string()); + } + if !gamma.is_finite() || gamma <= 1.0 { + return (alpha == 1.0 && gamma == f64::INFINITY && compactions >= 5) + .then_some(()) + .ok_or_else(|| "legacy gamma must be greater than one".to_string()); + } + if alpha == 1.0 { + return (compactions > 0 && 1.0 - 2.0 / (gamma + 1.0) == 1.0) + .then_some(()) + .ok_or_else(|| "legacy saturated mapping metadata is inconsistent".to_string()); + } + + for _ in 0..compactions { + alpha /= 1.0 + (1.0 - alpha * alpha).sqrt(); + gamma = gamma.sqrt(); + } + let expected_gamma = (1.0 + alpha) / (1.0 - alpha); + let relative_difference = (gamma - expected_gamma).abs() / expected_gamma; + if relative_difference > 1e-10 { + return Err("legacy mapping metadata is inconsistent".to_string()); + } + Ok(()) +} + +fn legacy_bucket_key_is_attainable(gamma: f64, key: LegacyBucketKey) -> bool { + let index = match key { + LegacyBucketKey::Zero => return true, + LegacyBucketKey::Negative(index) | LegacyBucketKey::Positive(index) => index, + LegacyBucketKey::Invalid => return false, + }; + if index == i64::MAX { + return true; + } + + let minimum = f64::from_bits(1).log(gamma).ceil() as i64; + let maximum = f64::MAX.log(gamma).ceil() as i64; + (minimum..=maximum).contains(&index) +} + +fn legacy_bucket_value(alpha: f64, gamma: f64, key: LegacyBucketKey) -> Result { + let magnitude = |index: i64| gamma.powf(index as f64 - 1.0) * (1.0 + alpha); + match key { + LegacyBucketKey::Negative(index) => Ok(-magnitude(index)), + LegacyBucketKey::Zero => Ok(0.0), + LegacyBucketKey::Positive(index) => Ok(magnitude(index)), + LegacyBucketKey::Invalid => Err("legacy bucket chain contains the end marker".to_string()), + } +} + +impl LegacyBucketStore { + fn ordered(self) -> Result, String> { + if self.map.len() > MAX_BUCKETS { + return Err("legacy populated bucket count exceeds decode limit".to_string()); + } + if self.map.is_empty() { + if self.head != LegacyBucketKey::Invalid { + return Err("legacy empty bucket store has a nonempty head".to_string()); + } + return Ok(Vec::new()); + } + + let mut buckets = Vec::with_capacity(self.map.len()); + let mut visited = HashSet::with_capacity(self.map.len()); + let mut key = self.head; + while key != LegacyBucketKey::Invalid { + if !visited.insert(key) { + return Err("legacy bucket chain contains a cycle".to_string()); + } + let bucket = self + .map + .get(&key) + .ok_or_else(|| "legacy bucket chain references a missing bucket".to_string())?; + if bucket.count == 0 { + return Err("legacy bucket has a zero count".to_string()); + } + buckets.push((key, bucket.count)); + key = bucket.next; + } + if buckets.len() != self.map.len() || self.map.contains_key(&LegacyBucketKey::Invalid) { + return Err("legacy bucket store contains unreachable buckets".to_string()); + } + if !buckets.windows(2).all(|pair| key_lt(pair[0].0, pair[1].0)) { + return Err("legacy bucket chain is not strictly ordered".to_string()); + } + Ok(buckets) + } +} + +fn mapping(initial_error: f64, compactions: u8) -> Result<(f64, f64), String> { + if !initial_error.is_finite() || !(1e-12..1.0).contains(&initial_error) { + return Err("legacy initial error is outside [1e-12, 1)".to_string()); + } + let mut alpha = initial_error; + let mut gamma = (1.0 + initial_error) / (1.0 - initial_error); + for _ in 0..compactions { + gamma *= gamma; + alpha = 2.0 * alpha / (1.0 + alpha.powi(2)); + } + Ok((alpha, gamma)) +} + +fn encode_canonical( + max_buckets: u32, + initial_error: f64, + compactions: u8, + count: u64, + sum: f64, + buckets: &[(LegacyBucketKey, u64)], +) -> Result, String> { + let negative = buckets + .iter() + .filter_map(|(key, count)| match key { + LegacyBucketKey::Negative(index) => Some((*index, *count)), + _ => None, + }) + .collect::>(); + let zero_count = buckets + .iter() + .find_map(|(key, count)| (*key == LegacyBucketKey::Zero).then_some(*count)) + .unwrap_or(0); + let positive = buckets + .iter() + .filter_map(|(key, count)| match key { + LegacyBucketKey::Positive(index) => Some((*index, *count)), + _ => None, + }) + .collect::>(); + + let mut encoded = vec![0; HEADER_LEN]; + put_varint(&mut encoded, negative.len() as u64); + put_varint(&mut encoded, zero_count); + put_varint(&mut encoded, positive.len() as u64); + encode_section(&mut encoded, &negative, true)?; + encode_section(&mut encoded, &positive, false)?; + + let payload_len = u32::try_from(encoded.len() - HEADER_LEN) + .map_err(|_| "legacy canonical payload exceeds u32".to_string())?; + encoded[0..4].copy_from_slice(b"UDDS"); + encoded[4] = 1; + encoded[6] = compactions; + encoded[8..12].copy_from_slice(&max_buckets.to_le_bytes()); + encoded[12..16].copy_from_slice(&(buckets.len() as u32).to_le_bytes()); + encoded[16..24].copy_from_slice(&initial_error.to_bits().to_le_bytes()); + encoded[24..32].copy_from_slice(&count.to_le_bytes()); + encoded[32..40].copy_from_slice(&sum.to_bits().to_le_bytes()); + encoded[40..44].copy_from_slice(&payload_len.to_le_bytes()); + Ok(encoded) +} + +fn encode_section( + output: &mut Vec, + buckets: &[(i64, u64)], + descending: bool, +) -> Result<(), String> { + let mut previous = None; + for &(index, count) in buckets { + let encoded_index = match previous { + None => zigzag(index), + Some((previous_index, _)) => { + let delta = if descending { + previous_index as i128 - index as i128 + } else { + index as i128 - previous_index as i128 + }; + u64::try_from(delta) + .ok() + .filter(|delta| *delta != 0) + .ok_or_else(|| "legacy bucket indices are not strictly ordered".to_string())? + } + }; + let encoded_count = match previous { + None => count, + Some((_, previous_count)) => zigzag(count.wrapping_sub(previous_count) as i64), + }; + put_varint(output, encoded_index); + put_varint(output, encoded_count); + previous = Some((index, count)); + } + Ok(()) +} + +fn put_varint(output: &mut Vec, value: u64) { + let mut buffer = [0; 9]; + let len = vu128::encode_u64(&mut buffer, value); + output.extend_from_slice(&buffer[..len]); +} + +const fn zigzag(value: i64) -> u64 { + (value.wrapping_shl(1) ^ (value >> 63)) as u64 +} + +fn key_lt(left: LegacyBucketKey, right: LegacyBucketKey) -> bool { + use LegacyBucketKey::*; + match (left, right) { + (Negative(left), Negative(right)) => left > right, + (Negative(_), Zero | Positive(_)) | (Zero, Positive(_)) => true, + (Positive(left), Positive(right)) => left < right, + _ => false, + } +} + +#[cfg(test)] +pub(crate) const LEGACY_STATE: &[u8] = &[ + 4, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 116, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 116, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 116, + 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 116, 0, 0, 0, 0, 0, 0, 0, + 123, 20, 174, 71, 225, 122, 132, 63, 253, 74, 129, 90, 191, 82, 240, 63, 0, 0, 0, 0, 128, 0, 0, + 0, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 240, 63, 123, 20, 174, 71, 225, 122, + 132, 63, +]; + +#[cfg(test)] +pub(crate) const COMPACTED_LEGACY_SKETCH: &[u8] = &[ + 6, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 255, 255, 255, 255, 255, 255, 255, 255, 36, 0, 0, 0, 0, 0, + 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 36, 0, 0, 0, 0, + 0, 0, 0, 2, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 254, 255, 255, 255, 255, 255, 255, + 255, 29, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 255, 255, 255, 255, 255, 255, 255, 255, 2, 0, 0, 0, + 2, 0, 0, 0, 0, 0, 0, 0, 36, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, + 0, 3, 0, 0, 0, 0, 0, 0, 0, 29, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 2, 0, 0, 0, 1, 0, 0, 0, 0, 0, + 0, 0, 35, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 254, 255, 255, + 255, 255, 255, 255, 255, 0, 0, 0, 0, 0, 0, 240, 63, 169, 137, 186, 120, 1, 63, 82, 71, 12, 0, + 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 201, 0, 0, 0, 0, 0, 0, 0, 112, 103, 108, 212, 220, 81, 180, 84, +]; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn current_format_is_decoded_without_legacy_fallback() { + let mut sketch = UddSketch::new(128, 0.01).unwrap(); + sketch.add(42.0).unwrap(); + let expected = sketch.quantile(0.5).unwrap(); + let encoded = sketch.encode().unwrap(); + + assert!(decode_legacy_state(&encoded).is_err()); + assert_eq!(decode(&encoded).unwrap(), sketch); + assert_eq!(quantile(&encoded, 0.5).unwrap(), expected); + } + + #[test] + fn legacy_decoder_rejects_oversized_bucket_count_before_deserializing() { + let mut encoded = LEGACY_STATE.to_vec(); + encoded[..8].copy_from_slice(&((MAX_BUCKETS as u64) + 1).to_le_bytes()); + + assert_eq!( + decode_legacy_state(&encoded).unwrap_err(), + "legacy populated bucket count exceeds decode limit" + ); + } + + #[test] + fn legacy_decoder_reads_bare_sketch_from_scalar_callers() { + let bare_sketch = &LEGACY_STATE[..LEGACY_STATE.len() - std::mem::size_of::()]; + + assert_eq!( + quantile(bare_sketch, 0.5).unwrap(), + Some(0.9900000000000001) + ); + } + + #[test] + fn legacy_decoder_reads_compacted_bare_sketch() { + assert!(decode_legacy_state(COMPACTED_LEGACY_SKETCH).is_err()); + assert!(quantile(COMPACTED_LEGACY_SKETCH, 0.5).unwrap().is_some()); + } + + #[test] + fn legacy_quantile_rejects_invalid_mapping() { + let bare_len = LEGACY_STATE.len() - std::mem::size_of::(); + let mut invalid = LEGACY_STATE[..bare_len].to_vec(); + let alpha_offset = bare_len - 44; + invalid[alpha_offset..alpha_offset + 8].copy_from_slice(&f64::NAN.to_le_bytes()); + + assert!(quantile(&invalid, 0.5).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + sketch.gamma = 2.0; + assert!(sketch.quantile(0.5).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + sketch.alpha = 1.0; + sketch.gamma = 2.0; + assert!(sketch.quantile(0.5).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + sketch.alpha = 1e-12; + sketch.gamma = 1.0 + 1e-9; + sketch.compactions = 0; + assert!(sketch.quantile(0.5).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + sketch.alpha = 1.0; + sketch.gamma = 1e20; + sketch.compactions = 0; + assert!(sketch.quantile(0.5).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + sketch.alpha = 1.0; + sketch.gamma = f64::INFINITY; + sketch.compactions = 1; + assert!(validate_current_mapping(sketch.alpha, sketch.gamma, sketch.compactions).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + sketch.alpha = 0.99999999; + sketch.gamma = 2.0; + sketch.compactions = 0; + assert!(sketch.quantile(0.5).is_err()); + + let mut sketch = decode_legacy_sketch(&LEGACY_STATE[..bare_len]).unwrap(); + let LegacyBucketKey::Negative(index) = sketch.buckets.head else { + panic!("Expected negative head bucket"); + }; + let bucket = sketch + .buckets + .map + .remove(&LegacyBucketKey::Negative(index)) + .unwrap(); + sketch + .buckets + .map + .insert(LegacyBucketKey::Negative(i64::MIN), bucket); + sketch.buckets.head = LegacyBucketKey::Negative(i64::MIN); + assert!(sketch.quantile(0.5).is_err()); + } + + #[test] + fn legacy_quantile_accepts_saturated_mapping() { + let sketch = LegacyUddSketch { + buckets: LegacyBucketStore { + map: HashMap::from([( + LegacyBucketKey::Positive(0), + LegacyBucket { + count: 1, + next: LegacyBucketKey::Invalid, + }, + )]), + head: LegacyBucketKey::Positive(0), + }, + alpha: 1.0, + gamma: f64::INFINITY, + compactions: 63, + max_buckets: 7, + count: 1, + sum: f64::INFINITY, + }; + + assert_eq!(sketch.quantile(0.5).unwrap(), Some(0.0)); + } + + #[test] + fn legacy_quantile_handles_maximum_count_at_one() { + let sketch = LegacyUddSketch { + buckets: LegacyBucketStore { + map: HashMap::from([( + LegacyBucketKey::Zero, + LegacyBucket { + count: u64::MAX, + next: LegacyBucketKey::Invalid, + }, + )]), + head: LegacyBucketKey::Zero, + }, + alpha: 0.01, + gamma: (1.0 + 0.01) / (1.0 - 0.01), + compactions: 0, + max_buckets: 128, + count: u64::MAX, + sum: 0.0, + }; + + assert_eq!(sketch.quantile(1.0).unwrap(), Some(0.0)); + } +} diff --git a/tests/compatibility/cases/uddsketch_legacy_state/case.toml b/tests/compatibility/cases/uddsketch_legacy_state/case.toml new file mode 100644 index 00000000000..b22670037d1 --- /dev/null +++ b/tests/compatibility/cases/uddsketch_legacy_state/case.toml @@ -0,0 +1,9 @@ +name = "uddsketch_legacy_state" +reason = "Verify UDDSketch state bytes persisted by old binaries (wrapped legacy bincode UddSketchState) are read directly by uddsketch_calc and uddsketch_merge on the new binary, merged into canonical v1, and merged together with newly generated canonical states." +introduced_by = "PR #8867" +topologies = ["distributed", "standalone"] +from_range = [">=v1.0.0"] +to_range = ["*"] +features = ["table", "query", "udf", "aggregate"] +owner = "query" +namespace = "uddsketch_legacy_state" diff --git a/tests/compatibility/cases/uddsketch_legacy_state/setup.sql b/tests/compatibility/cases/uddsketch_legacy_state/setup.sql new file mode 100644 index 00000000000..abe681a7b98 --- /dev/null +++ b/tests/compatibility/cases/uddsketch_legacy_state/setup.sql @@ -0,0 +1,31 @@ +CREATE TABLE uddsketch_values ( + seq_id INT PRIMARY KEY, + val DOUBLE, + ts TIMESTAMP TIME INDEX DEFAULT now() +); + +INSERT INTO uddsketch_values (seq_id, val) VALUES + (1, 10.0), + (2, 20.0), + (3, 30.0), + (4, 40.0), + (5, 50.0), + (6, 60.0), + (7, 70.0), + (8, 80.0), + (9, 90.0), + (10, 100.0); + +CREATE TABLE uddsketch_states ( + state BINARY, + grp INT PRIMARY KEY, + ts TIMESTAMP TIME INDEX DEFAULT now() +); + +INSERT INTO uddsketch_states (state, grp) +SELECT uddsketch_state(128, 0.01, val), seq_id / 5 * 5 AS grp +FROM uddsketch_values +GROUP BY grp; + +ADMIN FLUSH_TABLE('uddsketch_values'); +ADMIN FLUSH_TABLE('uddsketch_states'); diff --git a/tests/compatibility/cases/uddsketch_legacy_state/verify.result b/tests/compatibility/cases/uddsketch_legacy_state/verify.result new file mode 100644 index 00000000000..b4c8a3d1566 --- /dev/null +++ b/tests/compatibility/cases/uddsketch_legacy_state/verify.result @@ -0,0 +1,63 @@ +-- (1) direct uddsketch_calc of each persisted legacy row +SELECT grp, uddsketch_calc(0.5, state) AS p50 +FROM uddsketch_states +ORDER BY grp; + ++-----+--------------------+ +| grp | p50 | ++-----+--------------------+ +| 0 | 30.267171338721894 | +| 5 | 70.11183939140265 | +| 10 | 100.49456770856492 | ++-----+--------------------+ + +-- (2) uddsketch_merge over the persisted legacy rows, then calc. +-- Succeeds only if every legacy row decodes to the exact declared +-- parameters (128, 0.01); the merge parameter check is the guard. +SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 +FROM uddsketch_states; + ++--------------------+ +| p50 | ++--------------------+ +| 59.745049810145126 | ++--------------------+ + +-- (3) generate fresh canonical states with the new binary and merge mixed legacy + canonical rows +CREATE TABLE uddsketch_new_states ( + state BINARY, + grp INT PRIMARY KEY, + ts TIMESTAMP TIME INDEX DEFAULT now() +); + +Affected Rows: 0 + +INSERT INTO uddsketch_new_states (state, grp) +SELECT uddsketch_state(128, 0.01, val), seq_id / 5 * 5 AS grp +FROM uddsketch_values +GROUP BY grp; + +Affected Rows: 3 + +SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 +FROM ( + SELECT state FROM uddsketch_states + UNION ALL + SELECT state FROM uddsketch_new_states +) AS all_states; + ++--------------------+ +| p50 | ++--------------------+ +| 59.745049810145126 | ++--------------------+ + +-- (4) canonical-only control +SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 +FROM uddsketch_new_states; + ++--------------------+ +| p50 | ++--------------------+ +| 59.745049810145126 | ++--------------------+ diff --git a/tests/compatibility/cases/uddsketch_legacy_state/verify.sql b/tests/compatibility/cases/uddsketch_legacy_state/verify.sql new file mode 100644 index 00000000000..8cf84cfbbc5 --- /dev/null +++ b/tests/compatibility/cases/uddsketch_legacy_state/verify.sql @@ -0,0 +1,33 @@ +-- (1) direct uddsketch_calc of each persisted legacy row +SELECT grp, uddsketch_calc(0.5, state) AS p50 +FROM uddsketch_states +ORDER BY grp; + +-- (2) uddsketch_merge over the persisted legacy rows, then calc. +-- Succeeds only if every legacy row decodes to the exact declared +-- parameters (128, 0.01); the merge parameter check is the guard. +SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 +FROM uddsketch_states; + +-- (3) generate fresh canonical states with the new binary and merge mixed legacy + canonical rows +CREATE TABLE uddsketch_new_states ( + state BINARY, + grp INT PRIMARY KEY, + ts TIMESTAMP TIME INDEX DEFAULT now() +); + +INSERT INTO uddsketch_new_states (state, grp) +SELECT uddsketch_state(128, 0.01, val), seq_id / 5 * 5 AS grp +FROM uddsketch_values +GROUP BY grp; + +SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 +FROM ( + SELECT state FROM uddsketch_states + UNION ALL + SELECT state FROM uddsketch_new_states +) AS all_states; + +-- (4) canonical-only control +SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 +FROM uddsketch_new_states; From 4126cf99b6cd1d58d74b26e92490039198ecac6e Mon Sep 17 00:00:00 2001 From: Yingwen Date: Mon, 17 Aug 2026 07:21:29 +0000 Subject: [PATCH 56/92] refactor(mito2): rename pk index to series index (#8893) Signed-off-by: evenyag --- src/mito2/AGENTS.md | 2 +- src/mito2/src/lib.rs | 2 +- .../src/{pk_index.rs => series_index.rs} | 6 +- .../src/{pk_index => series_index}/writer.rs | 148 +++++++++--------- 4 files changed, 80 insertions(+), 78 deletions(-) rename src/mito2/src/{pk_index.rs => series_index.rs} (80%) rename src/mito2/src/{pk_index => series_index}/writer.rs (89%) diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index af60d784f1c..e68ca586096 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -26,7 +26,7 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`. | `compaction` | `src/mito2/src/compaction/` | Compaction scheduler (`scheduler.rs` + `scheduler/`), TWCS picker, strict-window manual picker, compactor, memory control | | `access_layer` | `src/mito2/src/access_layer.rs` | SST read/write over the object store | | `sst` | `src/mito2/src/sst/` | Parquet format, file metadata, index layout | -| `pk_index` | `src/mito2/src/pk_index/` | Incremental writer for aggregate primary-key index files | +| `series_index` | `src/mito2/src/series_index/` | Incremental writer for aggregate series index files | | `read` | `src/mito2/src/read/` | `ScanRegion`, merge, dedup, projection, streaming | | `manifest` | `src/mito2/src/manifest/` | `RegionManifestManager`, manifest actions/edits | | `cache` | `src/mito2/src/cache.rs` | Write/file/page caches | diff --git a/src/mito2/src/lib.rs b/src/mito2/src/lib.rs index 9dfa478060b..4f4409afd8a 100644 --- a/src/mito2/src/lib.rs +++ b/src/mito2/src/lib.rs @@ -38,13 +38,13 @@ pub mod gc; pub mod manifest; pub mod memtable; mod metrics; -pub mod pk_index; pub mod read; pub mod region; mod region_write_ctx; pub mod remap_manifest; pub mod request; pub mod schedule; +pub mod series_index; pub mod sst; mod time_provider; pub mod wal; diff --git a/src/mito2/src/pk_index.rs b/src/mito2/src/series_index.rs similarity index 80% rename from src/mito2/src/pk_index.rs rename to src/mito2/src/series_index.rs index 61318026d6c..21ced00cf64 100644 --- a/src/mito2/src/pk_index.rs +++ b/src/mito2/src/series_index.rs @@ -12,8 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Primary-key index writer. +//! Series index writer. mod writer; -pub use writer::{PkIndexWriter, PkIndexWriterMetrics, PkIndexWriterOptions, pk_columns_schema}; +pub use writer::{ + SeriesIndexWriter, SeriesIndexWriterMetrics, SeriesIndexWriterOptions, series_index_schema, +}; diff --git a/src/mito2/src/pk_index/writer.rs b/src/mito2/src/series_index/writer.rs similarity index 89% rename from src/mito2/src/pk_index/writer.rs rename to src/mito2/src/series_index/writer.rs index f2407d97423..ce5b7405f21 100644 --- a/src/mito2/src/pk_index/writer.rs +++ b/src/mito2/src/series_index/writer.rs @@ -49,9 +49,9 @@ use crate::sst::parquet::DEFAULT_ROW_GROUP_SIZE; use crate::sst::parquet::flat_format::{primary_key_column_index, time_index_column_index}; use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY}; -const MIN_TS_COLUMN: &str = "__pk_min_ts"; -const MAX_TS_COLUMN: &str = "__pk_max_ts"; -const ROW_COUNT_COLUMN: &str = "__pk_row_count"; +const MIN_TS_COLUMN: &str = "__series_min_ts"; +const MAX_TS_COLUMN: &str = "__series_max_ts"; +const ROW_COUNT_COLUMN: &str = "__series_row_count"; const TABLE_ID_COLUMN: &str = "__table_id"; const TSID_COLUMN: &str = "__tsid"; const WRITE_BATCH_SIZE: usize = 1024; @@ -106,14 +106,14 @@ impl AsyncFileWriter for AsyncWriter { } } -/// Options for writing a primary-key index. +/// Options for writing a series index. #[derive(Debug, Clone)] -pub struct PkIndexWriterOptions { +pub struct SeriesIndexWriterOptions { /// Maximum number of rows in a Parquet row group. pub row_group_size: usize, } -impl Default for PkIndexWriterOptions { +impl Default for SeriesIndexWriterOptions { fn default() -> Self { Self { row_group_size: DEFAULT_ROW_GROUP_SIZE, @@ -121,9 +121,9 @@ impl Default for PkIndexWriterOptions { } } -/// Metrics collected by a [`PkIndexWriter`]. +/// Metrics collected by a [`SeriesIndexWriter`]. #[derive(Debug, Clone, Default)] -pub struct PkIndexWriterMetrics { +pub struct SeriesIndexWriterMetrics { /// Number of input record batches passed to the writer. pub input_batches: usize, /// Number of logical rows passed to the writer. @@ -146,7 +146,7 @@ pub struct PkIndexWriterMetrics { pub aborted: bool, } -impl PkIndexWriterMetrics { +impl SeriesIndexWriterMetrics { /// Returns time spent by writer-owned work. pub fn total_elapsed(&self) -> Duration { self.open_elapsed @@ -158,7 +158,7 @@ impl PkIndexWriterMetrics { } #[derive(Debug)] -struct PkIndexRow { +struct SeriesIndexRow { min_ts: i64, max_ts: i64, row_count: u64, @@ -167,8 +167,8 @@ struct PkIndexRow { tags: Vec>, } -/// Incrementally aggregates sorted flat record batches and writes `pk_columns` Parquet. -pub struct PkIndexWriter { +/// Incrementally aggregates sorted flat record batches into series-index Parquet files. +pub struct SeriesIndexWriter { codec: Arc, tag_columns: Vec<(ColumnId, String)>, schema: SchemaRef, @@ -176,13 +176,13 @@ pub struct PkIndexWriter { file_name: String, writer: Option, current_primary_key: Option>, - current_row: Option, - buffered_rows: Vec, - metrics: PkIndexWriterMetrics, + current_row: Option, + buffered_rows: Vec, + metrics: SeriesIndexWriterMetrics, failed: bool, } -impl PkIndexWriter { +impl SeriesIndexWriter { /// Creates a writer for `path` in `object_store`. /// /// The file name in `path` must be unique in the object store so aborting @@ -191,16 +191,16 @@ impl PkIndexWriter { metadata: RegionMetadataRef, object_store: ObjectStore, path: &str, - options: PkIndexWriterOptions, + options: SeriesIndexWriterOptions, ) -> Result { let open_start = Instant::now(); ensure!( options.row_group_size > 0, InvalidMetaSnafu { - reason: "primary-key index row group size must be greater than zero", + reason: "series index row group size must be greater than zero", } ); - let schema = pk_columns_schema(&metadata)?; + let schema = series_index_schema(&metadata)?; let tag_columns = tag_columns(&metadata); let file_name = path.rsplit('/').next().unwrap_or(path).to_string(); let output = object_store @@ -231,7 +231,7 @@ impl PkIndexWriter { current_primary_key: None, current_row: None, buffered_rows: Vec::with_capacity(WRITE_BATCH_SIZE), - metrics: PkIndexWriterMetrics { + metrics: SeriesIndexWriterMetrics { open_elapsed: open_start.elapsed(), ..Default::default() }, @@ -240,7 +240,7 @@ impl PkIndexWriter { } /// Returns the metrics collected so far. - pub fn metrics(&self) -> &PkIndexWriterMetrics { + pub fn metrics(&self) -> &SeriesIndexWriterMetrics { &self.metrics } @@ -251,7 +251,7 @@ impl PkIndexWriter { ensure!( !self.failed, InvalidRecordBatchSnafu { - reason: "cannot write to a failed primary-key index writer", + reason: "cannot write to a failed series index writer", } ); @@ -269,10 +269,10 @@ impl PkIndexWriter { } /// Finishes and commits the index file. - pub async fn finish(mut self) -> Result { + pub async fn finish(mut self) -> Result { if self.failed { let error = InvalidRecordBatchSnafu { - reason: "cannot finish a failed primary-key index writer", + reason: "cannot finish a failed series index writer", } .build(); self.cleanup().await; @@ -287,7 +287,7 @@ impl PkIndexWriter { } /// Aborts the writer and removes incomplete output files. - pub async fn abort(mut self) -> Result { + pub async fn abort(mut self) -> Result { self.metrics.aborted = true; self.cleanup().await; Ok(self.metrics) @@ -301,7 +301,7 @@ impl PkIndexWriter { batch.num_columns() >= 4, InvalidRecordBatchSnafu { reason: format!( - "primary-key index input has too few columns: {}", + "series index input has too few columns: {}", batch.num_columns() ), } @@ -322,7 +322,7 @@ impl PkIndexWriter { ensure!( array.null_count() == 0, InvalidRecordBatchSnafu { - reason: "primary-key index input contains null primary keys", + reason: "series index input contains null primary keys", } ); self.write_binary_primary_keys(array, timestamps.values()) @@ -334,7 +334,7 @@ impl PkIndexWriter { ensure!( array.null_count() == 0, InvalidRecordBatchSnafu { - reason: "primary-key index input contains null primary keys", + reason: "series index input contains null primary keys", } ); self.write_dictionary_primary_keys(array, timestamps.values()) @@ -342,7 +342,7 @@ impl PkIndexWriter { } else { InvalidRecordBatchSnafu { reason: format!( - "primary-key index requires Binary or Dictionary(UInt32, Binary) primary keys, got {:?}", + "series index requires Binary or Dictionary(UInt32, Binary) primary keys, got {:?}", primary_keys.data_type() ), } @@ -419,13 +419,13 @@ impl PkIndexWriter { match primary_key.cmp(current) { Ordering::Less => { return InvalidRecordBatchSnafu { - reason: "primary-key index input is not sorted by primary key", + reason: "series index input is not sorted by primary key", } .fail(); } Ordering::Equal => { let row = self.current_row.as_mut().context(InvalidRecordBatchSnafu { - reason: "primary-key index aggregation state is incomplete", + reason: "series index aggregation state is incomplete", })?; row.min_ts = row.min_ts.min(min_ts); row.max_ts = row.max_ts.max(max_ts); @@ -471,7 +471,7 @@ impl PkIndexWriter { .writer .as_mut() .context(UnexpectedSnafu { - reason: "primary-key index Parquet writer is closed", + reason: "series index Parquet writer is closed", })? .write(&batch) .await @@ -494,13 +494,13 @@ impl PkIndexWriter { self.writer .as_mut() .context(UnexpectedSnafu { - reason: "primary-key index Parquet writer is closed", + reason: "series index Parquet writer is closed", })? .finish() .await .context(WriteParquetSnafu)?; let writer = self.writer.take().context(UnexpectedSnafu { - reason: "primary-key index Parquet writer is closed", + reason: "series index Parquet writer is closed", })?; self.metrics.output_bytes = writer.into_inner().output_bytes(); self.metrics.finish_elapsed += finish_start.elapsed(); @@ -512,7 +512,7 @@ impl PkIndexWriter { if let Some(writer) = self.writer.take() { let mut writer = writer.into_inner().into_inner(); if let Err(error) = writer.abort().await { - common_telemetry::warn!(error; "Failed to abort primary-key index writer"); + common_telemetry::warn!(error; "Failed to abort series index writer"); } } self.current_primary_key = None; @@ -525,8 +525,8 @@ impl PkIndexWriter { } } -/// Returns the Arrow schema of a `pk_columns` primary-key index. -pub fn pk_columns_schema(metadata: &RegionMetadataRef) -> Result { +/// Returns the Arrow schema of a series index. +pub fn series_index_schema(metadata: &RegionMetadataRef) -> Result { validate_metadata(metadata)?; let mut fields = vec![ Field::new(MIN_TS_COLUMN, DataType::Int64, false), @@ -552,7 +552,7 @@ fn validate_metadata(metadata: &RegionMetadataRef) -> Result<()> { ), InvalidMetaSnafu { reason: format!( - "primary-key index internal column name {} is already in use", + "series index internal column name {} is already in use", column.column_schema.name ), } @@ -561,7 +561,7 @@ fn validate_metadata(metadata: &RegionMetadataRef) -> Result<()> { ensure!( metadata.primary_key_encoding == PrimaryKeyEncoding::Sparse, InvalidMetaSnafu { - reason: "primary-key index only supports sparse primary-key encoding", + reason: "series index only supports sparse primary-key encoding", } ); ensure!( @@ -569,24 +569,24 @@ fn validate_metadata(metadata: &RegionMetadataRef) -> Result<()> { .primary_key .starts_with(&[ReservedColumnId::table_id(), ReservedColumnId::tsid()]), InvalidMetaSnafu { - reason: "primary-key index requires (__table_id, __tsid) as the primary-key prefix", + reason: "series index requires (__table_id, __tsid) as the primary-key prefix", } ); let table_id = metadata .column_by_id(ReservedColumnId::table_id()) .context(InvalidMetaSnafu { - reason: "primary-key index metadata is missing __table_id", + reason: "series index metadata is missing __table_id", })?; let tsid = metadata .column_by_id(ReservedColumnId::tsid()) .context(InvalidMetaSnafu { - reason: "primary-key index metadata is missing __tsid", + reason: "series index metadata is missing __tsid", })?; ensure!( table_id.column_schema.data_type == ConcreteDataType::uint32_datatype() && tsid.column_schema.data_type == ConcreteDataType::uint64_datatype(), InvalidMetaSnafu { - reason: "primary-key index requires UInt32 __table_id and UInt64 __tsid", + reason: "series index requires UInt32 __table_id and UInt64 __tsid", } ); for column in metadata.primary_key_columns() { @@ -597,7 +597,7 @@ fn validate_metadata(metadata: &RegionMetadataRef) -> Result<()> { column.column_schema.data_type == ConcreteDataType::string_datatype(), InvalidMetaSnafu { reason: format!( - "primary-key index requires string tag column {}, got {}", + "series index requires string tag column {}, got {}", column.column_schema.name, column.column_schema.data_type ), } @@ -626,7 +626,7 @@ fn timestamp_values(array: &ArrayRef) -> Result { .map(|(array, _)| array) .with_context(|| InvalidRecordBatchSnafu { reason: format!( - "primary-key index requires an Int64 or timestamp time index, got {:?}", + "series index requires an Int64 or timestamp time index, got {:?}", array.data_type() ), })? @@ -634,7 +634,7 @@ fn timestamp_values(array: &ArrayRef) -> Result { ensure!( timestamps.null_count() == 0, InvalidRecordBatchSnafu { - reason: "primary-key index input contains null timestamps", + reason: "series index input contains null timestamps", } ); Ok(timestamps) @@ -648,7 +648,7 @@ fn decode_primary_key( max_ts: i64, row_count: u64, tag_columns: &[(ColumnId, String)], -) -> Result { +) -> Result { let CompositeValues::Sparse(values) = codec.decode(primary_key).context(DecodeSnafu)? else { return InvalidRecordBatchSnafu { reason: "decoded primary key is not sparse", @@ -686,7 +686,7 @@ fn decode_primary_key( .fail(), }) .collect::>>()?; - Ok(PkIndexRow { + Ok(SeriesIndexRow { min_ts, max_ts, row_count, @@ -696,7 +696,7 @@ fn decode_primary_key( }) } -fn rows_to_batch(schema: &SchemaRef, rows: &[PkIndexRow]) -> Result { +fn rows_to_batch(schema: &SchemaRef, rows: &[SeriesIndexRow]) -> Result { let mut arrays: Vec = vec![ Arc::new(Int64Array::from_iter_values( rows.iter().map(|row| row.min_ts), @@ -805,11 +805,11 @@ mod tests { } #[test] - fn test_pk_columns_schema() { + fn test_series_index_schema() { let metadata = Arc::new(sst_region_metadata_with_encoding( PrimaryKeyEncoding::Sparse, )); - let schema = pk_columns_schema(&metadata).unwrap(); + let schema = series_index_schema(&metadata).unwrap(); assert_eq!( schema .fields() @@ -817,9 +817,9 @@ mod tests { .map(|field| field.name().as_str()) .collect::>(), [ - "__pk_min_ts", - "__pk_max_ts", - "__pk_row_count", + "__series_min_ts", + "__series_max_ts", + "__series_row_count", "__table_id", "__tsid", "tag_0", @@ -830,11 +830,11 @@ mod tests { assert!(schema.field(5).is_nullable()); let dense = Arc::new(sst_region_metadata_with_encoding(PrimaryKeyEncoding::Dense)); - assert!(pk_columns_schema(&dense).is_err()); + assert!(series_index_schema(&dense).is_err()); } #[tokio::test] - async fn test_reject_pk_index_internal_column_names_before_opening_writer() { + async fn test_reject_series_index_internal_column_names_before_opening_writer() { for (index, name) in [MIN_TS_COLUMN, MAX_TS_COLUMN, ROW_COUNT_COLUMN] .into_iter() .enumerate() @@ -850,11 +850,11 @@ mod tests { let metadata = Arc::new(builder.build().unwrap()); let store = object_store(); let path = format!("collision-{index}.parquet"); - let error = PkIndexWriter::try_new( + let error = SeriesIndexWriter::try_new( metadata, store.clone(), &path, - PkIndexWriterOptions::default(), + SeriesIndexWriterOptions::default(), ) .await .err() @@ -901,11 +901,11 @@ mod tests { let primary_key_1 = new_sparse_primary_key(&["a", "x"], &metadata, 1, 10); let primary_key_2 = new_sparse_primary_key(&["b", "y"], &metadata, 1, 20); let store = object_store(); - let mut writer = PkIndexWriter::try_new( + let mut writer = SeriesIndexWriter::try_new( metadata, store.clone(), - "pk.parquet", - PkIndexWriterOptions { row_group_size: 2 }, + "series.parquet", + SeriesIndexWriterOptions { row_group_size: 2 }, ) .await .unwrap(); @@ -946,7 +946,7 @@ mod tests { assert_eq!(metrics.num_series, 2); assert!(!metrics.aborted); - let (output_bytes, row_groups, batches) = read_index(&store, "pk.parquet").await; + let (output_bytes, row_groups, batches) = read_index(&store, "series.parquet").await; assert_eq!(metrics.output_bytes, output_bytes); assert_eq!(row_groups, 1); assert_eq!(batches.len(), 1); @@ -1008,11 +1008,11 @@ mod tests { PrimaryKeyEncoding::Sparse, )); let store = object_store(); - let mut writer = PkIndexWriter::try_new( + let mut writer = SeriesIndexWriter::try_new( metadata.clone(), store.clone(), "groups.parquet", - PkIndexWriterOptions { row_group_size: 2 }, + SeriesIndexWriterOptions { row_group_size: 2 }, ) .await .unwrap(); @@ -1028,11 +1028,11 @@ mod tests { let (_, row_groups, _) = read_index(&store, "groups.parquet").await; assert_eq!(row_groups, 3); - let empty = PkIndexWriter::try_new( + let empty = SeriesIndexWriter::try_new( metadata, store.clone(), "empty.parquet", - PkIndexWriterOptions::default(), + SeriesIndexWriterOptions::default(), ) .await .unwrap() @@ -1054,11 +1054,11 @@ mod tests { let primary_key_1 = new_sparse_primary_key(&["a", "x"], &metadata, 1, 10); let primary_key_2 = new_sparse_primary_key(&["b", "y"], &metadata, 1, 20); let store = object_store(); - let mut writer = PkIndexWriter::try_new( + let mut writer = SeriesIndexWriter::try_new( metadata.clone(), store.clone(), "abort.parquet", - PkIndexWriterOptions { row_group_size: 1 }, + SeriesIndexWriterOptions { row_group_size: 1 }, ) .await .unwrap(); @@ -1082,11 +1082,11 @@ mod tests { Bytes::from_static(b"existing") ); - let mut writer = PkIndexWriter::try_new( + let mut writer = SeriesIndexWriter::try_new( metadata, store.clone(), "dictionary-abort.parquet", - PkIndexWriterOptions { row_group_size: 1 }, + SeriesIndexWriterOptions { row_group_size: 1 }, ) .await .unwrap(); @@ -1137,11 +1137,11 @@ mod tests { ) .unwrap(); let store = object_store(); - let mut writer = PkIndexWriter::try_new( + let mut writer = SeriesIndexWriter::try_new( metadata.clone(), store.clone(), "nullable.parquet", - PkIndexWriterOptions::default(), + SeriesIndexWriterOptions::default(), ) .await .unwrap(); @@ -1159,11 +1159,11 @@ mod tests { assert!(tag_1.is_null(0)); assert!( - PkIndexWriter::try_new( + SeriesIndexWriter::try_new( metadata, store, "invalid.parquet", - PkIndexWriterOptions { row_group_size: 0 }, + SeriesIndexWriterOptions { row_group_size: 0 }, ) .await .is_err() From 09c0b23a23450d0aa3bb2d89938ad0f3f5be253c Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 17 Aug 2026 09:53:59 +0000 Subject: [PATCH 57/92] feat: manage semantic table options via ALTER TABLE SET/UNSET (#8880) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(meta): actually acquire logical table locks in alter-logical-tables procedure The procedure listed its logical table locks from table_info_values, which is only filled during Prepare, while procedure lock keys are fixed at submission — so the logical locks were never acquired. Today every writer of a logical table's info is serialized by the physical table lock, which hides the problem; a metadata-only alter procedure targeting a single logical table would race it. Resolve the logical table ids at submission, persist them in the procedure state (serde(default): state dumped by older versions keeps the previous behavior), lock physical + logical tables, and re-check the resolved ids against the locked set at Prepare so a table dropped and recreated after submission cannot be mutated without a lock. Signed-off-by: Dennis Zhuang * feat: manage semantic table options via ALTER TABLE SET/UNSET CREATE TABLE accepts greptime.semantic.* options, but ALTER TABLE SET routed every option through SetRegionOption, whose closed match rejects them — tables auto-created by ingestion could never receive semantic declarations after the fact. Semantic options are pure metadata markers no region consumes, so they now take a metadata-only alter, following the repartition-hint precedent: - New AlterKind::SetAnnotations/UnsetAnnotations carrying an AnnotationFamily (currently only Semantic), so future marker-style option families reuse the same machinery. The converter classifies a SET/UNSET batch by key prefix and rejects batches that mix annotation keys with regular options. - The procedure reuses the MetadataOnly flow: no region dispatch, table-info update plus cache invalidation only. - Validation lives in the table-meta mutation layer, so it runs at frontend verification and again in the procedure's prepare step under the table lock: SET is strict (known key, value domain, entity columns exist and render as strings); UNSET is lenient inside the namespace so stale keys can be cleaned up. ModifyColumnTypes re-checks columns referenced by entity declarations at the same layer, closing a verify-then-execute race. - Logical metric tables are supported: an annotation alter submits a regular alter-table task locking only the logical table, and the DDL manager's physical-route guard admits it. - create_table_info re-checks semantic value domains for gRPC-built expressions that bypass the SQL parser. Signed-off-by: Dennis Zhuang * refactor(table): centralize annotation option classification and validation Address review feedback on the AnnotationFamily abstraction: with only one variant that every consumer immediately destructured, the generality was fake. Make it real and exhaustive instead: - AnnotationFamily gains RepartitionHint: repartition.column.hint is the same kind of marker option (pure metadata, no region consumes it) and previously had a hand-rolled special case in the converter, the metadata-only classifier, and a dedicated AlterKind pair — all deleted, one classification API remains. Per-family logical-table eligibility (allows_logical_tables) replaces the hard-coded Semantic check in the DDL manager guard. - One validation core in the table crate (check_annotation) serves both DDL entry points. CREATE and ALTER previously duplicated the rules; each keeps its existing error variants, status codes and messages via thin adapters over a typed error (ALTER missing column stays 4002 TableColumnNotFound, CREATE stays InvalidArguments). - The batch classifier returns Result instead of swallowing the mixed-batch error: a mixed SET on a logical table now reports the actual problem instead of UnexpectedLogicalRouteTable, and the flow classifiers propagate instead of guessing. The converter also moves its owned payloads instead of cloning them. Signed-off-by: Dennis Zhuang * test(meta): cover logical-table annotation alter routing The route-guard branch admitting metadata-only annotation alters on logical tables was only exercised end to end by sqlness. Pin it at the DDL manager level: a semantic SET on a logical table succeeds, updates only the logical table's metadata and dispatches nothing to datanodes; a mixed batch reports its own error instead of the route guard's; the repartition hint stays rejected on logical routes. Signed-off-by: Dennis Zhuang * fix(table): keep entity guard on ADD COLUMN and report missing columns first Review follow-ups: the old verify_alter loop scanned the post-alter schema, so it also caught DROP COLUMN followed by re-adding the declared column with a non-string type — the mutation-layer move only kept the MODIFY path. Guard add_columns the same way (this also covers ingestion auto-alter). And run the MODIFY drift check after the existence lookup, so altering a dropped-but-still-declared column reports ColumnNotExists (4002) like every other MODIFY on a missing column. Signed-off-by: Dennis Zhuang * style(grpc-expr): drop a test comment restating the classifier doc Signed-off-by: Dennis Zhuang * refactor(table): rename annotation validation helpers per review check_annotation* validated and normalized; align the names with the validate_and_normalize_* convention nearby, and spell out AnnotationContext (Cx is not used in this repo). Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/common/grpc-expr/src/alter.rs | 256 +++++++-- src/common/grpc-expr/src/lib.rs | 2 +- src/common/meta/AGENTS.md | 5 + .../meta/src/ddl/alter_logical_tables.rs | 57 +- src/common/meta/src/ddl/alter_table.rs | 42 +- .../meta/src/ddl/alter_table/executor.rs | 4 +- .../src/ddl/tests/alter_logical_tables.rs | 103 +++- src/common/meta/src/ddl/tests/alter_table.rs | 59 +- src/common/meta/src/ddl/tests/event/table.rs | 1 + src/common/meta/src/ddl_manager.rs | 162 +++++- src/operator/src/statement/ddl.rs | 209 +++---- src/table/src/metadata.rs | 510 ++++++++++++++++-- src/table/src/requests.rs | 217 +++++++- src/table/src/requests/semantic.rs | 26 +- .../information_schema/table_semantics.result | 118 ++++ .../information_schema/table_semantics.sql | 62 +++ .../common/system/semantic_graph.result | 55 ++ .../common/system/semantic_graph.sql | 29 + 18 files changed, 1630 insertions(+), 287 deletions(-) diff --git a/src/common/grpc-expr/src/alter.rs b/src/common/grpc-expr/src/alter.rs index 46de85da4c7..85de11d5cdb 100644 --- a/src/common/grpc-expr/src/alter.rs +++ b/src/common/grpc-expr/src/alter.rs @@ -29,8 +29,8 @@ use snafu::{OptionExt, ResultExt, ensure}; use store_api::region_request::{SetRegionOption, UnsetRegionOption}; use table::metadata::{TableId, TableMeta}; use table::requests::{ - AddColumnRequest, AlterKind, AlterTableRequest, ModifyColumnTypeRequest, - REPARTITION_COLUMN_HINT_KEY, SetDefaultRequest, SetIndexOption, UnsetIndexOption, + AddColumnRequest, AlterKind, AlterTableRequest, AnnotationFamily, ModifyColumnTypeRequest, + SetDefaultRequest, SetIndexOption, UnsetIndexOption, }; use crate::error::{ @@ -44,6 +44,57 @@ use crate::error::{ const LOCATION_TYPE_FIRST: i32 = LocationType::First as i32; const LOCATION_TYPE_AFTER: i32 = LocationType::After as i32; +/// Classifies a SET/UNSET key batch: `Ok(Some(family))` when every key belongs +/// to the same annotation family, `Ok(None)` when none does, and an error on a +/// mixed batch — annotation alters skip region dispatch, so they cannot share +/// a statement with options that regions must see. +fn annotation_family_of_keys<'a>( + mut keys: impl Iterator, +) -> Result> { + let Some(first) = keys.next() else { + return Ok(None); + }; + let family = AnnotationFamily::of_key(first); + let mut count = 1usize; + for key in keys { + count += 1; + let this = AnnotationFamily::of_key(key); + if this != family { + return error::InvalidTableOptionRequestSnafu { + err_msg: family.or(this).unwrap().mixed_batch_error(), + } + .fail(); + } + } + if let Some(family) = family + && family.requires_single_key() + && count > 1 + { + return error::InvalidTableOptionRequestSnafu { + err_msg: family.mixed_batch_error(), + } + .fail(); + } + Ok(family) +} + +/// Returns the annotation family when `kind` is a SET/UNSET whose keys all +/// belong to one family — the alters that only rewrite table metadata and skip +/// region dispatch. A mixed batch is an error; never interpret it as "not an +/// annotation alter", or the batch falls through to a path that reports a +/// misleading error (or dispatches to regions). +pub fn annotation_alter_family(kind: &Kind) -> Result> { + match kind { + Kind::SetTableOptions(api::v1::SetTableOptions { table_options }) => { + annotation_family_of_keys(table_options.iter().map(|option| option.key.as_str())) + } + Kind::UnsetTableOptions(api::v1::UnsetTableOptions { keys }) => { + annotation_family_of_keys(keys.iter().map(|key| key.as_str())) + } + _ => Ok(None), + } +} + fn set_index_option_from_proto(set_index: api::v1::SetIndex) -> Result { let options = set_index.options.context(MissingAlterIndexOptionSnafu)?; Ok(match options { @@ -164,20 +215,15 @@ pub fn alter_expr_to_request( AlterKind::RenameTable { new_table_name } } Kind::SetTableOptions(api::v1::SetTableOptions { table_options }) => { - let repartition_column_hint = table_options - .iter() - .find(|option| option.key == REPARTITION_COLUMN_HINT_KEY); - if let Some(option) = repartition_column_hint { - ensure!( - table_options.len() == 1, - error::InvalidTableOptionRequestSnafu { - err_msg: format!( - "{REPARTITION_COLUMN_HINT_KEY} must be altered separately" - ), - } - ); - AlterKind::SetRepartitionColumnHint { - column_name: option.value.clone(), + if let Some(family) = annotation_family_of_keys( + table_options.iter().map(|option| option.key.as_str()), + )? { + AlterKind::SetAnnotations { + family, + options: table_options + .into_iter() + .map(|option| (option.key, option.value)) + .collect(), } } else { AlterKind::SetTableOptions { @@ -190,19 +236,8 @@ pub fn alter_expr_to_request( } } Kind::UnsetTableOptions(api::v1::UnsetTableOptions { keys }) => { - let unset_repartition_column_hint = keys - .iter() - .any(|key| key.as_str() == REPARTITION_COLUMN_HINT_KEY); - if unset_repartition_column_hint { - ensure!( - keys.len() == 1, - error::InvalidTableOptionRequestSnafu { - err_msg: format!( - "{REPARTITION_COLUMN_HINT_KEY} must be altered separately" - ), - } - ); - AlterKind::UnsetRepartitionColumnHint + if let Some(family) = annotation_family_of_keys(keys.iter().map(|key| key.as_str()))? { + AlterKind::UnsetAnnotations { family, keys } } else { AlterKind::UnsetTableOptions { keys: keys @@ -359,6 +394,7 @@ mod tests { Option as PbOption, SemanticType, SetTableOptions, UnsetTableOptions, }; use datatypes::prelude::ConcreteDataType; + use table::requests::REPARTITION_COLUMN_HINT_KEY; use super::*; @@ -563,8 +599,12 @@ mod tests { let alter_request = alter_expr_to_request(1, expr, None).unwrap(); match alter_request.alter_kind { - AlterKind::SetRepartitionColumnHint { column_name } => { - assert_eq!("host", column_name); + AlterKind::SetAnnotations { family, options } => { + assert_eq!(AnnotationFamily::RepartitionHint, family); + assert_eq!( + vec![(REPARTITION_COLUMN_HINT_KEY.to_string(), "host".to_string())], + options + ); } _ => unreachable!(), } @@ -595,6 +635,30 @@ mod tests { err.to_string() .contains("repartition.column.hint must be altered separately") ); + + // Duplicate hint entries are not a meaningful batch either. + let dup = AlterTableExpr { + catalog_name: "test_catalog".to_string(), + schema_name: "test_schema".to_string(), + table_name: "monitor".to_string(), + kind: Some(Kind::SetTableOptions(SetTableOptions { + table_options: vec![ + PbOption { + key: REPARTITION_COLUMN_HINT_KEY.to_string(), + value: "host".to_string(), + }, + PbOption { + key: REPARTITION_COLUMN_HINT_KEY.to_string(), + value: "region".to_string(), + }, + ], + })), + }; + let err = alter_expr_to_request(1, dup, None).unwrap_err(); + assert!( + err.to_string() + .contains("repartition.column.hint must be altered separately") + ); } #[test] @@ -608,10 +672,138 @@ mod tests { })), }; + let alter_request = alter_expr_to_request(1, expr, None).unwrap(); + match alter_request.alter_kind { + AlterKind::UnsetAnnotations { family, keys } => { + assert_eq!(AnnotationFamily::RepartitionHint, family); + assert_eq!(vec![REPARTITION_COLUMN_HINT_KEY.to_string()], keys); + } + _ => unreachable!(), + } + } + + #[test] + fn test_semantic_options_classified_as_annotations() { + let expr = AlterTableExpr { + catalog_name: "test_catalog".to_string(), + schema_name: "test_schema".to_string(), + table_name: "monitor".to_string(), + kind: Some(Kind::SetTableOptions(SetTableOptions { + table_options: vec![ + PbOption { + key: "greptime.semantic.signal_type".to_string(), + value: "trace".to_string(), + }, + PbOption { + key: "greptime.semantic.entity.host.id".to_string(), + value: "host".to_string(), + }, + ], + })), + }; + + let alter_request = alter_expr_to_request(1, expr, None).unwrap(); + let AlterKind::SetAnnotations { family, options } = alter_request.alter_kind else { + panic!( + "expected SetAnnotations, got {:?}", + alter_request.alter_kind + ); + }; + assert_eq!(family, AnnotationFamily::Semantic); + assert_eq!(options.len(), 2); + + let expr = AlterTableExpr { + catalog_name: "test_catalog".to_string(), + schema_name: "test_schema".to_string(), + table_name: "monitor".to_string(), + kind: Some(Kind::UnsetTableOptions(UnsetTableOptions { + keys: vec!["greptime.semantic.signal_type".to_string()], + })), + }; let alter_request = alter_expr_to_request(1, expr, None).unwrap(); assert!(matches!( alter_request.alter_kind, - AlterKind::UnsetRepartitionColumnHint + AlterKind::UnsetAnnotations { + family: AnnotationFamily::Semantic, + .. + } )); } + + #[test] + fn test_semantic_options_reject_mixed_batch() { + let mixed_set = AlterTableExpr { + catalog_name: "test_catalog".to_string(), + schema_name: "test_schema".to_string(), + table_name: "monitor".to_string(), + kind: Some(Kind::SetTableOptions(SetTableOptions { + table_options: vec![ + PbOption { + key: "greptime.semantic.signal_type".to_string(), + value: "trace".to_string(), + }, + PbOption { + key: table::requests::TTL_KEY.to_string(), + value: "7d".to_string(), + }, + ], + })), + }; + let err = alter_expr_to_request(1, mixed_set, None).unwrap_err(); + assert!(err.to_string().contains("altered separately"), "{err}"); + + let mixed_unset = AlterTableExpr { + catalog_name: "test_catalog".to_string(), + schema_name: "test_schema".to_string(), + table_name: "monitor".to_string(), + kind: Some(Kind::UnsetTableOptions(UnsetTableOptions { + keys: vec![ + "ttl".to_string(), + "greptime.semantic.signal_type".to_string(), + ], + })), + }; + let err = alter_expr_to_request(1, mixed_unset, None).unwrap_err(); + assert!(err.to_string().contains("altered separately"), "{err}"); + } + + #[test] + fn test_annotation_alter_family() { + let mixed = Kind::SetTableOptions(SetTableOptions { + table_options: vec![ + PbOption { + key: "greptime.semantic.signal_type".to_string(), + value: "trace".to_string(), + }, + PbOption { + key: "ttl".to_string(), + value: "7d".to_string(), + }, + ], + }); + let err = annotation_alter_family(&mixed).unwrap_err(); + assert!( + err.to_string().contains("must be altered separately"), + "{err}" + ); + + // Two annotation families cannot share a batch either. + let cross = Kind::SetTableOptions(SetTableOptions { + table_options: vec![ + PbOption { + key: "greptime.semantic.signal_type".to_string(), + value: "trace".to_string(), + }, + PbOption { + key: REPARTITION_COLUMN_HINT_KEY.to_string(), + value: "host".to_string(), + }, + ], + }); + let err = annotation_alter_family(&cross).unwrap_err(); + assert!( + err.to_string().contains("must be altered separately"), + "{err}" + ); + } } diff --git a/src/common/grpc-expr/src/lib.rs b/src/common/grpc-expr/src/lib.rs index 0941cb650d0..9b9994a8f02 100644 --- a/src/common/grpc-expr/src/lib.rs +++ b/src/common/grpc-expr/src/lib.rs @@ -16,4 +16,4 @@ mod alter; pub mod error; pub mod util; -pub use alter::{alter_expr_to_request, create_table_schema}; +pub use alter::{alter_expr_to_request, annotation_alter_family, create_table_schema}; diff --git a/src/common/meta/AGENTS.md b/src/common/meta/AGENTS.md index fd5da08a933..3684f8f9db3 100644 --- a/src/common/meta/AGENTS.md +++ b/src/common/meta/AGENTS.md @@ -26,6 +26,11 @@ and durable DDL procedures. Metasrv process and service wiring live in invalidation. - Procedure changes must preserve persisted state and `TYPE_NAME`; new procedures need loader registration in `src/common/meta/src/ddl_manager.rs`. +- Marker-style table options that no region consumes (`greptime.semantic.*`, + `repartition.column.hint`) go through `AnnotationFamily` in + `src/table/src/requests.rs` — classification, validation, logical-table + eligibility, and the metadata-only alter flow come with it. Do not add + per-option special cases to the alter conversion or DDL routing. - `KvBackend` behavior changes should extend the shared tests in `src/common/meta/src/kv_backend/test.rs`. diff --git a/src/common/meta/src/ddl/alter_logical_tables.rs b/src/common/meta/src/ddl/alter_logical_tables.rs index 9207f910577..ccf130cec8c 100644 --- a/src/common/meta/src/ddl/alter_logical_tables.rs +++ b/src/common/meta/src/ddl/alter_logical_tables.rs @@ -24,7 +24,7 @@ use common_procedure::{Context, EventContext, EventTrigger, LockKey, Procedure, use common_telemetry::{debug, error, info, warn}; pub use executor::make_alter_region_request; use serde::{Deserialize, Serialize}; -use snafu::ResultExt; +use snafu::{ResultExt, ensure}; use store_api::metadata::ColumnMetadata; use store_api::metric_engine_consts::ALTER_PHYSICAL_EXTENSION_KEY; use strum::AsRefStr; @@ -40,7 +40,7 @@ use crate::ddl::event::table::{ TableDdlEvent, TableDdlEventType, TableDdlLocator, alter_table_kind_name, }; use crate::ddl::utils::{extract_column_metadatas, map_to_procedure_error, sync_follower_regions}; -use crate::error::Result; +use crate::error::{self, Result}; use crate::instruction::CacheIdent; use crate::key::DeserializedValueWithBytes; use crate::key::table_info::TableInfoValue; @@ -88,6 +88,7 @@ impl AlterLogicalTablesProcedure { pub fn new( tasks: Vec, physical_table_id: TableId, + logical_table_ids: Vec, context: DdlContext, ) -> Self { Self { @@ -97,6 +98,7 @@ impl AlterLogicalTablesProcedure { tasks, table_info_values: vec![], physical_table_id, + logical_table_ids, physical_table_info: None, physical_columns: vec![], table_cache_keys_to_invalidate: vec![], @@ -147,6 +149,26 @@ impl AlterLogicalTablesProcedure { ); } + // Locks are fixed at submission. A logical table that resolves to an + // id outside the locked set here was dropped and recreated in between, + // so this procedure holds no lock for it. Procedures restored from + // pre-upgrade state have no recorded ids and keep the old behavior. + if !self.data.logical_table_ids.is_empty() { + for value in &table_info_values { + let table_id = value.get_inner_ref().table_info.ident.table_id; + ensure!( + self.data.logical_table_ids.contains(&table_id), + error::UnexpectedSnafu { + err_msg: format!( + "logical table {} (id {table_id}) is not covered by the \ + procedure locks; retry the statement", + value.get_inner_ref().table_info.name + ), + } + ); + } + } + // Updates the procedure state. retain_unskipped(&mut self.data.tasks, &skip_alter); self.data.physical_table_info = Some(physical_table_info); @@ -305,17 +327,29 @@ impl Procedure for AlterLogicalTablesProcedure { // CatalogLock, SchemaLock, // TableLock // TableNameLock(s) - let mut lock_key = Vec::with_capacity(2 + 1 + self.data.tasks.len()); + let mut lock_key = Vec::with_capacity(2 + 1 + self.data.logical_table_ids.len()); let table_ref = self.data.tasks[0].table_ref(); lock_key.push(CatalogLock::Read(table_ref.catalog).into()); lock_key.push(SchemaLock::read(table_ref.catalog, table_ref.schema).into()); lock_key.push(TableLock::Write(self.data.physical_table_id).into()); - lock_key.extend( - self.data - .table_info_values - .iter() - .map(|table| TableLock::Write(table.table_info.ident.table_id).into()), - ); + if self.data.logical_table_ids.is_empty() { + // Pre-upgrade procedure state has no `logical_table_ids`, but a + // dump taken after `Prepare` still carries the resolved table + // snapshots — recover the logical locks from them. + lock_key.extend( + self.data + .table_info_values + .iter() + .map(|table| TableLock::Write(table.table_info.ident.table_id).into()), + ); + } else { + lock_key.extend( + self.data + .logical_table_ids + .iter() + .map(|table_id| TableLock::Write(*table_id).into()), + ); + } LockKey::new(lock_key) } @@ -366,6 +400,11 @@ pub struct AlterTablesData { table_info_values: Vec>, /// Physical table info physical_table_id: TableId, + /// Logical table ids resolved at submission time, so `lock_key` can name + /// them before `Prepare` runs (procedure locks are fixed at submission). + /// Empty when restored from pre-upgrade procedure state. + #[serde(default)] + logical_table_ids: Vec, physical_table_info: Option>, physical_columns: Vec, table_cache_keys_to_invalidate: Vec, diff --git a/src/common/meta/src/ddl/alter_table.rs b/src/common/meta/src/ddl/alter_table.rs index 63342a7bbb6..8e24f3ae2ab 100644 --- a/src/common/meta/src/ddl/alter_table.rs +++ b/src/common/meta/src/ddl/alter_table.rs @@ -21,7 +21,7 @@ use std::vec; use api::region::RegionResponse; use api::v1::alter_table_expr::Kind; -use api::v1::{RenameTable, SetTableOptions, UnsetTableOptions}; +use api::v1::{RenameTable, SetTableOptions}; use async_trait::async_trait; use common_catalog::consts::{METRIC_ENGINE, MITO_ENGINE}; use common_error::ext::BoxedError; @@ -38,7 +38,7 @@ use store_api::metric_engine_consts::TABLE_COLUMN_METADATA_EXTENSION_KEY; use store_api::storage::RegionId; use strum::AsRefStr; use table::metadata::{TableId, TableInfo}; -use table::requests::{REPARTITION_COLUMN_HINT_KEY, SKIP_WAL_KEY}; +use table::requests::SKIP_WAL_KEY; use table::table_reference::TableReference; use crate::ddl::DdlContext; @@ -51,7 +51,8 @@ use crate::ddl::utils::{ sync_follower_regions, }; use crate::error::{ - AbortProcedureSnafu, NoLeaderSnafu, PutPoisonSnafu, Result, RetryLaterSnafu, UnsupportedSnafu, + AbortProcedureSnafu, ConvertAlterTableRequestSnafu, NoLeaderSnafu, PutPoisonSnafu, Result, + RetryLaterSnafu, UnsupportedSnafu, }; use crate::key::table_info::TableInfoValue; use crate::key::{DeserializedValueWithBytes, RegionDistribution}; @@ -201,7 +202,7 @@ impl AlterTableProcedure { self.data.region_distribution = Some(region_distribution(&physical_table_route.region_routes)); } - self.data.state = self.data.flow().after_prepare(); + self.data.state = self.data.flow()?.after_prepare(); Ok(Status::executing(true)) } @@ -249,7 +250,7 @@ impl AlterTableProcedure { ensure!(!leaders.is_empty(), NoLeaderSnafu { table_id }); // Puts the poison before submitting alter region requests to datanodes. self.put_poison(ctx_provider, procedure_id).await?; - let flow = self.data.flow(); + let flow = self.data.flow()?; if flow == AlterTableFlow::MetadataFirst { let results = self .executor @@ -340,7 +341,7 @@ impl AlterTableProcedure { "altering table result doesn't contains extension key `{TABLE_COLUMN_METADATA_EXTENSION_KEY}`,leaving the table's column metadata unchanged" ); } - self.data.state = self.data.flow().after_regions(); + self.data.state = self.data.flow()?.after_regions(); Ok(()) } @@ -372,7 +373,7 @@ impl AlterTableProcedure { let table_info_value = self.data.table_info_value.as_ref().unwrap(); // Safety: Checked in `AlterTableProcedure::new`. let alter_kind = self.data.task.alter_table.kind.as_ref().unwrap(); - let flow = self.data.flow(); + let flow = self.data.flow()?; let metadata_only_alter = flow == AlterTableFlow::MetadataOnly; // Gets the table info from the cache or builds it. @@ -469,17 +470,12 @@ pub(crate) fn only_enables_skip_wal(alter_kind: &Kind) -> bool { && table_options[0].value == "true" } -fn is_metadata_only_alter(alter_kind: &Kind) -> bool { - match alter_kind { - Kind::RenameTable { .. } => true, - Kind::SetTableOptions(SetTableOptions { table_options }) => { - table_options.len() == 1 && table_options[0].key.as_str() == REPARTITION_COLUMN_HINT_KEY - } - Kind::UnsetTableOptions(UnsetTableOptions { keys }) => { - keys.len() == 1 && keys[0].as_str() == REPARTITION_COLUMN_HINT_KEY - } - _ => false, - } +fn is_metadata_only_alter(alter_kind: &Kind) -> Result { + // A mixed annotation batch is an error, never "not metadata-only": falling + // through to the region-first flow would dispatch it to regions. + let family = common_grpc_expr::annotation_alter_family(alter_kind) + .context(ConvertAlterTableRequestSnafu)?; + Ok(family.is_some() || matches!(alter_kind, Kind::RenameTable { .. })) } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -490,14 +486,14 @@ enum AlterTableFlow { } impl AlterTableFlow { - fn from_kind(kind: &Kind) -> Self { - if only_enables_skip_wal(kind) { + fn from_kind(kind: &Kind) -> Result { + Ok(if only_enables_skip_wal(kind) { Self::MetadataFirst - } else if is_metadata_only_alter(kind) { + } else if is_metadata_only_alter(kind)? { Self::MetadataOnly } else { Self::RegionFirst - } + }) } fn after_prepare(self) -> AlterTableState { @@ -647,7 +643,7 @@ impl AlterTableData { .map(|value| &value.table_info) } - fn flow(&self) -> AlterTableFlow { + fn flow(&self) -> Result { // Safety: Checked in `AlterTableProcedure::new`. AlterTableFlow::from_kind(self.task.alter_table.kind.as_ref().unwrap()) } diff --git a/src/common/meta/src/ddl/alter_table/executor.rs b/src/common/meta/src/ddl/alter_table/executor.rs index fbacdd1b2f5..65fa45d0a8f 100644 --- a/src/common/meta/src/ddl/alter_table/executor.rs +++ b/src/common/meta/src/ddl/alter_table/executor.rs @@ -335,8 +335,8 @@ fn build_new_table_info( | AlterKind::ModifyColumnTypes { .. } | AlterKind::SetTableOptions { .. } | AlterKind::UnsetTableOptions { .. } - | AlterKind::SetRepartitionColumnHint { .. } - | AlterKind::UnsetRepartitionColumnHint + | AlterKind::SetAnnotations { .. } + | AlterKind::UnsetAnnotations { .. } | AlterKind::SetIndexes { .. } | AlterKind::UnsetIndexes { .. } | AlterKind::DropDefaults { .. } diff --git a/src/common/meta/src/ddl/tests/alter_logical_tables.rs b/src/common/meta/src/ddl/tests/alter_logical_tables.rs index 351c0fd96ec..619e02bb9b9 100644 --- a/src/common/meta/src/ddl/tests/alter_logical_tables.rs +++ b/src/common/meta/src/ddl/tests/alter_logical_tables.rs @@ -43,6 +43,7 @@ use crate::error::Error::{AlterLogicalTablesInvalidArguments, TableNotFound}; use crate::error::Result; use crate::key::table_name::TableNameKey; use crate::key::table_route::{PhysicalTableRouteValue, TableRouteValue}; +use crate::lock_key::TableLock; use crate::rpc::ddl::AlterTableTask; use crate::rpc::router::{Region, RegionRoute}; use crate::test_util::{MockDatanodeManager, new_ddl_context}; @@ -162,7 +163,8 @@ async fn test_on_prepare_check_schema() { ), ]; let physical_table_id = 1024u32; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); + let mut procedure = + AlterLogicalTablesProcedure::new(tasks, physical_table_id, vec![], ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, AlterLogicalTablesInvalidArguments { .. }); } @@ -177,7 +179,8 @@ async fn test_on_prepare_check_alter_kind() { "new_table1", )]; let physical_table_id = 1024u32; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, physical_table_id, ddl_context); + let mut procedure = + AlterLogicalTablesProcedure::new(tasks, physical_table_id, vec![], ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, AlterLogicalTablesInvalidArguments { .. }); } @@ -197,7 +200,7 @@ async fn test_on_prepare_different_physical_table() { make_alter_logical_table_add_column_task(None, "table2", vec!["column2".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy1_id, ddl_context); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy1_id, vec![], ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, AlterLogicalTablesInvalidArguments { .. }); } @@ -218,7 +221,7 @@ async fn test_on_prepare_logical_table_not_exists() { make_alter_logical_table_add_column_task(None, "table2", vec!["column2".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, vec![], ddl_context); let err = procedure.on_prepare().await.unwrap_err(); assert_matches!(err, TableNotFound { .. }); } @@ -241,7 +244,7 @@ async fn test_on_prepare() { make_alter_logical_table_add_column_task(None, "table3", vec!["column3".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, vec![], ddl_context); let result = procedure.on_prepare().await; assert_matches!( result, @@ -277,7 +280,8 @@ async fn test_on_update_metadata() { make_alter_logical_table_add_column_task(None, "table3", vec!["new_col".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context.clone()); + let mut procedure = + AlterLogicalTablesProcedure::new(tasks, phy_id, vec![], ddl_context.clone()); let mut status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -362,7 +366,8 @@ async fn test_on_part_duplicate_alter_request() { make_alter_logical_table_add_column_task(None, "table2", vec!["col_0".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context.clone()); + let mut procedure = + AlterLogicalTablesProcedure::new(tasks, phy_id, vec![], ddl_context.clone()); let mut status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -448,7 +453,8 @@ async fn test_on_part_duplicate_alter_request() { ), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context.clone()); + let mut procedure = + AlterLogicalTablesProcedure::new(tasks, phy_id, vec![], ddl_context.clone()); let mut status = procedure.on_prepare().await.unwrap(); assert_matches!( status, @@ -620,7 +626,7 @@ async fn test_on_submit_alter_region_request() { make_alter_logical_table_add_column_task(None, "table2", vec!["mew_col".to_string()]), ]; - let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, ddl_context); + let mut procedure = AlterLogicalTablesProcedure::new(tasks, phy_id, vec![], ddl_context); procedure.on_prepare().await.unwrap(); procedure.on_submit_alter_region_requests().await.unwrap(); let mut results = Vec::new(); @@ -646,3 +652,82 @@ async fn test_on_submit_alter_region_request() { }) ); } + +#[test] +fn test_lock_key_covers_submission_resolved_logical_tables() { + let node_manager = Arc::new(MockDatanodeManager::new(())); + let ddl_context = new_ddl_context(node_manager); + let tasks = vec![make_alter_logical_table_add_column_task( + None, + "table1", + vec!["column1".to_string()], + )]; + let procedure = AlterLogicalTablesProcedure::new(tasks, 1024, vec![1025, 1026], ddl_context); + let keys = procedure.lock_key(); + let keys = keys.keys_to_lock().collect::>(); + for table_id in [1024u32, 1025, 1026] { + let expected: common_procedure::StringKey = TableLock::Write(table_id).into(); + assert!(keys.contains(&&expected), "missing lock for {table_id}"); + } +} + +#[tokio::test] +async fn test_from_json_without_logical_table_ids() { + let node_manager = Arc::new(MockDatanodeManager::new(())); + let ddl_context = new_ddl_context(node_manager); + let phy_id = create_physical_table(&ddl_context, "phy").await; + let logical_id = create_logical_table(ddl_context.clone(), phy_id, "table1").await; + let tasks = vec![make_alter_logical_table_add_column_task( + None, + "table1", + vec!["column1".to_string()], + )]; + let mut procedure = + AlterLogicalTablesProcedure::new(tasks, phy_id, vec![logical_id], ddl_context.clone()); + procedure.on_prepare().await.unwrap(); + let json = procedure.dump().unwrap(); + + // Pre-upgrade procedure state carries no `logical_table_ids`, but a + // post-Prepare dump still holds the resolved table snapshots. + let mut value: serde_json::Value = serde_json::from_str(&json).unwrap(); + value + .as_object_mut() + .unwrap() + .remove("logical_table_ids") + .unwrap(); + let restored = AlterLogicalTablesProcedure::from_json(&value.to_string(), ddl_context).unwrap(); + + let keys = restored.lock_key(); + let keys = keys.keys_to_lock().collect::>(); + for table_id in [phy_id, logical_id] { + let expected: common_procedure::StringKey = TableLock::Write(table_id).into(); + assert!(keys.contains(&&expected), "missing lock for {table_id}"); + } +} + +#[tokio::test] +async fn test_on_prepare_rejects_table_outside_locked_set() { + let node_manager = Arc::new(MockDatanodeManager::new(())); + let ddl_context = new_ddl_context(node_manager); + + let phy_id = create_physical_table(&ddl_context, "phy").await; + create_logical_table(ddl_context.clone(), phy_id, "table1").await; + + let tasks = vec![make_alter_logical_table_add_column_task( + None, + "table1", + vec!["column1".to_string()], + )]; + + // The table resolved at prepare time is not in the locked set, as if it + // had been dropped and recreated after submission. + let stale_locked_id = phy_id + 1000; + let mut procedure = + AlterLogicalTablesProcedure::new(tasks, phy_id, vec![stale_locked_id], ddl_context); + let err = procedure.on_prepare().await.unwrap_err(); + assert!( + err.to_string() + .contains("not covered by the procedure locks"), + "{err}" + ); +} diff --git a/src/common/meta/src/ddl/tests/alter_table.rs b/src/common/meta/src/ddl/tests/alter_table.rs index b35666e52e1..30910eb0ba5 100644 --- a/src/common/meta/src/ddl/tests/alter_table.rs +++ b/src/common/meta/src/ddl/tests/alter_table.rs @@ -29,7 +29,7 @@ use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_procedure::store::poison_store::PoisonStore; use common_procedure::{Procedure, ProcedureId, Status}; -use common_procedure_test::MockContextProvider; +use common_procedure_test::{MockContextProvider, execute_procedure_until_done}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; use store_api::metadata::ColumnMetadata; @@ -1121,3 +1121,60 @@ async fn test_on_submit_alter_request_with_exist_poison() { .unwrap_err(); assert_matches!(err, Error::PutPoison { .. }); } + +#[tokio::test] +async fn test_semantic_annotation_alter_is_metadata_only() { + let (tx, mut rx) = mpsc::channel(8); + let node_manager = Arc::new(MockDatanodeManager::new(DatanodeWatcher::new(tx))); + let ddl_context = new_ddl_context(node_manager); + let table_id = 1024; + let table_name = "foo"; + let task = test_create_table_task(table_name, table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info.clone(), + prepare_table_route(table_id), + HashMap::new(), + ) + .await + .unwrap(); + + let alter_table_task = AlterTableTask { + alter_table: AlterTableExpr { + catalog_name: DEFAULT_CATALOG_NAME.to_string(), + schema_name: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + kind: Some(Kind::SetTableOptions(SetTableOptions { + table_options: vec![api::v1::Option { + key: "greptime.semantic.signal_type".to_string(), + value: "metric".to_string(), + }], + })), + }, + }; + let mut procedure = + AlterTableProcedure::new(table_id, alter_table_task, ddl_context.clone()).unwrap(); + execute_procedure_until_done(&mut procedure).await; + + // Metadata-only: no region request reaches any datanode. + rx.try_recv().unwrap_err(); + + let table_info = ddl_context + .table_metadata_manager + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap() + .into_inner() + .table_info; + assert_eq!( + table_info + .meta + .options + .extra_options + .get("greptime.semantic.signal_type"), + Some(&"metric".to_string()) + ); +} diff --git a/src/common/meta/src/ddl/tests/event/table.rs b/src/common/meta/src/ddl/tests/event/table.rs index 24182fd19af..92fa10ee9e0 100644 --- a/src/common/meta/src/ddl/tests/event/table.rs +++ b/src/common/meta/src/ddl/tests/event/table.rs @@ -391,6 +391,7 @@ fn procedure_cases() -> Vec { ), ], 43, + vec![], test_context(), ); let drop_table = DropTableProcedure::new( diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index 7ba408ca2d9..d88f08918e9 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -53,7 +53,7 @@ use crate::ddl::truncate_table::TruncateTableProcedure; use crate::ddl::undrop_table::UndropTableProcedure; use crate::ddl::{DdlContext, utils}; use crate::error::{ - self, CreateRepartitionProcedureSnafu, EmptyDdlTasksSnafu, + self, ConvertAlterTableRequestSnafu, CreateRepartitionProcedureSnafu, EmptyDdlTasksSnafu, PersistRepartitionGcRequirementSnafu, ProcedureOutputSnafu, RegisterProcedureLoaderSnafu, RegisterRepartitionProcedureLoaderSnafu, Result, SubmitProcedureSnafu, TableInfoNotFoundSnafu, TableNotFoundSnafu, TableRouteNotFoundSnafu, UnexpectedLogicalRouteTableSnafu, @@ -554,8 +554,26 @@ impl DdlManager { ) -> Result<(ProcedureId, Option)> { let context = self.create_context(); - let procedure = - AlterLogicalTablesProcedure::new(alter_table_tasks, physical_table_id, context); + // Resolve the logical table ids up front: procedure locks are fixed + // at submission, so `lock_key` cannot derive them during `Prepare`. + let logical_table_ids = { + let table_refs = alter_table_tasks + .iter() + .map(|task| task.table_ref()) + .collect::>(); + utils::table_id::get_all_table_ids_by_names( + self.table_metadata_manager().table_name_manager(), + &table_refs, + ) + .await? + }; + + let procedure = AlterLogicalTablesProcedure::new( + alter_table_tasks, + physical_table_id, + logical_table_ids, + context, + ); let procedure_with_id = ProcedureWithId::with_random_id(Box::new(procedure)).with_context(procedure_context); @@ -1053,8 +1071,17 @@ async fn handle_alter_table_task( .get(table_id) .await? .context(TableRouteNotFoundSnafu { table_id })?; + // Classify before the route guard: a mixed annotation batch must surface + // its own error here, not a misleading "non-physical route" one. Families + // that only rewrite the logical table's metadata may target logical tables. + let annotation_family = match alter_table_task.alter_table.kind.as_ref() { + Some(kind) => common_grpc_expr::annotation_alter_family(kind) + .context(ConvertAlterTableRequestSnafu)?, + None => None, + }; ensure!( - table_route_value.is_physical(), + table_route_value.is_physical() + || annotation_family.is_some_and(|family| family.allows_logical_tables()), UnexpectedLogicalRouteTableSnafu { err_msg: format!("{:?} is a non-physical TableRouteValue.", table_ref), } @@ -1499,7 +1526,9 @@ mod tests { use common_event_recorder::{PersistentEventContext, ProcedureEventInput, TriggerReason}; use common_procedure::local::LocalManager; use common_procedure::test_util::InMemoryPoisonStore; - use common_procedure::{BoxedProcedure, ProcedureContext, ProcedureManagerRef}; + use common_procedure::{ + BoxedProcedure, ProcedureContext, ProcedureManager, ProcedureManagerRef, + }; use store_api::storage::TableId; use table::table_name::TableName; @@ -1900,4 +1929,127 @@ mod tests { assert!(matches!(err, crate::error::Error::Unsupported { .. })); } } + + async fn ddl_manager_with_context(ddl_context: DdlContext) -> DdlManager { + let kv_backend = Arc::new(MemoryKvBackend::new()); + let state_store = Arc::new(KvStateStore::new(kv_backend.clone())); + let poison_manager = Arc::new(InMemoryPoisonStore::default()); + let procedure_manager = Arc::new(LocalManager::new( + Default::default(), + state_store, + poison_manager, + None, + None, + )); + procedure_manager.start().await.unwrap(); + let ddl_manager = DdlManager::new( + ddl_context, + procedure_manager, + Arc::new(DummyRepartitionProcedureFactory), + ); + ddl_manager.register_loaders().unwrap(); + ddl_manager + } + + fn set_options_expr(table_name: &str, options: &[(&str, &str)]) -> api::v1::AlterTableExpr { + api::v1::AlterTableExpr { + catalog_name: common_catalog::consts::DEFAULT_CATALOG_NAME.to_string(), + schema_name: common_catalog::consts::DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + kind: Some(api::v1::alter_table_expr::Kind::SetTableOptions( + api::v1::SetTableOptions { + table_options: options + .iter() + .map(|(key, value)| api::v1::Option { + key: key.to_string(), + value: value.to_string(), + }) + .collect(), + }, + )), + } + } + + #[tokio::test] + async fn test_logical_table_annotation_alter_routing() { + let (tx, mut rx) = tokio::sync::mpsc::channel(8); + let node_manager = Arc::new(crate::test_util::MockDatanodeManager::new( + crate::ddl::test_util::datanode_handler::DatanodeWatcher::new(tx), + )); + let ddl_context = crate::test_util::new_ddl_context(node_manager); + let phy_id = crate::ddl::test_util::create_physical_table(&ddl_context, "phy").await; + let logical_id = + crate::ddl::test_util::create_logical_table(ddl_context.clone(), phy_id, "logical") + .await; + let ddl_manager = ddl_manager_with_context(ddl_context.clone()).await; + + // A semantic alter on a logical table passes the route guard, updates + // only the logical table's metadata, and dispatches nothing. + ddl_manager + .submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + SubmitDdlTaskRequest::new(DdlTask::new_alter_table(set_options_expr( + "logical", + &[("greptime.semantic.signal_type", "metric")], + ))), + ) + .await + .unwrap(); + rx.try_recv().unwrap_err(); + let table_info = ddl_manager + .table_metadata_manager() + .table_info_manager() + .get(logical_id) + .await + .unwrap() + .unwrap() + .into_inner() + .table_info; + assert_eq!( + table_info + .meta + .options + .extra_options + .get("greptime.semantic.signal_type"), + Some(&"metric".to_string()) + ); + + // A mixed batch fails with its own error, not the route guard's. + let err = ddl_manager + .submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + SubmitDdlTaskRequest::new(DdlTask::new_alter_table(set_options_expr( + "logical", + &[("greptime.semantic.source", "prometheus"), ("ttl", "7d")], + ))), + ) + .await + .unwrap_err(); + let msg = common_error::ext::ErrorExt::output_msg(&err); + assert!(msg.contains("must be altered separately"), "{msg}"); + + // The repartition hint drives physical repartitioning; on a logical + // route it stays rejected by the guard. + let err = ddl_manager + .submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + SubmitDdlTaskRequest::new(DdlTask::new_alter_table(set_options_expr( + "logical", + &[("repartition.column.hint", "host")], + ))), + ) + .await + .unwrap_err(); + let msg = common_error::ext::ErrorExt::output_msg(&err); + assert!(msg.contains("non-physical TableRouteValue"), "{msg}"); + } } diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 3c10186ca4e..f89f5de73a7 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -98,8 +98,8 @@ use table::TableRef; use table::dist_table::DistTable; use table::metadata::{self, TableId, TableInfo, TableMeta, TableType}; use table::requests::{ - AlterKind, AlterTableRequest, COMMENT_KEY, DDL_TIMEOUT, DDL_WAIT, REPARTITION_COLUMN_HINT_KEY, - TableOptions, parse_entity_columns, parse_entity_option_key, + AlterKind, AlterTableRequest, AnnotationContext, COMMENT_KEY, DDL_TIMEOUT, DDL_WAIT, + TableOptions, validate_and_normalize_annotation_options, }; use table::table_name::TableName; use table::table_reference::TableReference; @@ -1881,8 +1881,20 @@ impl StatementExecutor { (req, invalidate_keys) } else { - // This is logical table - let req = SubmitDdlTaskRequest::new(DdlTask::new_alter_logical_tables(vec![expr])); + // This is logical table. Annotation alters only rewrite its own + // metadata; `AlterLogicalTablesProcedure` only handles column adds. + let annotation_alter = match expr.kind.as_ref() { + Some(kind) => common_grpc_expr::annotation_alter_family(kind) + .context(AlterExprToRequestSnafu)? + .is_some_and(|family| family.allows_logical_tables()), + None => false, + }; + let task = if annotation_alter { + DdlTask::new_alter_table(expr) + } else { + DdlTask::new_alter_logical_tables(vec![expr]) + }; + let req = SubmitDdlTaskRequest::new(task); let mut invalidate_keys = vec![ CacheIdent::TableId(physical_table_id), @@ -2338,19 +2350,6 @@ pub fn verify_alter( .context(error::BuildTableMetaSnafu { table_name })?; validate_json2_columns_append_mode(&new_meta.schema, &new_meta.options)?; - // A column referenced by an entity declaration may be dropped (the - // read-time derivation skips the stale declaration), but must not change - // to a type without a stable string form. - for (key, value) in &new_meta.options.extra_options { - if parse_entity_option_key(key).is_none() { - continue; - } - for col in &parse_entity_columns(value) { - if let Some(column) = new_meta.schema.column_schema_by_name(col) { - ensure_entity_column_renders_as_string(key, col, &column.data_type)?; - } - } - } Ok(true) } @@ -2402,14 +2401,7 @@ pub fn create_table_info( validate_json2_columns_append_mode(&schema, &table_options)?; - validate_repartition_column_hint( - &mut table_options, - &column_name_to_index_map, - &partition_key_indices, - &create_table.time_index, - )?; - - validate_entity_semantic_options(&table_options, &schema)?; + validate_and_normalize_annotations(&mut table_options, &schema, &partition_key_indices)?; let meta = TableMeta { schema, @@ -2469,61 +2461,6 @@ fn validate_json2_columns_append_mode(schema: &Schema, table_options: &TableOpti Ok(()) } -fn validate_repartition_column_hint( - table_options: &mut TableOptions, - column_name_to_index_map: &HashMap, - partition_key_indices: &[usize], - time_index: &str, -) -> Result<()> { - let Some(column_name) = table_options - .extra_options - .get(REPARTITION_COLUMN_HINT_KEY) - .map(|value| value.trim().to_string()) - else { - return Ok(()); - }; - - ensure!( - !column_name.is_empty(), - InvalidPartitionRuleSnafu { - reason: format!("{REPARTITION_COLUMN_HINT_KEY} expects exactly one column name"), - } - ); - - ensure!( - !column_name.contains(','), - InvalidPartitionRuleSnafu { - reason: format!("{REPARTITION_COLUMN_HINT_KEY} expects exactly one column name"), - } - ); - - ensure!( - partition_key_indices.is_empty(), - InvalidPartitionRuleSnafu { - reason: format!( - "cannot set {REPARTITION_COLUMN_HINT_KEY} on a table with partition metadata" - ), - } - ); - - column_name_to_index_map - .get(&column_name) - .context(ColumnNotFoundSnafu { msg: &column_name })?; - - ensure!( - column_name != time_index, - InvalidPartitionRuleSnafu { - reason: format!("cannot set {REPARTITION_COLUMN_HINT_KEY} to the time index column"), - } - ); - - table_options - .extra_options - .insert(REPARTITION_COLUMN_HINT_KEY.to_string(), column_name); - - Ok(()) -} - /// Rejects DDL against read-only schemas and computed entity-graph tables. fn ensure_table_writable(schema: &str, table: &str) -> Result<()> { ensure!( @@ -2557,57 +2494,34 @@ fn ensure_table_definition_writable(schema: &str, table: &str) -> Result<()> { Ok(()) } -/// Whether `CAST(column AS Utf8)` yields a stable entity-id string — the -/// binary-backed and nested types do not, and without the DDL check the -/// failure would surface only when the graph is scanned. -fn has_stable_string_form(data_type: &datatypes::prelude::ConcreteDataType) -> bool { - use datatypes::prelude::ConcreteDataType; - !matches!( - data_type, - ConcreteDataType::Binary(_) - | ConcreteDataType::Json(_) - | ConcreteDataType::Vector(_) - | ConcreteDataType::List(_) - | ConcreteDataType::Struct(_) - | ConcreteDataType::Dictionary(_) - | ConcreteDataType::Null(_) - ) -} - -fn ensure_entity_column_renders_as_string( - key: &str, - col: &str, - data_type: &datatypes::prelude::ConcreteDataType, +/// CREATE-side annotation validation: one rule source in the table crate, +/// mapped onto this crate's existing error variants so client-visible codes +/// and messages stay put. +fn validate_and_normalize_annotations( + options: &mut TableOptions, + schema: &Schema, + partition_key_indices: &[usize], ) -> Result<()> { - ensure!( - has_stable_string_form(data_type), - InvalidSqlSnafu { - err_msg: format!( - "entity column `{col}` (option `{key}`) has type `{data_type}`, which cannot \ - render as a string", - ), + use table::requests::AnnotationValidationError as CheckError; + let cx = AnnotationContext { + schema, + partition_key_indices, + }; + validate_and_normalize_annotation_options(options, &cx).map_err(|e| match e { + CheckError::ColumnNotFound { column } => ColumnNotFoundSnafu { msg: column }.build(), + e @ (CheckError::UnknownKey { .. } + | CheckError::InvalidValue { .. } + | CheckError::ColumnNotStringForm { .. }) => InvalidSqlSnafu { + err_msg: e.to_string(), } - ); - Ok(()) -} - -/// Validates `greptime.semantic.entity..{id|descriptive|scope}` options -/// against the table schema: every named column must exist (tag or field) and -/// have a stable string form — the derivation renders id, scope and -/// descriptive values as strings. -fn validate_entity_semantic_options(table_options: &TableOptions, schema: &Schema) -> Result<()> { - for (key, value) in &table_options.extra_options { - if parse_entity_option_key(key).is_none() { - continue; - }; - for col in &parse_entity_columns(value) { - let column = schema - .column_schema_by_name(col) - .context(ColumnNotFoundSnafu { msg: col })?; - ensure_entity_column_renders_as_string(key, col, &column.data_type)?; + .build(), + e @ (CheckError::NotSingleColumn + | CheckError::PartitionMetadataConflict + | CheckError::TimeIndexConflict) => InvalidPartitionRuleSnafu { + reason: e.to_string(), } - } - Ok(()) + .build(), + }) } fn find_partition_columns(partitions: &Option) -> Result> { @@ -2904,6 +2818,7 @@ mod test { use sql::parser::{ParseOptions, ParserContext}; use sql::statements::statement::Statement; use sqlparser::parser::Parser; + use table::requests::REPARTITION_COLUMN_HINT_KEY; use super::*; use crate::expr_helper; @@ -3133,7 +3048,7 @@ mod test { } #[test] - fn test_validate_entity_semantic_options() { + fn test_validate_and_normalize_annotations() { let schema = Schema::new(vec![ ColumnSchema::new("service_name", ConcreteDataType::string_datatype(), true), ColumnSchema::new("host_id", ConcreteDataType::string_datatype(), true), @@ -3147,6 +3062,10 @@ mod test { .collect(), ..Default::default() }; + let check = |pairs: &[(&str, &str)]| { + let mut options = opts(pairs); + validate_and_normalize_annotations(&mut options, &schema, &[]).map(|()| options) + }; // Any existing column with a string form may be an id, tag or field. for pairs in [ @@ -3156,22 +3075,27 @@ mod test { )][..], &[("greptime.semantic.entity.service.id", "value")][..], ] { - assert!(validate_entity_semantic_options(&opts(pairs), &schema).is_ok()); + assert!(check(pairs).is_ok()); } - let missing = validate_entity_semantic_options( - &opts(&[("greptime.semantic.entity.service.id", "nope")]), - &schema, - ) - .unwrap_err(); + // Missing columns keep this crate's error variant and status code. + let missing = check(&[("greptime.semantic.entity.service.id", "nope")]).unwrap_err(); assert!(matches!(missing, error::Error::ColumnNotFound { .. })); + assert_eq!( + common_error::status_code::StatusCode::InvalidArguments, + common_error::ext::ErrorExt::status_code(&missing) + ); - let binary_id = validate_entity_semantic_options( - &opts(&[("greptime.semantic.entity.service.id", "payload")]), - &schema, - ) - .unwrap_err(); + let binary_id = check(&[("greptime.semantic.entity.service.id", "payload")]).unwrap_err(); assert!(matches!(binary_id, error::Error::InvalidSql { .. })); + + // The SQL parser checks keys and value domains, but gRPC expressions + // bypass it. + let bad_value = check(&[("greptime.semantic.signal_type", "garbage")]).unwrap_err(); + assert!(matches!(bad_value, error::Error::InvalidSql { .. })); + + let unknown_key = check(&[("greptime.semantic.nonsense", "x")]).unwrap_err(); + assert!(matches!(unknown_key, error::Error::InvalidSql { .. })); } #[test] @@ -3559,9 +3483,10 @@ WITH ('repartition.column.hint' = ' host ')", }], })); let err = verify_alter(1, table_info.clone(), modify).unwrap_err(); + let msg = common_error::ext::ErrorExt::output_msg(&err); assert!( - err.to_string().contains("cannot render as a string"), - "{err}" + msg.contains("must keep a type that renders as a string"), + "{msg}" ); // Dropping a declared column stays allowed: the derivation skips the diff --git a/src/table/src/metadata.rs b/src/table/src/metadata.rs index 7b8958e22ce..4f7ed1c057d 100644 --- a/src/table/src/metadata.rs +++ b/src/table/src/metadata.rs @@ -37,8 +37,10 @@ use store_api::storage::{ColumnDescriptor, ColumnDescriptorBuilder, ColumnId}; use crate::error::{self, Result}; use crate::requests::{ - AddColumnRequest, AlterKind, ModifyColumnTypeRequest, REPARTITION_COLUMN_HINT_KEY, - SetDefaultRequest, SetIndexOption, TableOptions, UnsetIndexOption, + AddColumnRequest, AlterKind, AnnotationContext, AnnotationFamily, AnnotationValidationError, + ModifyColumnTypeRequest, REPARTITION_COLUMN_HINT_KEY, SetDefaultRequest, SetIndexOption, + TableOptions, UnsetIndexOption, has_stable_string_form, parse_entity_columns, + parse_entity_option_key, validate_and_normalize_annotation, }; use crate::table_reference::TableReference; @@ -331,10 +333,12 @@ impl TableMeta { AlterKind::RenameTable { .. } => Ok(self.new_meta_builder()), AlterKind::SetTableOptions { options } => self.set_table_options(options), AlterKind::UnsetTableOptions { keys } => self.unset_table_options(keys), - AlterKind::SetRepartitionColumnHint { column_name } => { - self.set_repartition_column_hint(table_name, column_name) + AlterKind::SetAnnotations { family, options } => { + self.set_annotations(table_name, *family, options) + } + AlterKind::UnsetAnnotations { family, keys } => { + self.unset_annotations(table_name, *family, keys) } - AlterKind::UnsetRepartitionColumnHint => self.unset_repartition_column_hint(), AlterKind::SetIndexes { options } => self.set_indexes(table_name, options), AlterKind::UnsetIndexes { options } => self.unset_indexes(table_name, options), AlterKind::DropDefaults { names } => self.drop_defaults(table_name, names), @@ -422,67 +426,92 @@ impl TableMeta { self.set_table_options(&requests) } - fn set_repartition_column_hint( + /// Applies an annotation SET. Validation lives here, on the mutation + /// path, so it runs both at frontend verification and again inside the + /// alter procedure's prepare step — under the table lock, against fresh + /// metadata. + fn set_annotations( &self, table_name: &str, - column_name: &str, + family: AnnotationFamily, + options: &[(String, String)], ) -> Result { - let column_name = column_name.trim(); ensure!( - !column_name.is_empty() && !column_name.contains(','), + !family.requires_single_key() || options.len() == 1, error::InvalidAlterRequestSnafu { table: table_name, - err: format!("{REPARTITION_COLUMN_HINT_KEY} expects exactly one column name"), + err: family.mixed_batch_error(), } ); - - ensure!( - self.partition_key_indices.is_empty(), - error::InvalidAlterRequestSnafu { - table: table_name, - err: format!( - "cannot set {REPARTITION_COLUMN_HINT_KEY} on a table with partition metadata" - ), - } - ); - - let column_index = self - .schema - .column_index_by_name(column_name) - .with_context(|| error::ColumnNotExistsSnafu { - column_name, - table_name, - })?; - - if let Some(time_index) = self.schema.timestamp_index() { + let cx = AnnotationContext { + schema: &self.schema, + partition_key_indices: &self.partition_key_indices, + }; + let mut new_options = self.options.clone(); + for (key, value) in options { ensure!( - column_index != time_index, + AnnotationFamily::of_key(key) == Some(family), error::InvalidAlterRequestSnafu { table: table_name, err: format!( - "cannot set {REPARTITION_COLUMN_HINT_KEY} to the time index column" + "`{key}` is outside the `{}` annotation namespace", + family.namespace() ), } ); + let checked = validate_and_normalize_annotation(family, &cx, key, value).map_err( + |e| match e { + AnnotationValidationError::ColumnNotFound { column } => { + error::ColumnNotExistsSnafu { + column_name: column, + table_name, + } + .build() + } + other => error::InvalidAlterRequestSnafu { + table: table_name, + err: other.to_string(), + } + .build(), + }, + )?; + new_options.extra_options.insert(key.clone(), checked); } - - let mut new_options = self.options.clone(); - new_options.extra_options.insert( - REPARTITION_COLUMN_HINT_KEY.to_string(), - column_name.to_string(), - ); - let mut builder = self.new_meta_builder(); builder.options(new_options); Ok(builder) } - fn unset_repartition_column_hint(&self) -> Result { + /// Applies an annotation UNSET. Deliberately lenient inside the family's + /// namespace: keys this version does not recognise may still be removed, + /// so options left behind by other versions can be cleaned up. + fn unset_annotations( + &self, + table_name: &str, + family: AnnotationFamily, + keys: &[String], + ) -> Result { + ensure!( + !family.requires_single_key() || keys.len() == 1, + error::InvalidAlterRequestSnafu { + table: table_name, + err: family.mixed_batch_error(), + } + ); let mut new_options = self.options.clone(); - new_options - .extra_options - .remove(REPARTITION_COLUMN_HINT_KEY); - + for key in keys { + ensure!( + AnnotationFamily::of_key(key) == Some(family), + error::InvalidAlterRequestSnafu { + table: table_name, + err: format!( + "`{key}` is outside the `{}` annotation namespace", + family.namespace() + ), + } + ); + new_options.extra_options.remove(key); + } let mut builder = self.new_meta_builder(); builder.options(new_options); Ok(builder) @@ -773,6 +802,24 @@ impl TableMeta { }, ); + // A dropped column may leave a stale entity declaration behind; + // re-adding it must not hand the declaration a type without a + // stable string form. + if !has_stable_string_form(&col_to_add.column_schema.data_type) + && let Some(key) = + entity_option_referencing(&self.options, &col_to_add.column_schema.name) + { + return error::InvalidAlterRequestSnafu { + table: table_name, + err: format!( + "column `{}` is referenced by entity option `{key}` and must \ + keep a type that renders as a string, got `{}`", + col_to_add.column_schema.name, col_to_add.column_schema.data_type + ), + } + .fail(); + } + new_columns.push(col_to_add.clone()); } } @@ -966,6 +1013,25 @@ impl TableMeta { table_name, })?; + // A column referenced by an entity declaration may be dropped (the + // read-time derivation skips the stale declaration), but must not + // change to a type without a stable string form. Checked after the + // existence lookup so a missing column keeps reporting + // `ColumnNotExists` regardless of stale declarations. + if !has_stable_string_form(&col_to_change.target_type) + && let Some(key) = entity_option_referencing(&self.options, change_column_name) + { + return error::InvalidAlterRequestSnafu { + table: table_name, + err: format!( + "column `{change_column_name}` is referenced by entity option \ + `{key}` and must keep a type that renders as a string, got `{}`", + col_to_change.target_type + ), + } + .fail(); + } + let column = &table_schema.column_schemas()[index]; ensure!( @@ -1408,6 +1474,14 @@ impl TableInfo { } } +fn entity_option_referencing<'a>(options: &'a TableOptions, column: &str) -> Option<&'a str> { + options.extra_options.iter().find_map(|(key, value)| { + (parse_entity_option_key(key).is_some() + && parse_entity_columns(value).iter().any(|c| c == column)) + .then_some(key.as_str()) + }) +} + /// Set column fulltext options if it passed the validation. /// /// Options allowed to modify: @@ -1738,8 +1812,12 @@ mod tests { .build() .unwrap(); - let alter_kind = AlterKind::SetRepartitionColumnHint { - column_name: " col1 ".to_string(), + let alter_kind = AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: vec![( + REPARTITION_COLUMN_HINT_KEY.to_string(), + " col1 ".to_string(), + )], }; let new_meta = meta .builder_with_alter_kind("my_table", &alter_kind) @@ -1767,8 +1845,9 @@ mod tests { .build() .unwrap(); - let alter_kind = AlterKind::SetRepartitionColumnHint { - column_name: " ".to_string(), + let alter_kind = AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: vec![(REPARTITION_COLUMN_HINT_KEY.to_string(), " ".to_string())], }; let err = meta .builder_with_alter_kind("my_table", &alter_kind) @@ -1791,8 +1870,12 @@ mod tests { .build() .unwrap(); - let alter_kind = AlterKind::SetRepartitionColumnHint { - column_name: "col1,col2".to_string(), + let alter_kind = AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: vec![( + REPARTITION_COLUMN_HINT_KEY.to_string(), + "col1,col2".to_string(), + )], }; let err = meta .builder_with_alter_kind("my_table", &alter_kind) @@ -1815,8 +1898,12 @@ mod tests { .build() .unwrap(); - let alter_kind = AlterKind::SetRepartitionColumnHint { - column_name: "missing".to_string(), + let alter_kind = AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: vec![( + REPARTITION_COLUMN_HINT_KEY.to_string(), + "missing".to_string(), + )], }; let err = meta .builder_with_alter_kind("my_table", &alter_kind) @@ -1836,8 +1923,9 @@ mod tests { .build() .unwrap(); - let alter_kind = AlterKind::SetRepartitionColumnHint { - column_name: "ts".to_string(), + let alter_kind = AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: vec![(REPARTITION_COLUMN_HINT_KEY.to_string(), "ts".to_string())], }; let err = meta .builder_with_alter_kind("my_table", &alter_kind) @@ -1861,8 +1949,9 @@ mod tests { .build() .unwrap(); - let alter_kind = AlterKind::SetRepartitionColumnHint { - column_name: "col1".to_string(), + let alter_kind = AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: vec![(REPARTITION_COLUMN_HINT_KEY.to_string(), "col1".to_string())], }; let err = meta .builder_with_alter_kind("my_table", &alter_kind) @@ -1891,7 +1980,13 @@ mod tests { .unwrap(); let new_meta = meta - .builder_with_alter_kind("my_table", &AlterKind::UnsetRepartitionColumnHint) + .builder_with_alter_kind( + "my_table", + &AlterKind::UnsetAnnotations { + family: AnnotationFamily::RepartitionHint, + keys: vec![REPARTITION_COLUMN_HINT_KEY.to_string()], + }, + ) .unwrap() .build() .unwrap(); @@ -2627,4 +2722,305 @@ mod tests { }; assert_eq!(actual, expected); } + + use crate::requests::{SEMANTIC_METRIC_UNIT, SEMANTIC_SIGNAL_TYPE}; + + /// `host` is a tag, `service` and `note` are string fields, `payload` is a + /// binary field. + fn semantic_test_meta() -> TableMeta { + let column_schemas = vec![ + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new("payload", ConcreteDataType::binary_datatype(), true), + ColumnSchema::new("service", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("note", ConcreteDataType::string_datatype(), true), + ]; + let schema = Arc::new( + SchemaBuilder::try_from(column_schemas) + .unwrap() + .build() + .unwrap(), + ); + TableMetaBuilder::empty() + .schema(schema) + .primary_key_indices(vec![0]) + .engine("engine") + .next_column_id(5) + .build() + .unwrap() + } + + #[test] + fn test_set_semantic_annotations() { + let meta = semantic_test_meta(); + // `service` is a plain field: entity columns may be tags or fields. + let alter_kind = AlterKind::SetAnnotations { + family: AnnotationFamily::Semantic, + options: vec![ + (SEMANTIC_SIGNAL_TYPE.to_string(), "trace".to_string()), + ( + "greptime.semantic.entity.service.id".to_string(), + "service".to_string(), + ), + ], + }; + let new_meta = meta + .builder_with_alter_kind("my_table", &alter_kind) + .unwrap() + .build() + .unwrap(); + assert_eq!( + new_meta.options.extra_options.get(SEMANTIC_SIGNAL_TYPE), + Some(&"trace".to_string()) + ); + assert_eq!( + new_meta + .options + .extra_options + .get("greptime.semantic.entity.service.id"), + Some(&"service".to_string()) + ); + } + + #[test] + fn test_repartition_hint_batch_must_be_single_key() { + let meta = TableMetaBuilder::empty() + .schema(Arc::new(new_test_schema())) + .primary_key_indices(vec![0]) + .engine("engine") + .next_column_id(3) + .build() + .unwrap(); + + // Direct gRPC can hand the mutation layer a duplicated batch the + // converter never saw; last-write-wins must not silently apply. + let dup_set = AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: vec![ + (REPARTITION_COLUMN_HINT_KEY.to_string(), "col1".to_string()), + (REPARTITION_COLUMN_HINT_KEY.to_string(), "col2".to_string()), + ], + }; + let err = meta + .builder_with_alter_kind("my_table", &dup_set) + .err() + .unwrap(); + assert!( + err.to_string().contains("must be altered separately"), + "{err}" + ); + + let dup_unset = AlterKind::UnsetAnnotations { + family: AnnotationFamily::RepartitionHint, + keys: vec![ + REPARTITION_COLUMN_HINT_KEY.to_string(), + REPARTITION_COLUMN_HINT_KEY.to_string(), + ], + }; + let err = meta + .builder_with_alter_kind("my_table", &dup_unset) + .err() + .unwrap(); + assert!( + err.to_string().contains("must be altered separately"), + "{err}" + ); + } + + #[test] + fn test_set_semantic_annotations_rejects_invalid() { + let meta = semantic_test_meta(); + let cases = [ + ( + "greptime.semantic.unknown_key", + "x", + "unknown semantic option", + ), + (SEMANTIC_SIGNAL_TYPE, "garbage", "invalid value"), + ( + "greptime.semantic.entity.host.id", + "no_such_column", + "no_such_column", + ), + ( + "greptime.semantic.entity.host.id", + "payload", + "cannot render as a string", + ), + ]; + for (key, value, needle) in cases { + let alter_kind = AlterKind::SetAnnotations { + family: AnnotationFamily::Semantic, + options: vec![(key.to_string(), value.to_string())], + }; + let err = meta + .builder_with_alter_kind("my_table", &alter_kind) + .err() + .unwrap(); + assert!( + err.to_string().contains(needle), + "key `{key}`: unexpected error `{err}`" + ); + } + + // ALTER keeps missing columns on the 4002 contract (pinned by sqlness); + // the shared validator must not collapse it into InvalidArguments. + let missing = meta + .builder_with_alter_kind( + "my_table", + &AlterKind::SetAnnotations { + family: AnnotationFamily::Semantic, + options: vec![( + "greptime.semantic.entity.host.id".to_string(), + "no_such_column".to_string(), + )], + }, + ) + .err() + .unwrap(); + assert_eq!( + common_error::status_code::StatusCode::TableColumnNotFound, + common_error::ext::ErrorExt::status_code(&missing) + ); + } + + #[test] + fn test_unset_semantic_annotations() { + let mut meta = semantic_test_meta(); + meta.options + .extra_options + .insert(SEMANTIC_SIGNAL_TYPE.to_string(), "trace".to_string()); + // A key this version does not recognise, still inside the namespace. + meta.options + .extra_options + .insert("greptime.semantic.future_key".to_string(), "x".to_string()); + + let alter_kind = AlterKind::UnsetAnnotations { + family: AnnotationFamily::Semantic, + keys: vec![ + SEMANTIC_SIGNAL_TYPE.to_string(), + "greptime.semantic.future_key".to_string(), + // Absent key: removal is a no-op, not an error. + SEMANTIC_METRIC_UNIT.to_string(), + ], + }; + let new_meta = meta + .builder_with_alter_kind("my_table", &alter_kind) + .unwrap() + .build() + .unwrap(); + assert!( + !new_meta + .options + .extra_options + .contains_key(SEMANTIC_SIGNAL_TYPE) + ); + assert!( + !new_meta + .options + .extra_options + .contains_key("greptime.semantic.future_key") + ); + + let outside = AlterKind::UnsetAnnotations { + family: AnnotationFamily::Semantic, + keys: vec!["ttl".to_string()], + }; + let err = meta + .builder_with_alter_kind("my_table", &outside) + .err() + .unwrap(); + assert!(err.to_string().contains("annotation namespace"), "{err}"); + } + + #[test] + fn test_modify_entity_column_type_keeps_string_form() { + let mut meta = semantic_test_meta(); + meta.options.extra_options.insert( + "greptime.semantic.entity.service.id".to_string(), + "service".to_string(), + ); + + let alter_kind = AlterKind::ModifyColumnTypes { + columns: vec![ModifyColumnTypeRequest { + column_name: "service".to_string(), + target_type: ConcreteDataType::binary_datatype(), + }], + }; + let err = meta + .builder_with_alter_kind("my_table", &alter_kind) + .err() + .unwrap(); + assert!( + err.to_string() + .contains("must keep a type that renders as a string"), + "{err}" + ); + + // An unreferenced column may still change to a non-string form. + let alter_kind = AlterKind::ModifyColumnTypes { + columns: vec![ModifyColumnTypeRequest { + column_name: "note".to_string(), + target_type: ConcreteDataType::binary_datatype(), + }], + }; + meta.builder_with_alter_kind("my_table", &alter_kind) + .unwrap(); + } + + #[test] + fn test_stale_entity_declaration_guards_readd_and_reports_missing_column() { + // A stale declaration: `gone` was dropped after being declared. + let mut meta = semantic_test_meta(); + meta.options.extra_options.insert( + "greptime.semantic.entity.service.id".to_string(), + "gone".to_string(), + ); + + // MODIFY on the missing column keeps the ColumnNotExists contract + // (4002), stale declaration or not. + let modify = AlterKind::ModifyColumnTypes { + columns: vec![ModifyColumnTypeRequest { + column_name: "gone".to_string(), + target_type: ConcreteDataType::binary_datatype(), + }], + }; + let err = meta + .builder_with_alter_kind("my_table", &modify) + .err() + .unwrap(); + assert_eq!( + common_error::status_code::StatusCode::TableColumnNotFound, + common_error::ext::ErrorExt::status_code(&err) + ); + + // Re-adding the declared column must not hand the declaration a + // non-string type... + let add = |ty: ConcreteDataType| AlterKind::AddColumns { + columns: vec![AddColumnRequest { + column_schema: ColumnSchema::new("gone", ty, true), + is_key: false, + location: None, + add_if_not_exists: false, + }], + }; + let err = meta + .builder_with_alter_kind("my_table", &add(ConcreteDataType::binary_datatype())) + .err() + .unwrap(); + assert!( + err.to_string() + .contains("must keep a type that renders as a string"), + "{err}" + ); + + // ...while a string re-add re-satisfies it. + meta.builder_with_alter_kind("my_table", &add(ConcreteDataType::string_datatype())) + .unwrap(); + } } diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index b4a7d5e052b..aea45130ed9 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -27,7 +27,7 @@ use common_time::range::TimestampRange; use datatypes::data_type::ConcreteDataType; use datatypes::prelude::VectorRef; use datatypes::schema::{ - ColumnDefaultConstraint, ColumnSchema, FulltextOptions, SkippingIndexOptions, + ColumnDefaultConstraint, ColumnSchema, FulltextOptions, Schema, SkippingIndexOptions, }; use greptime_proto::v1::region::compact_request; use once_cell::sync::Lazy; @@ -301,6 +301,211 @@ pub struct ModifyColumnTypeRequest { pub target_type: ConcreteDataType, } +/// A family of annotation table options: pure metadata markers that no region +/// consumes. Setting or unsetting them only rewrites the table's +/// `extra_options`, so the alter skips region dispatch entirely. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum AnnotationFamily { + /// `greptime.semantic.*` options (see the [`semantic`] module). + Semantic, + /// `repartition.column.hint`, consumed by the auto-repartition planner. + RepartitionHint, +} + +impl AnnotationFamily { + /// The key namespace: a prefix for [`Self::Semantic`], the exact key for + /// [`Self::RepartitionHint`]. + pub fn namespace(self) -> &'static str { + match self { + Self::Semantic => SEMANTIC_PREFIX, + Self::RepartitionHint => REPARTITION_COLUMN_HINT_KEY, + } + } + + pub fn of_key(key: &str) -> Option { + if key.starts_with(SEMANTIC_PREFIX) { + Some(Self::Semantic) + } else if key == REPARTITION_COLUMN_HINT_KEY { + Some(Self::RepartitionHint) + } else { + None + } + } + + /// Whether this family may be altered on logical metric tables. Only + /// families whose values nothing on the physical side consumes qualify; + /// the repartition hint drives physical region repartitioning. + pub fn allows_logical_tables(self) -> bool { + match self { + Self::Semantic => true, + Self::RepartitionHint => false, + } + } + + /// Whether this family's SET/UNSET batch must contain exactly one key. + /// The repartition hint is a single marker; a batch with several hint + /// entries (duplicates included) has no meaningful order. + pub fn requires_single_key(self) -> bool { + matches!(self, Self::RepartitionHint) + } + + /// The error for a SET/UNSET batch mixing this family with other options. + pub fn mixed_batch_error(self) -> String { + match self { + Self::Semantic => format!( + "`{SEMANTIC_PREFIX}*` options must be altered separately from other table options" + ), + Self::RepartitionHint => { + format!("{REPARTITION_COLUMN_HINT_KEY} must be altered separately") + } + } + } +} + +/// Table shape an annotation option is validated against. +pub struct AnnotationContext<'a> { + pub schema: &'a Schema, + pub partition_key_indices: &'a [usize], +} + +/// Why an annotation option was rejected. Typed so each DDL entry point maps +/// rules onto its existing error variants and status codes: ALTER keeps +/// missing columns as `TableColumnNotFound` (4002), CREATE keeps its +/// `InvalidArguments` family — the rules converge, the contracts do not. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AnnotationValidationError { + UnknownKey { + key: String, + }, + InvalidValue { + key: String, + value: String, + }, + ColumnNotFound { + column: String, + }, + ColumnNotStringForm { + key: String, + column: String, + ty: ConcreteDataType, + }, + NotSingleColumn, + PartitionMetadataConflict, + TimeIndexConflict, +} + +impl fmt::Display for AnnotationValidationError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::UnknownKey { key } => write!(f, "unknown semantic option `{key}`"), + Self::InvalidValue { key, value } => { + write!(f, "invalid value `{value}` for semantic option `{key}`") + } + Self::ColumnNotFound { column } => write!(f, "column `{column}` not found"), + Self::ColumnNotStringForm { key, column, ty } => write!( + f, + "entity column `{column}` (option `{key}`) has type `{ty}`, \ + which cannot render as a string" + ), + Self::NotSingleColumn => write!( + f, + "{REPARTITION_COLUMN_HINT_KEY} expects exactly one column name" + ), + Self::PartitionMetadataConflict => write!( + f, + "cannot set {REPARTITION_COLUMN_HINT_KEY} on a table with partition metadata" + ), + Self::TimeIndexConflict => write!( + f, + "cannot set {REPARTITION_COLUMN_HINT_KEY} to the time index column" + ), + } + } +} + +/// Validates one annotation option and returns the value to store — the +/// repartition hint is trimmed to the bare column name, semantic values pass +/// through unchanged. +pub(crate) fn validate_and_normalize_annotation( + family: AnnotationFamily, + cx: &AnnotationContext<'_>, + key: &str, + value: &str, +) -> std::result::Result { + match family { + AnnotationFamily::Semantic => { + if !is_semantic_option_key(key) { + return Err(AnnotationValidationError::UnknownKey { + key: key.to_string(), + }); + } + if !validate_semantic_option(key, value) { + return Err(AnnotationValidationError::InvalidValue { + key: key.to_string(), + value: value.to_string(), + }); + } + if parse_entity_option_key(key).is_some() { + for column in parse_entity_columns(value) { + let schema = cx.schema.column_schema_by_name(&column).ok_or_else(|| { + AnnotationValidationError::ColumnNotFound { + column: column.clone(), + } + })?; + if !has_stable_string_form(&schema.data_type) { + return Err(AnnotationValidationError::ColumnNotStringForm { + key: key.to_string(), + column, + ty: schema.data_type.clone(), + }); + } + } + } + Ok(value.to_string()) + } + AnnotationFamily::RepartitionHint => { + let column_name = value.trim(); + if column_name.is_empty() || column_name.contains(',') { + return Err(AnnotationValidationError::NotSingleColumn); + } + if !cx.partition_key_indices.is_empty() { + return Err(AnnotationValidationError::PartitionMetadataConflict); + } + let column_index = cx.schema.column_index_by_name(column_name).ok_or_else(|| { + AnnotationValidationError::ColumnNotFound { + column: column_name.to_string(), + } + })?; + if cx.schema.timestamp_index() == Some(column_index) { + return Err(AnnotationValidationError::TimeIndexConflict); + } + Ok(column_name.to_string()) + } + } +} + +/// CREATE-side entry: validates every annotation option present in `options` +/// and writes normalized values back in place. +pub fn validate_and_normalize_annotation_options( + options: &mut TableOptions, + cx: &AnnotationContext<'_>, +) -> std::result::Result<(), AnnotationValidationError> { + let mut normalized = Vec::new(); + for (key, value) in &options.extra_options { + let Some(family) = AnnotationFamily::of_key(key) else { + continue; + }; + let checked = validate_and_normalize_annotation(family, cx, key, value)?; + if checked != *value { + normalized.push((key.clone(), checked)); + } + } + for (key, value) in normalized { + options.extra_options.insert(key, value); + } + Ok(()) +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub enum AlterKind { AddColumns { @@ -321,10 +526,14 @@ pub enum AlterKind { UnsetTableOptions { keys: Vec, }, - SetRepartitionColumnHint { - column_name: String, + SetAnnotations { + family: AnnotationFamily, + options: Vec<(String, String)>, + }, + UnsetAnnotations { + family: AnnotationFamily, + keys: Vec, }, - UnsetRepartitionColumnHint, SetIndexes { options: Vec, }, diff --git a/src/table/src/requests/semantic.rs b/src/table/src/requests/semantic.rs index 5c0e4c3caed..2daa440c11d 100644 --- a/src/table/src/requests/semantic.rs +++ b/src/table/src/requests/semantic.rs @@ -30,6 +30,8 @@ //! [`crate::requests::validate_table_option`], so they are accepted both on the //! ingestion auto-create path and on explicit `CREATE TABLE ... WITH (...)` DDL. +use datatypes::prelude::ConcreteDataType; + /// Reserved prefix for every public semantic table-option key. pub const SEMANTIC_PREFIX: &str = "greptime.semantic."; @@ -90,8 +92,10 @@ pub const SEMANTIC_ENTITY_PREFIX: &str = "greptime.semantic.entity."; /// is the `service_name` tag column, declaring the logical `service` entity. pub const SEMANTIC_ENTITY_SERVICE_ID: &str = "greptime.semantic.entity.service.id"; -/// The role a set of columns plays for an entity: `id` (identifying attributes, -/// must be tag columns — enforced at DDL time), `descriptive`, or `scope`. +/// The role a set of columns plays for an entity: `id` (identifying +/// attributes), `descriptive`, or `scope`. Columns may be tags or fields; DDL +/// validation only requires that they exist and render as stable strings +/// ([`has_stable_string_form`]). #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum EntityRole { Id, @@ -182,6 +186,24 @@ pub fn is_entity_option_key(key: &str) -> bool { parse_entity_option_key(key).is_some() } +/// Returns true if a column of `data_type` renders as a stable string — the +/// requirement for entity id/descriptive/scope columns. The read-time +/// derivation casts them to strings, so a type without a stable string form +/// would fail only when the graph is scanned; DDL validation rejects it up +/// front instead. +pub fn has_stable_string_form(data_type: &ConcreteDataType) -> bool { + !matches!( + data_type, + ConcreteDataType::Binary(_) + | ConcreteDataType::Json(_) + | ConcreteDataType::Vector(_) + | ConcreteDataType::List(_) + | ConcreteDataType::Struct(_) + | ConcreteDataType::Dictionary(_) + | ConcreteDataType::Null(_) + ) +} + /// Tokenizes an entity option's comma-separated column list (trimmed, empty /// tokens dropped). [`validate_semantic_option`] rejects empty tokens at DDL /// time, so readers only ever drop what validation already refused. diff --git a/tests/cases/standalone/common/information_schema/table_semantics.result b/tests/cases/standalone/common/information_schema/table_semantics.result index 5aa39a91fbf..b57087b3f50 100644 --- a/tests/cases/standalone/common/information_schema/table_semantics.result +++ b/tests/cases/standalone/common/information_schema/table_semantics.result @@ -86,3 +86,121 @@ DROP TABLE plain_table; Affected Rows: 0 +-- ALTER TABLE manages semantic declarations on existing tables: SET appears in +-- the view, UNSET disappears from it. +CREATE TABLE altered_semantics ( + ts TIMESTAMP TIME INDEX, + svc STRING, + payload BINARY, + val DOUBLE, +); + +Affected Rows: 0 + +ALTER TABLE altered_semantics SET 'greptime.semantic.signal_type' = 'metric', 'greptime.semantic.entity.service.id' = 'svc'; + +Affected Rows: 0 + +SELECT table_name, signal_type, semantic_options +FROM information_schema.table_semantics +WHERE table_name = 'altered_semantics'; + ++-------------------+-------------+-----------------------------+ +| table_name | signal_type | semantic_options | ++-------------------+-------------+-----------------------------+ +| altered_semantics | metric | {"entity.service.id":"svc"} | ++-------------------+-------------+-----------------------------+ + +-- Semantic options never share a statement with regular table options. +ALTER TABLE altered_semantics SET 'greptime.semantic.source' = 'prometheus', 'ttl' = '7d'; + +Error: 1004(InvalidArguments), Invalid table option request: `greptime.semantic.*` options must be altered separately from other table options + +-- Unknown semantic keys are rejected on SET. +ALTER TABLE altered_semantics SET 'greptime.semantic.nonsense' = 'x'; + +Error: 1004(InvalidArguments), Invalid alter table(altered_semantics) request: unknown semantic option `greptime.semantic.nonsense` + +-- Values outside the key's domain are rejected. +ALTER TABLE altered_semantics SET 'greptime.semantic.signal_type' = 'garbage'; + +Error: 1004(InvalidArguments), Invalid alter table(altered_semantics) request: invalid value `garbage` for semantic option `greptime.semantic.signal_type` + +-- Entity columns must exist. +ALTER TABLE altered_semantics SET 'greptime.semantic.entity.host.id' = 'no_such_column'; + +Error: 4002(TableColumnNotFound), Column no_such_column not exists in table altered_semantics + +-- Entity columns must render as strings. +ALTER TABLE altered_semantics SET 'greptime.semantic.entity.host.id' = 'payload'; + +Error: 1004(InvalidArguments), Invalid alter table(altered_semantics) request: entity column `payload` (option `greptime.semantic.entity.host.id`) has type `Binary`, which cannot render as a string + +ALTER TABLE altered_semantics UNSET 'greptime.semantic.entity.service.id'; + +Affected Rows: 0 + +SELECT table_name, signal_type, semantic_options +FROM information_schema.table_semantics +WHERE table_name = 'altered_semantics'; + ++-------------------+-------------+------------------+ +| table_name | signal_type | semantic_options | ++-------------------+-------------+------------------+ +| altered_semantics | metric | | ++-------------------+-------------+------------------+ + +DROP TABLE altered_semantics; + +Affected Rows: 0 + +-- Logical metric tables take the same metadata-only path. +CREATE TABLE phy_sem (ts TIMESTAMP TIME INDEX, val DOUBLE) engine=metric with ("physical_metric_table" = ""); + +Affected Rows: 0 + +CREATE TABLE logical_sem (ts TIMESTAMP TIME INDEX, val DOUBLE, host STRING PRIMARY KEY) engine=metric with ("on_physical_table" = "phy_sem"); + +Affected Rows: 0 + +ALTER TABLE logical_sem SET 'greptime.semantic.signal_type' = 'metric', 'greptime.semantic.entity.host.id' = 'host'; + +Affected Rows: 0 + +SELECT table_name, signal_type, semantic_options +FROM information_schema.table_semantics +WHERE table_name = 'logical_sem'; + ++-------------+-------------+---------------------------+ +| table_name | signal_type | semantic_options | ++-------------+-------------+---------------------------+ +| logical_sem | metric | {"entity.host.id":"host"} | ++-------------+-------------+---------------------------+ + +-- Regular options still cannot be altered on logical tables. +ALTER TABLE logical_sem SET 'ttl' = '7d'; + +Error: 1004(InvalidArguments), Alter logical tables invalid arguments: Only support add columns operation + +ALTER TABLE logical_sem UNSET 'greptime.semantic.entity.host.id'; + +Affected Rows: 0 + +SELECT table_name, signal_type, semantic_options +FROM information_schema.table_semantics +WHERE table_name = 'logical_sem'; + ++-------------+-------------+------------------+ +| table_name | signal_type | semantic_options | ++-------------+-------------+------------------+ +| logical_sem | metric | | ++-------------+-------------+------------------+ + +DROP TABLE logical_sem; + +Affected Rows: 0 + +DROP TABLE phy_sem; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/information_schema/table_semantics.sql b/tests/cases/standalone/common/information_schema/table_semantics.sql index 8fb7c78aefe..4acb15294da 100644 --- a/tests/cases/standalone/common/information_schema/table_semantics.sql +++ b/tests/cases/standalone/common/information_schema/table_semantics.sql @@ -45,3 +45,65 @@ DROP TABLE metrics_tagged; DROP TABLE traces_tagged; DROP TABLE plain_table; + +-- ALTER TABLE manages semantic declarations on existing tables: SET appears in +-- the view, UNSET disappears from it. +CREATE TABLE altered_semantics ( + ts TIMESTAMP TIME INDEX, + svc STRING, + payload BINARY, + val DOUBLE, +); + +ALTER TABLE altered_semantics SET 'greptime.semantic.signal_type' = 'metric', 'greptime.semantic.entity.service.id' = 'svc'; + +SELECT table_name, signal_type, semantic_options +FROM information_schema.table_semantics +WHERE table_name = 'altered_semantics'; + +-- Semantic options never share a statement with regular table options. +ALTER TABLE altered_semantics SET 'greptime.semantic.source' = 'prometheus', 'ttl' = '7d'; + +-- Unknown semantic keys are rejected on SET. +ALTER TABLE altered_semantics SET 'greptime.semantic.nonsense' = 'x'; + +-- Values outside the key's domain are rejected. +ALTER TABLE altered_semantics SET 'greptime.semantic.signal_type' = 'garbage'; + +-- Entity columns must exist. +ALTER TABLE altered_semantics SET 'greptime.semantic.entity.host.id' = 'no_such_column'; + +-- Entity columns must render as strings. +ALTER TABLE altered_semantics SET 'greptime.semantic.entity.host.id' = 'payload'; + +ALTER TABLE altered_semantics UNSET 'greptime.semantic.entity.service.id'; + +SELECT table_name, signal_type, semantic_options +FROM information_schema.table_semantics +WHERE table_name = 'altered_semantics'; + +DROP TABLE altered_semantics; + +-- Logical metric tables take the same metadata-only path. +CREATE TABLE phy_sem (ts TIMESTAMP TIME INDEX, val DOUBLE) engine=metric with ("physical_metric_table" = ""); + +CREATE TABLE logical_sem (ts TIMESTAMP TIME INDEX, val DOUBLE, host STRING PRIMARY KEY) engine=metric with ("on_physical_table" = "phy_sem"); + +ALTER TABLE logical_sem SET 'greptime.semantic.signal_type' = 'metric', 'greptime.semantic.entity.host.id' = 'host'; + +SELECT table_name, signal_type, semantic_options +FROM information_schema.table_semantics +WHERE table_name = 'logical_sem'; + +-- Regular options still cannot be altered on logical tables. +ALTER TABLE logical_sem SET 'ttl' = '7d'; + +ALTER TABLE logical_sem UNSET 'greptime.semantic.entity.host.id'; + +SELECT table_name, signal_type, semantic_options +FROM information_schema.table_semantics +WHERE table_name = 'logical_sem'; + +DROP TABLE logical_sem; + +DROP TABLE phy_sem; diff --git a/tests/cases/standalone/common/system/semantic_graph.result b/tests/cases/standalone/common/system/semantic_graph.result index ed6a3d3f16f..6af28cfa2b5 100644 --- a/tests/cases/standalone/common/system/semantic_graph.result +++ b/tests/cases/standalone/common/system/semantic_graph.result @@ -119,6 +119,61 @@ drop table graph_app_metrics; Affected Rows: 0 +-- Declarations can be added after the fact: a table created without semantic +-- options joins the graph once ALTER TABLE declares its entities, and leaves +-- it again on UNSET. +create table graph_late_metrics ( + ts timestamp time index, + svc string, + env string, + latency double +); + +Affected Rows: 0 + +insert into graph_late_metrics values (now(), 'checkout', 'eu-1', 1.5); + +Affected Rows: 1 + +select entity_type, entity_id from greptime_private.semantic_entities order by entity_type, entity_id; + +++ +++ + +alter table graph_late_metrics set 'greptime.semantic.entity.service.id' = 'svc', 'greptime.semantic.entity.service.scope' = 'env'; + +Affected Rows: 0 + +select entity_type, entity_id, scope from greptime_private.semantic_entities order by entity_type, entity_id; + ++-------------+-----------+-------+ +| entity_type | entity_id | scope | ++-------------+-----------+-------+ +| service | checkout | eu-1 | ++-------------+-----------+-------+ + +-- A column referenced by a declaration must keep a string-renderable type. +alter table graph_late_metrics modify column svc binary; + +Error: 1004(InvalidArguments), Invalid alter table(graph_late_metrics) request: column `svc` is referenced by entity option `greptime.semantic.entity.service.id` and must keep a type that renders as a string, got `Binary` + +alter table graph_late_metrics unset 'greptime.semantic.entity.service.id'; + +Affected Rows: 0 + +alter table graph_late_metrics unset 'greptime.semantic.entity.service.scope'; + +Affected Rows: 0 + +select entity_type, entity_id from greptime_private.semantic_entities order by entity_type, entity_id; + +++ +++ + +drop table graph_late_metrics; + +Affected Rows: 0 + -- Calls derivation over trace-v1 tables, all branches in one scan: the client -- span and its child server span land in different tables and still pair into -- one edge with RED metrics; the epoch-timestamped pair falls outside the diff --git a/tests/cases/standalone/common/system/semantic_graph.sql b/tests/cases/standalone/common/system/semantic_graph.sql index 181a8da66a9..bfc9ea8fd7b 100644 --- a/tests/cases/standalone/common/system/semantic_graph.sql +++ b/tests/cases/standalone/common/system/semantic_graph.sql @@ -66,6 +66,35 @@ order by rel_type, src_id; drop table graph_app_metrics; +-- Declarations can be added after the fact: a table created without semantic +-- options joins the graph once ALTER TABLE declares its entities, and leaves +-- it again on UNSET. +create table graph_late_metrics ( + ts timestamp time index, + svc string, + env string, + latency double +); + +insert into graph_late_metrics values (now(), 'checkout', 'eu-1', 1.5); + +select entity_type, entity_id from greptime_private.semantic_entities order by entity_type, entity_id; + +alter table graph_late_metrics set 'greptime.semantic.entity.service.id' = 'svc', 'greptime.semantic.entity.service.scope' = 'env'; + +select entity_type, entity_id, scope from greptime_private.semantic_entities order by entity_type, entity_id; + +-- A column referenced by a declaration must keep a string-renderable type. +alter table graph_late_metrics modify column svc binary; + +alter table graph_late_metrics unset 'greptime.semantic.entity.service.id'; + +alter table graph_late_metrics unset 'greptime.semantic.entity.service.scope'; + +select entity_type, entity_id from greptime_private.semantic_entities order by entity_type, entity_id; + +drop table graph_late_metrics; + -- Calls derivation over trace-v1 tables, all branches in one scan: the client -- span and its child server span land in different tables and still pair into -- one edge with RED metrics; the epoch-timestamped pair falls outside the From 10f587bc30071f230c979da6c0d3ebbf1d736698 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Mon, 17 Aug 2026 10:19:08 +0000 Subject: [PATCH 58/92] fix(query): preserve timestamp literal semantics in inserts (#8889) * fix(query): follow timestamp insert assignment lineage Signed-off-by: Dennis Zhuang * fix(query): fold constant insert timestamp literals at the assignment Following lineage by retyping the source column changed every output column that reads it: a string column sharing the literal was silently rewritten to a formatted timestamp, and a nanosecond column was truncated to the precision of whichever column was converted first. Resolve the constant read-only and fold it into the assignment expression instead, which leaves the source query untouched and also covers literals behind WHERE, ORDER BY and DISTINCT. VALUES rows and UNION branches carry per-row values, so they keep the in-place rewrite, now guarded against columns with more than one consumer. Signed-off-by: Dennis Zhuang * test(query): strengthen insert lineage regression Signed-off-by: Dennis Zhuang * test(query): trim redundant insert coverage Signed-off-by: Dennis Zhuang * fix(query): rebuild insert unions loosely and cover UNION distinct Per review: rebuilding a rewritten union with the strict constructor rejected legal pre-coercion plans whose untouched columns still differ across branches. Use try_new_with_loose_types, matching the SQL planner. Distinct::All joins the rewrite passthrough so UNION (distinct) literals get session-timezone parsing like UNION ALL; deduplication then keys on parsed instants instead of raw strings. The top-level rewrite path gains the same single-consumer guard as rewrite_projection for hand-built DML plans that share a source column between targets. Signed-off-by: Dennis Zhuang * style(query): tighten insert assignment comments Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/query/src/optimizer/type_conversion.rs | 160 +-------- .../type_conversion/insert_assignment.rs | 336 ++++++++++++++++++ src/query/src/planner.rs | 203 ++++++++++- .../insert/insert_default_timezone.result | 53 ++- .../common/insert/insert_default_timezone.sql | 30 +- 5 files changed, 618 insertions(+), 164 deletions(-) create mode 100644 src/query/src/optimizer/type_conversion/insert_assignment.rs diff --git a/src/query/src/optimizer/type_conversion.rs b/src/query/src/optimizer/type_conversion.rs index 366a7d8a31d..3941fedfe8a 100644 --- a/src/query/src/optimizer/type_conversion.rs +++ b/src/query/src/optimizer/type_conversion.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod insert_assignment; + use std::sync::Arc; use common_time::Timezone; @@ -19,7 +21,7 @@ use common_time::timestamp::{TimeUnit, Timestamp}; use datafusion::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; use datafusion_common::{DFSchemaRef, DataFusionError, Result, ScalarValue}; -use datafusion_expr::expr::{Cast, InList}; +use datafusion_expr::expr::InList; use datafusion_expr::{ Between, BinaryExpr, Expr, ExprSchemable, Filter, LogicalPlan, Operator, TableScan, WriteOp, }; @@ -29,6 +31,7 @@ use session::context::QueryContextRef; use crate::QueryEngineContext; use crate::optimizer::ExtensionAnalyzerRule; +use crate::optimizer::type_conversion::insert_assignment::rewrite_insert_assignments; use crate::plan::ExtractExpr; /// TypeConversionRule converts some literal values in logical plan to other types according @@ -307,120 +310,6 @@ impl TreeNodeRewriter for TypeConverter { } } -/// Rewrites timestamp assignment casts at the INSERT boundary. The source query -/// is intentionally left untouched so explicit casts keep DataFusion semantics. -fn rewrite_insert_assignments( - plan: LogicalPlan, - query_ctx: QueryContextRef, -) -> Result { - let LogicalPlan::Projection(projection) = plan else { - return Ok(plan); - }; - - // VALUES type coercion is performed inside the Values node. Only rewrite - // the direct source of the assignment projection, not nested query plans. - let input = rewrite_insert_values(projection.input.as_ref().clone(), query_ctx.clone())?; - let input_literals = projected_literals(&input); - let mut converter = InsertAssignmentConverter { - query_ctx, - input_literals, - }; - let plan = LogicalPlan::Projection(projection); - let expr = plan - .expressions_consider_join() - .into_iter() - .map(|expr| expr.rewrite(&mut converter).map(|x| x.data)) - .collect::>>()?; - - plan.with_new_exprs(expr, vec![input]) -} - -fn rewrite_insert_values(plan: LogicalPlan, query_ctx: QueryContextRef) -> Result { - let LogicalPlan::Values(mut values) = plan else { - return Ok(plan); - }; - let mut converter = InsertAssignmentConverter { - query_ctx, - input_literals: None, - }; - values.values = values - .values - .into_iter() - .map(|row| { - row.into_iter() - .map(|expr| expr.rewrite(&mut converter).map(|x| x.data)) - .collect::>>() - }) - .collect::>>()?; - - Ok(LogicalPlan::Values(values)) -} - -fn projected_literals(plan: &LogicalPlan) -> Option<(DFSchemaRef, Vec>)> { - let LogicalPlan::Projection(projection) = plan else { - return None; - }; - let literals = projection.expr.iter().map(projected_literal).collect(); - Some((projection.schema.clone(), literals)) -} - -fn projected_literal(expr: &Expr) -> Option { - match expr { - Expr::Literal(value, _) => Some(value.clone()), - Expr::Alias(alias) => projected_literal(&alias.expr), - _ => None, - } -} - -struct InsertAssignmentConverter { - query_ctx: QueryContextRef, - input_literals: Option<(DFSchemaRef, Vec>)>, -} - -impl InsertAssignmentConverter { - fn assignment_literal(&self, expr: &Expr) -> Option { - let Expr::Column(column) = expr else { - return expr.as_literal().cloned(); - }; - let (schema, literals) = self.input_literals.as_ref()?; - let index = schema.maybe_index_of_column(column)?; - literals.get(index)?.clone() - } -} - -impl TreeNodeRewriter for InsertAssignmentConverter { - type Node = Expr; - - fn f_up(&mut self, expr: Expr) -> Result> { - let Expr::Cast(Cast { - expr: inner, - data_type: target_type, - }) = expr - else { - return Ok(Transformed::no(expr)); - }; - let original = || { - Expr::Cast(Cast { - expr: inner.clone(), - data_type: target_type.clone(), - }) - }; - if !matches!(&target_type, DataType::Timestamp(_, _)) { - return Ok(Transformed::no(original())); - } - let Some(ScalarValue::Utf8(Some(value))) = self.assignment_literal(inner.as_ref()) else { - return Ok(Transformed::no(original())); - }; - - match cast_string_to_timestamp(&value, &target_type, Some(&self.query_ctx.timezone())) { - Ok(value) if !value.is_null() => Ok(Transformed::yes(Expr::Literal(value, None))), - // Preserve DataFusion's existing parser for unsupported timestamp - // syntax and out-of-range values. - Ok(_) | Err(_) => Ok(Transformed::no(original())), - } - } -} - fn timestamp_to_timestamp_ms_expr(val: i64, unit: TimeUnit) -> Expr { let timestamp = match unit { TimeUnit::Second => val * 1_000, @@ -472,7 +361,7 @@ mod tests { use datafusion_common::arrow::datatypes::Field; use datafusion_common::{Column, DFSchema, NullEquality}; - use datafusion_expr::expr::Exists; + use datafusion_expr::expr::{Cast, Exists}; use datafusion_expr::{Join, JoinConstraint, JoinType, Literal, LogicalPlanBuilder, Subquery}; use datafusion_sql::TableReference; use session::context::QueryContext; @@ -594,45 +483,6 @@ mod tests { ); } - #[test] - fn test_insert_assignment_uses_query_timezone_and_target_precision() { - use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit; - - let query_ctx = QueryContext::arc(); - query_ctx.set_timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()); - let mut converter = InsertAssignmentConverter { - query_ctx, - input_literals: None, - }; - - let expr = Expr::Cast(Cast { - expr: Box::new("2009-02-13 23:31:30.123456789".lit()), - data_type: DataType::Timestamp(ArrowTimeUnit::Nanosecond, None), - }); - assert_eq!( - converter.f_up(expr).unwrap().data, - ScalarValue::TimestampNanosecond(Some(1_234_539_090_123_456_789), None).lit() - ); - } - - #[test] - fn test_insert_assignment_falls_back_for_unsupported_literal() { - use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit; - - let mut converter = InsertAssignmentConverter { - query_ctx: QueryContext::arc(), - input_literals: None, - }; - for literal in ["1970-01-01", "-8-01-01 00:00:01.5"] { - let expr = Expr::Cast(Cast { - expr: Box::new(literal.lit()), - data_type: DataType::Timestamp(ArrowTimeUnit::Nanosecond, None), - }); - - assert_eq!(converter.f_up(expr.clone()).unwrap().data, expr); - } - } - #[test] fn test_type_converter_leaves_explicit_timestamp_cast_unchanged() { use datafusion_common::arrow::datatypes::TimeUnit as ArrowTimeUnit; diff --git a/src/query/src/optimizer/type_conversion/insert_assignment.rs b/src/query/src/optimizer/type_conversion/insert_assignment.rs new file mode 100644 index 00000000000..11ab7e7829e --- /dev/null +++ b/src/query/src/optimizer/type_conversion/insert_assignment.rs @@ -0,0 +1,336 @@ +// Copyright 2023 Greptime Team +// +// 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. + +use std::sync::Arc; + +use datafusion_common::{Column, Result, ScalarValue}; +use datafusion_expr::expr::{Alias, Cast}; +use datafusion_expr::{Distinct, Expr, ExprSchemable, LogicalPlan, Projection, Union, Values}; +use datatypes::arrow::datatypes::DataType; +use session::context::QueryContextRef; + +use crate::optimizer::type_conversion::cast_string_to_timestamp; +use crate::plan::ExtractExpr; + +/// Rewrites string literals that feed timestamp columns at an INSERT boundary. +/// Constants are folded at the assignment to avoid changing source types; +/// `VALUES` and `UNION` inputs are rewritten per row or branch. Explicit casts +/// stay on DataFusion's existing path. +pub(super) fn rewrite_insert_assignments( + plan: LogicalPlan, + query_ctx: QueryContextRef, +) -> Result { + let LogicalPlan::Projection(assignment) = plan else { + return Ok(plan); + }; + + let converter = InsertAssignmentConverter { query_ctx }; + let mut exprs = assignment.expr.clone(); + let mut input = assignment.input.as_ref().clone(); + let mut changed = false; + for (output_idx, expr) in assignment.expr.iter().enumerate() { + let target_type = assignment.schema.field(output_idx).data_type(); + if !matches!(target_type, DataType::Timestamp(_, _)) { + continue; + } + + let Some(column) = assignment_input_column(expr) else { + continue; + }; + let Some(input_idx) = input.schema().maybe_index_of_column(column) else { + continue; + }; + + let literal = lineage_literal(&input, input_idx).cloned(); + if let Some(literal) = literal + && let Some(folded) = converter.convert_literal(&literal, target_type) + { + let (qualifier, field) = assignment.schema.qualified_field(output_idx); + exprs[output_idx] = folded.alias_qualified(qualifier.cloned(), field.name()); + changed = true; + continue; + } + + // A hand-built DML plan (e.g. from flow) may share one source column + // between targets; rewriting it in place would retype every reader. + if assignment + .expr + .iter() + .enumerate() + .any(|(idx, other)| idx != output_idx && other.column_refs().contains(column)) + { + continue; + } + if let Some(rewritten) = converter.rewrite_output_column(&input, input_idx, target_type)? { + input = rewritten; + changed = true; + } + } + + if !changed { + return Ok(LogicalPlan::Projection(assignment)); + } + Projection::try_new(exprs, Arc::new(input)).map(LogicalPlan::Projection) +} + +/// Resolves a literal when every row carries the same value at `output_idx`. +fn lineage_literal(plan: &LogicalPlan, output_idx: usize) -> Option<&ScalarValue> { + if output_idx >= plan.schema().fields().len() { + return None; + } + + match plan { + LogicalPlan::Projection(projection) => match unalias(&projection.expr[output_idx]) { + Expr::Literal(value, _) => Some(value), + Expr::Column(column) => { + let input_idx = projection.input.schema().maybe_index_of_column(column)?; + lineage_literal(projection.input.as_ref(), input_idx) + } + _ => None, + }, + // These nodes drop, reorder or deduplicate rows without touching the + // value a surviving row carries, and their schema stays positional. + LogicalPlan::Filter(_) + | LogicalPlan::Sort(_) + | LogicalPlan::Limit(_) + | LogicalPlan::SubqueryAlias(_) + | LogicalPlan::Distinct(Distinct::All(_)) => { + let inputs = plan.inputs(); + let [input] = inputs.as_slice() else { + return None; + }; + lineage_literal(input, output_idx) + } + _ => None, + } +} + +struct InsertAssignmentConverter { + query_ctx: QueryContextRef, +} + +impl InsertAssignmentConverter { + fn rewrite_output_column( + &self, + plan: &LogicalPlan, + output_idx: usize, + target_type: &DataType, + ) -> Result> { + if output_idx >= plan.schema().fields().len() { + return Ok(None); + } + + match plan { + // DataFusion pushes INSERT assignment casts into Values, so inspect + // them even though the Values schema already has the target type. + LogicalPlan::Values(values) => self.rewrite_values(values, output_idx, target_type), + // Once the source query has produced the target type, its casts and + // coercions belong to the source query rather than INSERT assignment. + _ if plan.schema().field(output_idx).data_type() == target_type => Ok(None), + LogicalPlan::Projection(projection) => { + self.rewrite_projection(projection, output_idx, target_type) + } + LogicalPlan::Union(union) => self.rewrite_union(union, output_idx, target_type), + // Filter and Sort are excluded here: their predicates and keys read + // the retyped column, which would change the source query. Constants + // still reach the assignment through `lineage_literal`. Distinct::All + // holds no expressions; allowing it shifts dedup keys from raw + // strings to parsed instants. + LogicalPlan::Limit(_) + | LogicalPlan::SubqueryAlias(_) + | LogicalPlan::Distinct(Distinct::All(_)) => { + self.rewrite_passthrough(plan, output_idx, target_type) + } + _ => Ok(None), + } + } + + fn rewrite_projection( + &self, + projection: &Projection, + output_idx: usize, + target_type: &DataType, + ) -> Result> { + match unalias(&projection.expr[output_idx]) { + Expr::Literal(value, _) => { + let Some(converted) = self.convert_literal(value, target_type) else { + return Ok(None); + }; + let (qualifier, field) = projection.schema.qualified_field(output_idx); + let mut exprs = projection.expr.clone(); + exprs[output_idx] = converted.alias_qualified(qualifier.cloned(), field.name()); + + Projection::try_new(exprs, projection.input.clone()) + .map(LogicalPlan::Projection) + .map(Some) + } + Expr::Column(column) => { + // Retyping the input column affects every output column reading + // it, so only follow lineage with a single consumer. + if projection + .expr + .iter() + .enumerate() + .any(|(idx, other)| idx != output_idx && other.column_refs().contains(column)) + { + return Ok(None); + } + + let input = projection.input.as_ref(); + let Some(input_idx) = input.schema().maybe_index_of_column(column) else { + return Ok(None); + }; + let Some(rewritten) = self.rewrite_output_column(input, input_idx, target_type)? + else { + return Ok(None); + }; + + Projection::try_new(projection.expr.clone(), Arc::new(rewritten)) + .map(LogicalPlan::Projection) + .map(Some) + } + _ => Ok(None), + } + } + + fn rewrite_values( + &self, + values: &Values, + output_idx: usize, + target_type: &DataType, + ) -> Result> { + let mut rewritten = values.clone(); + let mut changed = false; + for row in &mut rewritten.values { + let Some(expr) = row.get_mut(output_idx) else { + return Ok(None); + }; + + if let Expr::Cast(Cast { + expr: inner, + data_type, + }) = expr + && data_type == target_type + && let Expr::Literal(value, _) = inner.as_ref() + && let Some(value) = self.convert_literal(value, target_type) + { + *expr = value; + changed = true; + continue; + } + + if expr.get_type(values.schema.as_ref()).ok().as_ref() != Some(target_type) { + return Ok(None); + } + } + + Ok(changed.then_some(LogicalPlan::Values(rewritten))) + } + + fn rewrite_union( + &self, + union: &Union, + output_idx: usize, + target_type: &DataType, + ) -> Result> { + let mut inputs = Vec::with_capacity(union.inputs.len()); + // A union has one shared schema, so rewrite it only when the same output + // column is an INSERT-assigned literal in every branch. + for input in &union.inputs { + let Some(rewritten) = self.rewrite_output_column(input, output_idx, target_type)? + else { + return Ok(None); + }; + inputs.push(Arc::new(rewritten)); + } + + // Untouched columns may still have mismatched branch types before + // TypeCoercion runs, so rebuild loosely like the SQL planner does. + Union::try_new_with_loose_types(inputs) + .map(LogicalPlan::Union) + .map(Some) + } + + fn rewrite_passthrough( + &self, + plan: &LogicalPlan, + output_idx: usize, + target_type: &DataType, + ) -> Result> { + let inputs = plan.inputs(); + let [input] = inputs.as_slice() else { + return Ok(None); + }; + let Some(rewritten_input) = self.rewrite_output_column(input, output_idx, target_type)? + else { + return Ok(None); + }; + + plan.with_new_exprs(plan.expressions_consider_join(), vec![rewritten_input]) + .map(Some) + } + + fn convert_literal(&self, value: &ScalarValue, target_type: &DataType) -> Option { + let ScalarValue::Utf8(Some(value)) = value else { + return None; + }; + cast_string_to_timestamp(value, target_type, Some(&self.query_ctx.timezone())) + .ok() + .filter(|value| !value.is_null()) + .map(|value| Expr::Literal(value, None)) + } +} + +fn assignment_input_column(expr: &Expr) -> Option<&Column> { + let expr = match unalias(expr) { + Expr::Cast(Cast { expr, .. }) => expr.as_ref(), + expr => expr, + }; + let Expr::Column(column) = expr else { + return None; + }; + Some(column) +} + +fn unalias(expr: &Expr) -> &Expr { + match expr { + Expr::Alias(Alias { expr, .. }) => unalias(expr), + expr => expr, + } +} + +#[cfg(test)] +mod tests { + use datafusion_common::ScalarValue; + use datafusion_common::arrow::datatypes::TimeUnit; + use session::context::QueryContext; + + use super::*; + + #[test] + fn test_convert_literal_falls_back_for_unsupported_literal() { + let converter = InsertAssignmentConverter { + query_ctx: QueryContext::arc(), + }; + let target_type = DataType::Timestamp(TimeUnit::Nanosecond, None); + + for literal in ["1970-01-01", "-8-01-01 00:00:01.5"] { + assert_eq!( + converter + .convert_literal(&ScalarValue::Utf8(Some(literal.to_string())), &target_type,), + None + ); + } + } +} diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index c941839c0ad..a7d4a84cb28 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -711,12 +711,18 @@ mod tests { ConcreteDataType::timestamp_millisecond_datatype(), false, ), + ColumnSchema::new("note", ConcreteDataType::string_datatype(), true), + ColumnSchema::new( + "ts_ns", + ConcreteDataType::timestamp_nanosecond_datatype(), + true, + ), ]; let schema = Arc::new(Schema::new(columns)); let table_meta = TableMetaBuilder::empty() .schema(schema) .primary_key_indices(vec![]) - .value_indices(vec![0, 1]) + .value_indices(vec![0, 1, 2, 3]) .next_column_id(1024) .build() .unwrap(); @@ -869,16 +875,69 @@ mod tests { ); let engine = create_timestamp_test_engine().await; - for (sql, expected_timestamp) in [ + for (sql, expected_timestamps) in [ ( "INSERT INTO timestamps (ts, st) \ VALUES ('2026-08-02 12:00:00.001', now())", - 1_785_643_200_001_i64, + &[1_785_643_200_001_i64][..], ), ( "INSERT INTO timestamps (ts, st) \ SELECT '2026-08-03 12:00:00.001', now()", - 1_785_729_600_001_i64, + &[1_785_729_600_001_i64][..], + ), + ( + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-04 12:00:00.001', now() LIMIT 1", + &[1_785_816_000_001_i64][..], + ), + ( + "INSERT INTO timestamps (ts, st) \ + SELECT * FROM (\ + SELECT '2026-08-05 12:00:00.001', now()\ + ) AS source", + &[1_785_902_400_001_i64][..], + ), + ( + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-06 12:00:00.001', now() \ + UNION ALL \ + SELECT '2026-08-07 12:00:00.001', now()", + &[1_785_988_800_001_i64, 1_786_075_200_001_i64][..], + ), + ( + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-16 12:00:00.001', now() \ + UNION \ + SELECT '2026-08-17 12:00:00.001', now()", + &[1_786_852_800_001_i64, 1_786_939_200_001_i64][..], + ), + ( + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-18 12:00:00.001', max(st) \ + FROM timestamps GROUP BY note", + &[1_787_025_600_001_i64][..], + ), + ( + "INSERT INTO timestamps (ts, st) \ + SELECT c, s FROM (\ + SELECT '2026-08-13 12:00:00.001' AS c, now() AS s\ + ) AS t WHERE c > '2026-01-01'", + &[1_786_593_600_001_i64][..], + ), + ( + "INSERT INTO timestamps (ts, st) \ + SELECT c, s FROM (\ + SELECT '2026-08-14 12:00:00.001' AS c, now() AS s\ + ) AS t ORDER BY c", + &[1_786_680_000_001_i64][..], + ), + ( + "INSERT INTO timestamps (ts, st) \ + SELECT DISTINCT c, s FROM (\ + SELECT '2026-08-15 12:00:00.001' AS c, now() AS s\ + ) AS t", + &[1_786_766_400_001_i64][..], ), ] { let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); @@ -890,6 +949,112 @@ mod tests { .display_indent() .to_string(); + for expected_timestamp in expected_timestamps { + assert!( + plan.contains(&format!("TimestampMillisecond({expected_timestamp}, None)")), + "{plan}" + ); + } + } + } + + #[tokio::test] + async fn test_insert_explicit_timestamp_cast_keeps_datafusion_semantics() { + let query_ctx = Arc::new( + QueryContextBuilder::default() + .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()) + .build(), + ); + let engine = create_timestamp_test_engine().await; + let sql = "INSERT INTO timestamps (ts, st) \ + VALUES (CAST('2026-08-08 12:00:00.001' AS TIMESTAMP), now())"; + let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); + let plan = engine + .planner() + .plan(&stmt, query_ctx) + .await + .unwrap() + .display_indent() + .to_string(); + + assert!( + plan.contains("arrow_cast(Utf8(\"2026-08-08 12:00:00.001\")"), + "{plan}" + ); + } + + #[tokio::test] + async fn test_insert_converts_source_literal_shared_by_several_columns() { + let query_ctx = Arc::new( + QueryContextBuilder::default() + .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()) + .build(), + ); + let engine = create_timestamp_test_engine().await; + + // Rewriting `c` in place would also retype `note` and truncate `ts_ns`. + for (sql, expected) in [ + ( + "INSERT INTO timestamps (ts, note) SELECT a, b FROM (\ + SELECT c AS a, c AS b FROM (\ + SELECT '2026-08-12 12:00:00.001' AS c\ + ) AS t1\ + ) AS t2", + &[ + "TimestampMillisecond(1786507200001, None)", + "[a:Utf8, b:Utf8]", + ][..], + ), + ( + "INSERT INTO timestamps (ts, ts_ns) SELECT a, b FROM (\ + SELECT c AS a, c AS b FROM (\ + SELECT '2026-08-12 12:00:00.123456789' AS c\ + ) AS t1\ + ) AS t2", + &[ + "TimestampMillisecond(1786507200123, None)", + "TimestampNanosecond(1786507200123456789, None)", + ][..], + ), + ] { + let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); + let plan = engine + .planner() + .plan(&stmt, query_ctx.clone()) + .await + .unwrap() + .display_indent_schema() + .to_string(); + + for expected in expected { + assert!(plan.contains(expected), "{plan}"); + } + } + } + + #[tokio::test] + async fn test_insert_union_tolerates_uncoerced_untouched_column() { + let query_ctx = Arc::new( + QueryContextBuilder::default() + .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()) + .build(), + ); + let engine = create_timestamp_test_engine().await; + // `st` stays Timestamp vs Null across branches until TypeCoercion runs. + let sql = "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-06 12:00:00.001', now() \ + UNION ALL \ + SELECT '2026-08-07 12:00:00.001', NULL"; + let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); + let plan = engine + .planner() + .plan(&stmt, query_ctx) + .await + .unwrap() + .display_indent() + .to_string(); + + for expected_timestamp in [1_785_988_800_001_i64, 1_786_075_200_001_i64] { assert!( plan.contains(&format!("TimestampMillisecond({expected_timestamp}, None)")), "{plan}" @@ -897,6 +1062,36 @@ mod tests { } } + #[tokio::test] + async fn test_insert_mixed_union_keeps_source_coercion() { + let query_ctx = Arc::new( + QueryContextBuilder::default() + .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()) + .build(), + ); + let engine = create_timestamp_test_engine().await; + let sql = "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-10 12:00:00.001', now() \ + UNION ALL \ + SELECT CAST('2026-08-11 12:00:00.001' AS TIMESTAMP), now()"; + let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); + let plan = engine + .planner() + .plan(&stmt, query_ctx) + .await + .unwrap() + .display_indent_schema() + .to_string(); + assert!( + !plan.contains("TimestampMillisecond(1786334400001, None)"), + "{plan}" + ); + assert!( + plan.contains("arrow_cast(Utf8(\"2026-08-11 12:00:00.001\")"), + "{plan}" + ); + } + #[tokio::test] async fn test_extract_placeholder_cast_types_multiple() { let plan = parse_sql_to_plan( diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.result b/tests/cases/standalone/common/insert/insert_default_timezone.result index da8b1261e7a..413eb5d46ae 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.result +++ b/tests/cases/standalone/common/insert/insert_default_timezone.result @@ -68,10 +68,48 @@ INSERT INTO test3 (ts, st) SELECT '2026-08-03 12:00:00.001', now(); Affected Rows: 1 +INSERT INTO test3 (ts, st) SELECT '2026-08-04 12:00:00.001', now() LIMIT 1; + +Affected Rows: 1 + +INSERT INTO test3 (ts, st) +SELECT '2026-08-06 12:00:00.001', now() +UNION ALL +SELECT '2026-08-07 12:00:00.001', now(); + +Affected Rows: 2 + +INSERT INTO test3 (ts, st) VALUES ( + CAST('2026-08-08 12:00:00.001' AS TIMESTAMP), + now() +); + +Affected Rows: 1 + +INSERT INTO test3 (ts, st) +SELECT '2026-08-10 12:00:00.001', now() +UNION ALL +SELECT CAST('2026-08-11 12:00:00.001' AS TIMESTAMP), now(); + +Affected Rows: 2 + +INSERT INTO test3 (ts, st) +SELECT '2026-08-16 12:00:00.001', now() +UNION +SELECT '2026-08-17 12:00:00.001', now(); + +Affected Rows: 2 + +INSERT INTO test3 (ts, ts_ns) SELECT a, b FROM ( + SELECT c AS a, c AS b FROM (SELECT '2026-08-12 12:00:00.123456789' AS c) AS t1 +) AS t2; + +Affected Rows: 1 + INSERT INTO test3 (ts, st, ts_ns) VALUES ( - '2026-08-04 12:00:00.001', + '2026-08-09 12:00:00.001', now(), - '2026-08-04 12:00:00.123456789' + '2026-08-09 12:00:00.123456789' ); Affected Rows: 1 @@ -84,7 +122,16 @@ SELECT ts, ts_ns FROM test3 ORDER BY ts; | 2026-08-01T04:00:00.001 | | | 2026-08-02T04:00:00.001 | | | 2026-08-03T04:00:00.001 | | -| 2026-08-04T04:00:00.001 | 2026-08-04T04:00:00.123456789 | +| 2026-08-04T04:00:00.001 | | +| 2026-08-06T04:00:00.001 | | +| 2026-08-07T04:00:00.001 | | +| 2026-08-08T12:00:00.001 | | +| 2026-08-09T04:00:00.001 | 2026-08-09T04:00:00.123456789 | +| 2026-08-10T12:00:00.001 | | +| 2026-08-11T12:00:00.001 | | +| 2026-08-12T04:00:00.123 | 2026-08-12T04:00:00.123456789 | +| 2026-08-16T04:00:00.001 | | +| 2026-08-17T04:00:00.001 | | +-------------------------+-------------------------------+ SET time_zone = 'UTC'; diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.sql b/tests/cases/standalone/common/insert/insert_default_timezone.sql index fba8c134fa0..7cbef15d6d9 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.sql +++ b/tests/cases/standalone/common/insert/insert_default_timezone.sql @@ -24,10 +24,36 @@ INSERT INTO test3 (ts, st) VALUES ('2026-08-02 12:00:00.001', now()); INSERT INTO test3 (ts, st) SELECT '2026-08-03 12:00:00.001', now(); +INSERT INTO test3 (ts, st) SELECT '2026-08-04 12:00:00.001', now() LIMIT 1; + +INSERT INTO test3 (ts, st) +SELECT '2026-08-06 12:00:00.001', now() +UNION ALL +SELECT '2026-08-07 12:00:00.001', now(); + +INSERT INTO test3 (ts, st) VALUES ( + CAST('2026-08-08 12:00:00.001' AS TIMESTAMP), + now() +); + +INSERT INTO test3 (ts, st) +SELECT '2026-08-10 12:00:00.001', now() +UNION ALL +SELECT CAST('2026-08-11 12:00:00.001' AS TIMESTAMP), now(); + +INSERT INTO test3 (ts, st) +SELECT '2026-08-16 12:00:00.001', now() +UNION +SELECT '2026-08-17 12:00:00.001', now(); + +INSERT INTO test3 (ts, ts_ns) SELECT a, b FROM ( + SELECT c AS a, c AS b FROM (SELECT '2026-08-12 12:00:00.123456789' AS c) AS t1 +) AS t2; + INSERT INTO test3 (ts, st, ts_ns) VALUES ( - '2026-08-04 12:00:00.001', + '2026-08-09 12:00:00.001', now(), - '2026-08-04 12:00:00.123456789' + '2026-08-09 12:00:00.123456789' ); SELECT ts, ts_ns FROM test3 ORDER BY ts; From d7f1233f775d54380318d6ad9ff62504a7cbcff1 Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Mon, 17 Aug 2026 11:24:44 +0000 Subject: [PATCH 59/92] refactor(json2): support JSON2 storage layout settings in DDL (#8895) * refactor(json2): add JSON2 storage layout settings Signed-off-by: luofucong * resolve PR comments Signed-off-by: luofucong --------- Signed-off-by: luofucong --- src/datatypes/src/error.rs | 8 + src/datatypes/src/extension/json.rs | 1 + src/datatypes/src/json.rs | 286 +++++++++++++++--- src/sql/src/error.rs | 11 +- src/sql/src/parsers/create_parser.rs | 4 +- src/sql/src/parsers/create_parser/json.rs | 154 ++++++++-- src/sql/src/statements.rs | 18 +- src/sql/src/statements/create.rs | 144 +++++---- .../json2_storage_layout_settings/case.toml | 8 + .../json2_storage_layout_settings/setup.sql | 11 + .../verify.result | 27 ++ .../json2_storage_layout_settings/verify.sql | 4 + 12 files changed, 532 insertions(+), 144 deletions(-) create mode 100644 tests/compatibility/cases/json2_storage_layout_settings/case.toml create mode 100644 tests/compatibility/cases/json2_storage_layout_settings/setup.sql create mode 100644 tests/compatibility/cases/json2_storage_layout_settings/verify.result create mode 100644 tests/compatibility/cases/json2_storage_layout_settings/verify.sql diff --git a/src/datatypes/src/error.rs b/src/datatypes/src/error.rs index cf76812c215..b34a5ffa9bf 100644 --- a/src/datatypes/src/error.rs +++ b/src/datatypes/src/error.rs @@ -196,6 +196,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Invalid JSON2 layout: {reason}"))] + InvalidJson2Layout { + reason: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Invalid Vector: {}", msg))] InvalidVector { msg: String, @@ -318,6 +325,7 @@ impl ErrorExt for Error { | InvalidTimestampPrecision { .. } | InvalidPrecisionOrScale { .. } | InvalidJson { .. } + | InvalidJson2Layout { .. } | InvalidJsonb { .. } | InvalidVector { .. } | InvalidFulltextOption { .. } diff --git a/src/datatypes/src/extension/json.rs b/src/datatypes/src/extension/json.rs index 26d470641f5..109aa853b21 100644 --- a/src/datatypes/src/extension/json.rs +++ b/src/datatypes/src/extension/json.rs @@ -22,6 +22,7 @@ use arrow_schema::{ArrowError, DataType, Field}; use serde::{Deserialize, Serialize}; use snafu::ResultExt; +pub use crate::json::JSON2_REMAINDER_FIELD_NAME; use crate::json::JsonSettings; const LEGACY_JSON_STRUCTURE_SETTINGS_KEY: &str = "json_structure_settings"; diff --git a/src/datatypes/src/json.rs b/src/datatypes/src/json.rs index 48a9174020e..ec9ca4ac789 100644 --- a/src/datatypes/src/json.rs +++ b/src/datatypes/src/json.rs @@ -29,21 +29,47 @@ use serde_json::{Map, Value as Json}; use snafu::ResultExt; use crate::data_type::ConcreteDataType; -use crate::error::{self, Result, UnsupportedJsonTypeSnafu}; +use crate::error::{self, InvalidJson2LayoutSnafu, Result, UnsupportedJsonTypeSnafu}; use crate::json::value::{JsonValue, JsonVariant, encode_serde_json_as_jsonb}; use crate::schema::ColumnDefaultConstraint; -use crate::types::json_type::JsonNativeType; +use crate::types::json_type::{JsonNativeType, JsonObjectType}; use crate::value::{ListValue, StructValue, Value}; /// Maximum number of JSON container levels represented as nested Arrow types. pub const JSON2_MAX_STRUCTURED_DEPTH: usize = 50; +/// Reserved physical field containing unexpanded JSON2 paths. +pub const JSON2_REMAINDER_FIELD_NAME: &str = "!__remainder__!"; /// JSON2 settings stored in column schema metadata and represented through /// Arrow extension metadata. -#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize)] pub struct JsonSettings { #[serde(default)] - pub type_hints: Vec, + type_hints: Vec, + /// Maximum number of unhinted JSON paths expanded into Arrow fields. + /// + /// `None` preserves the legacy unlimited layout. + #[serde(default, skip_serializing_if = "Option::is_none")] + max_auto_expanded_paths: Option, +} + +#[derive(Deserialize)] +struct JsonSettingsSerde { + #[serde(default)] + type_hints: Vec, + #[serde(default)] + max_auto_expanded_paths: Option, +} + +impl<'de> Deserialize<'de> for JsonSettings { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + let settings = JsonSettingsSerde::deserialize(deserializer)?; + Self::try_new(settings.type_hints, settings.max_auto_expanded_paths) + .map_err(serde::de::Error::custom) + } } /// Declares selected JSON2 subpaths as typed fields. @@ -79,8 +105,28 @@ pub struct JsonContext<'a> { } impl JsonSettings { - pub fn new(type_hints: Vec) -> Self { - Self { type_hints } + /// Creates and validates JSON2 settings. + pub fn try_new( + type_hints: Vec, + max_auto_expanded_paths: Option, + ) -> Result { + validate_type_hints(&type_hints)?; + Ok(Self { + type_hints, + max_auto_expanded_paths, + }) + } + + pub fn type_hints(&self) -> &[JsonTypeHint] { + &self.type_hints + } + + pub fn max_auto_expanded_paths(&self) -> Option { + self.max_auto_expanded_paths + } + + pub fn into_parts(self) -> (Vec, Option) { + (self.type_hints, self.max_auto_expanded_paths) } /// Decode an encoded StructValue back into a serde_json::Value. @@ -102,6 +148,87 @@ impl JsonSettings { } } +fn validate_type_hints(type_hints: &[JsonTypeHint]) -> Result<()> { + let mut object = JsonObjectType::new(); + for hint in type_hints { + if hint.path.len() > JSON2_MAX_STRUCTURED_DEPTH { + return InvalidJson2LayoutSnafu { + reason: format!( + "JSON2 type hint path cannot exceed {JSON2_MAX_STRUCTURED_DEPTH} segments" + ), + } + .fail(); + } + if hint + .path + .first() + .is_some_and(|x| x == JSON2_REMAINDER_FIELD_NAME) + { + return InvalidJson2LayoutSnafu { + reason: format!( + "JSON2 type hint path cannot be rooted at reserved field '{JSON2_REMAINDER_FIELD_NAME}'" + ), + } + .fail(); + } + let data_type = match &hint.data_type { + ConcreteDataType::Boolean(_) + | ConcreteDataType::UInt8(_) + | ConcreteDataType::UInt16(_) + | ConcreteDataType::UInt32(_) + | ConcreteDataType::UInt64(_) + | ConcreteDataType::Int8(_) + | ConcreteDataType::Int16(_) + | ConcreteDataType::Int32(_) + | ConcreteDataType::Int64(_) + | ConcreteDataType::Float32(_) + | ConcreteDataType::Float64(_) + | ConcreteDataType::String(_) => (&hint.data_type).into(), + data_type => { + return InvalidJson2LayoutSnafu { + reason: format!("unsupported JSON2 type hint data type: {data_type}"), + } + .fail(); + } + }; + validate_type_hint(&mut object, &hint.path, data_type)?; + } + Ok(()) +} + +fn validate_type_hint( + object: &mut JsonObjectType, + path: &[String], + data_type: JsonNativeType, +) -> Result<()> { + let Some((name, path)) = path.split_first() else { + return InvalidJson2LayoutSnafu { + reason: "JSON2 type hint path must not be empty".to_string(), + } + .fail(); + }; + if path.is_empty() { + if object.insert(name.clone(), data_type).is_some() { + return InvalidJson2LayoutSnafu { + reason: format!("duplicate JSON2 type hint path '{name}'"), + } + .fail(); + } + return Ok(()); + } + + let child = object + .entry(name.clone()) + .or_insert_with(|| JsonNativeType::Object(JsonObjectType::new())); + let JsonNativeType::Object(child) = child else { + return InvalidJson2LayoutSnafu { + reason: format!("conflicting JSON2 type hint path at '{name}'"), + } + .fail(); + }; + validate_type_hint(child, path, data_type) +} + impl<'a> JsonContext<'a> { fn type_hint(&self) -> Option<&'a JsonTypeHint> { self.settings @@ -483,7 +610,8 @@ mod tests { } #[test] - fn test_json_settings_forward_compatibility() { + fn test_json_settings_forward_compatibility() + -> std::result::Result<(), Box> { let json_str = r#"{ "type_hints": [ { @@ -512,11 +640,10 @@ mod tests { ] }"#; - let deserialized = serde_json::from_str::(json_str).unwrap(); + let deserialized = serde_json::from_str::(json_str)?; - assert_eq!( - deserialized, - JsonSettings::new(vec![ + let expected = JsonSettings::try_new( + vec![ JsonTypeHint { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), @@ -531,13 +658,16 @@ mod tests { default_constraint: None, inverted_index: false, }, - ]) - ); + ], + None, + )?; + assert_eq!(deserialized, expected); + Ok(()) } #[test] - fn test_json_settings_ser_de() { - let settings = JsonSettings::new(vec![ + fn test_json_settings_ser_de() -> std::result::Result<(), Box> { + let type_hints = vec![ JsonTypeHint { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), @@ -552,12 +682,76 @@ mod tests { default_constraint: None, inverted_index: false, }, - ]); + ]; - let serialized = serde_json::to_string(&settings).unwrap(); - let deserialized = serde_json::from_str::(&serialized).unwrap(); + for settings in [ + JsonSettings { + type_hints: type_hints.clone(), + max_auto_expanded_paths: None, + }, + JsonSettings { + type_hints, + max_auto_expanded_paths: Some(1), + }, + ] { + let serialized = serde_json::to_string(&settings)?; + let deserialized = serde_json::from_str::(&serialized)?; + assert_eq!(settings, deserialized); + } + Ok(()) + } - assert_eq!(settings, deserialized); + #[test] + fn test_json_settings_reject_invalid_type_hint_layout() { + for type_hints in [ + json!([{"path": [], "type": {"Int64": {}}, "nullable": true, "inverted_index": false}]), + json!([ + {"path": ["a"], "type": {"Int64": {}}, "nullable": true, "inverted_index": false}, + {"path": ["a", "b"], "type": {"Int64": {}}, "nullable": true, "inverted_index": false} + ]), + json!([{"path": [JSON2_REMAINDER_FIELD_NAME], "type": {"Int64": {}}, "nullable": true, "inverted_index": false}]), + ] { + let settings = json!({"type_hints": type_hints}); + assert!(serde_json::from_value::(settings).is_err()); + } + + let hint = |path, data_type| JsonTypeHint { + path, + data_type, + nullable: true, + default_constraint: None, + inverted_index: false, + }; + assert!( + JsonSettings::try_new( + vec![hint( + vec!["nested".to_string(), JSON2_REMAINDER_FIELD_NAME.to_string()], + ConcreteDataType::string_datatype(), + )], + None, + ) + .is_ok() + ); + assert!( + JsonSettings::try_new( + vec![hint( + vec!["nested".to_string(); JSON2_MAX_STRUCTURED_DEPTH + 1], + ConcreteDataType::string_datatype(), + )], + None, + ) + .is_err() + ); + assert!( + JsonSettings::try_new( + vec![hint( + vec!["date".to_string()], + ConcreteDataType::date_datatype(), + )], + None, + ) + .is_err() + ); } #[test] @@ -702,14 +896,15 @@ mod tests { } #[test] - fn test_encode_json_respects_type_hint() { - let settings = JsonSettings::new(vec![JsonTypeHint { + fn test_encode_json_respects_type_hint() -> Result<()> { + let type_hints = vec![JsonTypeHint { path: vec!["age".to_string()], data_type: ConcreteDataType::int64_datatype(), nullable: false, default_constraint: None, inverted_index: false, - }]); + }]; + let settings = JsonSettings::try_new(type_hints, None)?; let result = settings .encode(json!({ @@ -731,17 +926,19 @@ mod tests { })) .unwrap_err(); assert!(err.to_string().contains("does not match JSON2 type hint")); + Ok(()) } #[test] - fn test_encode_json_respects_unsigned_type_hint() { - let settings = JsonSettings::new(vec![JsonTypeHint { + fn test_encode_json_respects_unsigned_type_hint() -> Result<()> { + let type_hints = vec![JsonTypeHint { path: vec!["count".to_string()], data_type: ConcreteDataType::uint64_datatype(), nullable: false, default_constraint: None, inverted_index: false, - }]); + }]; + let settings = JsonSettings::try_new(type_hints, None)?; let result = settings .encode(json!({ @@ -765,17 +962,19 @@ mod tests { })) .unwrap_err(); assert!(err.to_string().contains("does not match JSON2 type hint")); + Ok(()) } #[test] - fn test_encode_json_fills_missing_type_hint_with_default() { - let settings = JsonSettings::new(vec![JsonTypeHint { + fn test_encode_json_fills_missing_type_hint_with_default() -> Result<()> { + let type_hints = vec![JsonTypeHint { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), nullable: false, default_constraint: Some(ColumnDefaultConstraint::Value(Value::Int64(7))), inverted_index: false, - }]); + }]; + let settings = JsonSettings::try_new(type_hints, None)?; let result = settings .encode(json!({})) @@ -790,17 +989,19 @@ mod tests { panic!("Expected user Struct value"); }; assert_eq!(struct_field_value(user, "age"), &Value::Int64(7)); + Ok(()) } #[test] - fn test_encode_json_fills_missing_nullable_type_hint_with_null() { - let settings = JsonSettings::new(vec![JsonTypeHint { + fn test_encode_json_fills_missing_nullable_type_hint_with_null() -> Result<()> { + let type_hints = vec![JsonTypeHint { path: vec!["user".to_string(), "name".to_string()], data_type: ConcreteDataType::string_datatype(), nullable: true, default_constraint: None, inverted_index: false, - }]); + }]; + let settings = JsonSettings::try_new(type_hints, None)?; let result = settings .encode(json!({ "user": {} })) @@ -815,28 +1016,31 @@ mod tests { panic!("Expected user Struct value"); }; assert_eq!(struct_field_value(user, "name"), &Value::Null); + Ok(()) } #[test] - fn test_encode_json_rejects_missing_non_null_type_hint() { - let settings = JsonSettings::new(vec![JsonTypeHint { + fn test_encode_json_rejects_missing_non_null_type_hint() -> Result<()> { + let type_hints = vec![JsonTypeHint { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), nullable: false, default_constraint: None, inverted_index: false, - }]); + }]; + let settings = JsonSettings::try_new(type_hints, None)?; let err = settings.encode(json!({})).unwrap_err(); assert!( err.to_string() .contains("missing non-null JSON2 type hint path user.age") ); + Ok(()) } #[test] - fn test_encode_json_merges_missing_type_hint_prefix() { - let settings = JsonSettings::new(vec![ + fn test_encode_json_merges_missing_type_hint_prefix() -> Result<()> { + let type_hints = vec![ JsonTypeHint { path: vec!["user".to_string(), "age".to_string()], data_type: ConcreteDataType::int64_datatype(), @@ -853,13 +1057,10 @@ mod tests { ))), inverted_index: false, }, - ]); + ]; + let settings = JsonSettings::try_new(type_hints, None)?; - let result = settings - .encode(json!({})) - .unwrap() - .into_json_inner() - .unwrap(); + let result = settings.encode(json!({}))?.into_json_inner().unwrap(); let Value::Struct(root) = result else { panic!("Expected Struct value"); @@ -872,6 +1073,7 @@ mod tests { struct_field_value(user, "name"), &Value::String("unknown".into()) ); + Ok(()) } #[test] diff --git a/src/sql/src/error.rs b/src/sql/src/error.rs index cc0b9f19506..dc819bda017 100644 --- a/src/sql/src/error.rs +++ b/src/sql/src/error.rs @@ -339,6 +339,13 @@ pub enum Error { #[snafu(implicit)] location: Location, }, + + #[snafu(transparent)] + Datatypes { + #[snafu(implicit)] + location: Location, + source: datatypes::error::Error, + }, } impl ErrorExt for Error { @@ -384,7 +391,9 @@ impl ErrorExt for Error { #[cfg(feature = "enterprise")] InvalidTriggerWebhookOption { .. } => StatusCode::InvalidArguments, - SerializeColumnDefaultConstraint { source, .. } => source.status_code(), + SerializeColumnDefaultConstraint { source, .. } | Datatypes { source, .. } => { + source.status_code() + } ConvertToGrpcDataType { source, .. } => source.status_code(), SqlCommon { source, .. } => source.status_code(), diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs index a257f029e35..be7ed0670dc 100644 --- a/src/sql/src/parsers/create_parser.rs +++ b/src/sql/src/parsers/create_parser.rs @@ -711,8 +711,8 @@ impl<'a> ParserContext<'a> { let mut extensions = ColumnExtensions::default(); let data_type = - if let Some((data_type, type_hints)) = json::parse_json2_type_and_hints(parser)? { - extensions.json_type_hints = type_hints; + if let Some((data_type, options)) = json::parse_json2_type_and_options(parser)? { + extensions.json2_options = options; data_type } else { parser.parse_data_type().context(SyntaxSnafu)? diff --git a/src/sql/src/parsers/create_parser/json.rs b/src/sql/src/parsers/create_parser/json.rs index e2f493719c1..63399f81a17 100644 --- a/src/sql/src/parsers/create_parser/json.rs +++ b/src/sql/src/parsers/create_parser/json.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use datatypes::extension::json::JSON2_REMAINDER_FIELD_NAME; use datatypes::json::JSON2_MAX_STRUCTURED_DEPTH; use snafu::{ResultExt, ensure}; use sqlparser::ast::{DataType, ExactNumberInfo, Expr, ObjectName, UnaryOperator}; @@ -22,14 +23,15 @@ use sqlparser::tokenizer::Token; use crate::ast::Ident; use crate::error::{InvalidSqlSnafu, Result, SyntaxSnafu}; use crate::parsers::create_parser::{INVERTED, SKIPPING}; -use crate::statements::create::JsonTypeHint; +use crate::statements::create::{Json2Options, JsonTypeHint}; use crate::statements::transform::type_alias::get_type_by_alias; const JSON2_TYPE_NAME: &str = "JSON2"; +const MAX_AUTO_EXPANDED_PATHS: &str = "max_auto_expanded_paths"; -pub(super) fn parse_json2_type_and_hints( +pub(super) fn parse_json2_type_and_options( parser: &mut Parser<'_>, -) -> Result)>> { +) -> Result)>> { let token = parser.peek_token(); let Token::Word(word) = &token.token else { return Ok(None); @@ -41,26 +43,62 @@ pub(super) fn parse_json2_type_and_hints( parser.next_token(); let data_type = DataType::Custom(ObjectName::from(vec![Ident::new(JSON2_TYPE_NAME)]), vec![]); - let type_hints = if parser.consume_token(&Token::LParen) { - parse_json2_type_hints(parser)? + let options = if parser.consume_token(&Token::LParen) { + parse_json2_options(parser)? } else { - vec![] + None }; - Ok(Some((data_type, type_hints))) + Ok(Some((data_type, options))) } -fn parse_json2_type_hints(parser: &mut Parser<'_>) -> Result> { - let mut hints = Vec::new(); - +fn parse_json2_options(parser: &mut Parser<'_>) -> Result> { if parser.consume_token(&Token::RParen) { - return Ok(hints); + return Ok(None); } + let mut max_auto_expanded_paths = None; + let mut type_hints = Vec::new(); loop { - let hint = parse_json2_type_hint(parser)?; - ensure_no_path_conflict(&hints, &hint.path)?; - hints.push(hint); + let token = parser.peek_token(); + let is_max_auto_expanded_paths = matches!( + &token.token, + Token::Word(word) + if word.quote_style.is_none() + && word.value.eq_ignore_ascii_case(MAX_AUTO_EXPANDED_PATHS) + ); + if is_max_auto_expanded_paths { + parser.next_token(); + ensure!( + max_auto_expanded_paths.is_none(), + InvalidSqlSnafu { + msg: format!("duplicated JSON2 option '{MAX_AUTO_EXPANDED_PATHS}'") + } + ); + parser.expect_token(&Token::Eq).context(SyntaxSnafu)?; + + let token = parser.next_token(); + let Token::Number(value, _) = token.token else { + return InvalidSqlSnafu { + msg: format!( + "JSON2 option '{MAX_AUTO_EXPANDED_PATHS}' expects a non-negative integer" + ), + } + .fail(); + }; + max_auto_expanded_paths = Some(value.parse::().map_err(|_| { + InvalidSqlSnafu { + msg: format!( + "JSON2 option '{MAX_AUTO_EXPANDED_PATHS}' expects a non-negative integer" + ), + } + .build() + })?); + } else { + let hint = parse_json2_type_hint(parser)?; + ensure_no_path_conflict(&type_hints, &hint.path)?; + type_hints.push(hint); + } if parser.consume_token(&Token::Comma) { if parser.consume_token(&Token::RParen) { @@ -72,11 +110,22 @@ fn parse_json2_type_hints(parser: &mut Parser<'_>) -> Result> } } - Ok(hints) + Ok(Some(Json2Options { + max_auto_expanded_paths, + type_hints, + })) } fn parse_json2_type_hint(parser: &mut Parser<'_>) -> Result { let path = parse_json2_path(parser)?; + ensure!( + path.first().is_none_or(|x| x != JSON2_REMAINDER_FIELD_NAME), + InvalidSqlSnafu { + msg: format!( + "JSON2 type hint path cannot be rooted at reserved field '{JSON2_REMAINDER_FIELD_NAME}'" + ) + } + ); ensure!( path.len() <= JSON2_MAX_STRUCTURED_DEPTH, InvalidSqlSnafu { @@ -309,7 +358,7 @@ CREATE TABLE traces ( column.column_def.data_type, DataType::Custom(_, _) )); - let hints = column.extensions.json_type_hints; + let hints = column.extensions.json2_options.unwrap().type_hints; assert_eq!(hints.len(), 4); assert_eq!(hints[0].path, vec!["service.name"]); @@ -339,6 +388,71 @@ CREATE TABLE traces ( assert!(hints[3].nullable); } + #[test] + fn test_parse_json2_max_auto_expanded_paths() { + let column = parse_json2_column( + r#" +CREATE TABLE traces ( + log_json_data JSON2 ( + http.method STRING, + max_auto_expanded_paths = 0 + ), + ts TIMESTAMP TIME INDEX, +)"#, + ); + + let options = column.extensions.json2_options.unwrap(); + assert_eq!(options.max_auto_expanded_paths, Some(0)); + assert_eq!(options.type_hints.len(), 1); + + let empty = parse_json2_column( + r#" +CREATE TABLE traces ( + log_json_data JSON2 (), + ts TIMESTAMP TIME INDEX, +)"#, + ); + assert!(empty.extensions.json2_options.is_none()); + + let quoted = parse_json2_column( + r#" +CREATE TABLE traces ( + log_json_data JSON2 ( + "max_auto_expanded_paths" STRING, + nested."!__remainder__!" STRING + ), + ts TIMESTAMP TIME INDEX, +)"#, + ); + let options = quoted.extensions.json2_options.unwrap(); + assert_eq!(options.max_auto_expanded_paths, None); + assert_eq!(options.type_hints.len(), 2); + } + + #[test] + fn test_parse_json2_max_auto_expanded_paths_rejects_invalid_options() { + for options in [ + "max_auto_expanded_paths = 0, max_auto_expanded_paths = 1", + "max_auto_expanded_paths = -1", + "max_auto_expanded_paths = 1.5", + "max_auto_expanded_paths = 4294967296", + r#""!__remainder__!".value STRING"#, + ] { + let sql = format!( + "CREATE TABLE traces (log_json_data JSON2 ({options}), ts TIMESTAMP TIME INDEX)" + ); + assert!( + ParserContext::create_with_dialect( + &sql, + &GreptimeDbDialect {}, + ParseOptions::default() + ) + .is_err(), + "{options}" + ); + } + } + #[test] fn test_parse_json2_type_hint_default_nullable() { let column = parse_json2_column( @@ -349,7 +463,7 @@ CREATE TABLE traces ( )"#, ); - let hints = column.extensions.json_type_hints; + let hints = column.extensions.json2_options.unwrap().type_hints; assert_eq!(hints.len(), 1); assert!(hints[0].nullable); } @@ -369,7 +483,7 @@ CREATE TABLE traces ( )"#, ); - let hints = column.extensions.json_type_hints; + let hints = column.extensions.json2_options.unwrap().type_hints; assert_eq!(hints.len(), 4); assert_eq!(hints[0].path, vec!["a", "b"]); assert_eq!(hints[1].path, vec!["x", "y"]); @@ -402,7 +516,7 @@ CREATE TABLE traces ( )"#, ); - let hints = column.extensions.json_type_hints; + let hints = column.extensions.json2_options.unwrap().type_hints; assert_eq!(hints.len(), 14); for hint in hints.iter().take(6) { assert_eq!(hint.data_type, DataType::BigInt(None)); @@ -428,7 +542,7 @@ CREATE TABLE traces ( )"#, ); - let hints = column.extensions.json_type_hints; + let hints = column.extensions.json2_options.unwrap().type_hints; assert_eq!(hints.len(), 2); assert_eq!( hints[0] diff --git a/src/sql/src/statements.rs b/src/sql/src/statements.rs index aafe62d2e6e..4f96b66f87f 100644 --- a/src/sql/src/statements.rs +++ b/src/sql/src/statements.rs @@ -718,11 +718,7 @@ mod tests { "true".to_string(), ), ])), - vector_options: None, - skipping_index_options: None, - inverted_index_options: None, - json_type_hints: vec![], - vector_index_options: None, + ..Default::default() }, }; @@ -749,17 +745,13 @@ mod tests { options: vec![], }, extensions: ColumnExtensions { - fulltext_index_options: None, - vector_options: None, - skipping_index_options: None, - inverted_index_options: None, - json_type_hints: vec![], vector_index_options: Some(OptionMap::from([ ("metric".to_string(), "cosine".to_string()), ("connectivity".to_string(), "32".to_string()), ("expansion_add".to_string(), "200".to_string()), ("expansion_search".to_string(), "100".to_string()), ])), + ..Default::default() }, }; @@ -790,12 +782,8 @@ mod tests { options: vec![], }, extensions: ColumnExtensions { - fulltext_index_options: None, - vector_options: None, - skipping_index_options: None, - inverted_index_options: None, - json_type_hints: vec![], vector_index_options: Some(OptionMap::default()), + ..Default::default() }, }; diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs index dc94c7f8889..f3fbc908362 100644 --- a/src/sql/src/statements/create.rs +++ b/src/sql/src/statements/create.rs @@ -136,7 +136,28 @@ pub struct ColumnExtensions { pub inverted_index_options: Option, /// Vector index options for HNSW-based vector similarity search. pub vector_index_options: Option, - pub json_type_hints: Vec, + /// JSON2-specific column options. + pub json2_options: Option, +} + +/// JSON2-specific options represented in the SQL AST. +#[derive(Debug, PartialEq, Eq, Clone, Visit, VisitMut, Serialize)] +pub struct Json2Options { + /// Maximum number of unhinted JSON2 paths expanded into Arrow fields. + pub(crate) max_auto_expanded_paths: Option, + /// Paths stored as explicitly typed JSON2 fields. + pub(crate) type_hints: Vec, +} + +impl Display for Json2Options { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + let mut options = Vec::with_capacity(self.type_hints.len() + 1); + if let Some(max) = self.max_auto_expanded_paths { + options.push(format!("max_auto_expanded_paths = {max}")); + } + options.extend(self.type_hints.iter().map(format_json_type_hint)); + write!(f, "(\n {}\n )", options.iter().join(",\n ")) + } } #[derive(Debug, PartialEq, Eq, Clone, Visit, VisitMut, Serialize)] @@ -180,12 +201,8 @@ impl Display for Column { } write!(f, "{} {}", self.column_def.name, self.column_def.data_type)?; - if !self.extensions.json_type_hints.is_empty() { - write!( - f, - "{}", - format_json_type_hints(&self.extensions.json_type_hints) - )?; + if let Some(options) = &self.extensions.json2_options { + write!(f, "{options}")?; } for option in &self.column_def.options { write!(f, " {option}")?; @@ -335,29 +352,30 @@ impl ColumnExtensions { } pub fn build_json_settings(&self) -> Result> { - if self.json_type_hints.is_empty() { + let Some(options) = &self.json2_options else { return Ok(None); - } + }; - Ok(Some(JsonSettings::new( - self.json_type_hints - .iter() - .map(|hint| { - Ok(datatypes::json::JsonTypeHint { - path: hint.path.clone(), - data_type: json_type_hint_concrete_data_type(&hint.data_type)?, - nullable: hint.nullable, - default_constraint: build_json_type_hint_default_constraint(hint)?, - inverted_index: hint.inverted_index, - }) + let type_hints = options + .type_hints + .iter() + .map(|hint| { + Ok(datatypes::json::JsonTypeHint { + path: hint.path.clone(), + data_type: json_type_hint_concrete_data_type(&hint.data_type)?, + nullable: hint.nullable, + default_constraint: build_json_type_hint_default_constraint(hint)?, + inverted_index: hint.inverted_index, }) - .collect::>>()?, - ))) + }) + .collect::>>()?; + let settings = JsonSettings::try_new(type_hints, options.max_auto_expanded_paths)?; + Ok(Some(settings)) } pub fn set_json_settings(&mut self, settings: JsonSettings) -> Result<()> { - self.json_type_hints = settings - .type_hints + let (type_hints, max_auto_expanded_paths) = settings.into_parts(); + let type_hints = type_hints .into_iter() .map(|hint| { let data_type = json_type_hint_sql_data_type(&hint.data_type)?; @@ -374,6 +392,11 @@ impl ColumnExtensions { }) }) .collect::>>()?; + self.json2_options = (max_auto_expanded_paths.is_some() || !type_hints.is_empty()) + .then_some(Json2Options { + max_auto_expanded_paths, + type_hints, + }); Ok(()) } } @@ -487,13 +510,6 @@ fn format_json_type_hint(hint: &JsonTypeHint) -> String { ) } -fn format_json_type_hints(hints: &[JsonTypeHint]) -> String { - format!( - "(\n {}\n )", - hints.iter().map(format_json_type_hint).join(",\n ") - ) -} - fn format_json_path_segment(segment: &str) -> String { format!("\"{}\"", segment.replace('"', "\"\"")) } @@ -786,6 +802,7 @@ mod tests { use datatypes::schema::ColumnDefaultConstraint; use datatypes::value::Value; + use super::*; use crate::dialect::GreptimeDbDialect; use crate::error::Error; use crate::parser::{ParseOptions, ParserContext}; @@ -1066,7 +1083,7 @@ ENGINE=mito .build_json_settings() .unwrap() .unwrap(); - let hints = settings.type_hints; + let hints = settings.type_hints(); assert_eq!(hints[0].data_type, ConcreteDataType::int64_datatype()); assert_eq!( @@ -1118,10 +1135,10 @@ ENGINE=mito } #[test] - fn test_set_json_settings_normalizes_type_hint_sql_types() { + fn test_set_json_settings_normalizes_type_hint_sql_types() -> Result<()> { let mut extensions = super::ColumnExtensions::default(); - extensions - .set_json_settings(JsonSettings::new(vec![ + let settings = JsonSettings::try_new( + vec![ DatatypeJsonTypeHint { path: vec!["i".to_string()], data_type: ConcreteDataType::int32_datatype(), @@ -1157,36 +1174,51 @@ ENGINE=mito default_constraint: None, inverted_index: false, }, - ])) - .unwrap(); + ], + None, + )?; + extensions.set_json_settings(settings)?; assert_eq!( extensions - .json_type_hints + .json2_options + .unwrap() + .type_hints .iter() .map(|hint| hint.data_type.to_string()) .collect::>(), vec!["BIGINT", "DOUBLE", "BIGINT UNSIGNED", "STRING", "BOOLEAN"] ); + Ok(()) } #[test] - fn test_set_json_settings_rejects_unsupported_type_hint_type() { - let mut extensions = super::ColumnExtensions::default(); - let err = extensions - .set_json_settings(JsonSettings::new(vec![DatatypeJsonTypeHint { + fn test_set_json_settings_rejects_unsupported_type_hint_type() -> Result<()> { + let err = JsonSettings::try_new( + vec![DatatypeJsonTypeHint { path: vec!["u".to_string()], data_type: ConcreteDataType::date_datatype(), nullable: true, default_constraint: None, inverted_index: false, - }])) - .unwrap_err(); + }], + None, + ) + .unwrap_err(); assert!( err.to_string() .contains("unsupported JSON2 type hint data type") ); + Ok(()) + } + + #[test] + fn test_set_empty_json_settings_omits_json2_options() -> Result<()> { + let mut extensions = ColumnExtensions::default(); + extensions.set_json_settings(JsonSettings::default())?; + assert!(extensions.json2_options.is_none()); + Ok(()) } #[test] @@ -1364,15 +1396,11 @@ AS SELECT number FROM numbers_input where number > 10"#, // Test zero connectivity should fail let extensions = ColumnExtensions { - fulltext_index_options: None, - vector_options: None, - skipping_index_options: None, - inverted_index_options: None, - json_type_hints: vec![], vector_index_options: Some(OptionMap::from([( "connectivity".to_string(), "0".to_string(), )])), + ..Default::default() }; let result = extensions.build_vector_index_options(); assert!(result.is_err()); @@ -1385,15 +1413,11 @@ AS SELECT number FROM numbers_input where number > 10"#, // Test zero expansion_add should fail let extensions = ColumnExtensions { - fulltext_index_options: None, - vector_options: None, - skipping_index_options: None, - inverted_index_options: None, - json_type_hints: vec![], vector_index_options: Some(OptionMap::from([( "expansion_add".to_string(), "0".to_string(), )])), + ..Default::default() }; let result = extensions.build_vector_index_options(); assert!(result.is_err()); @@ -1406,15 +1430,11 @@ AS SELECT number FROM numbers_input where number > 10"#, // Test zero expansion_search should fail let extensions = ColumnExtensions { - fulltext_index_options: None, - vector_options: None, - skipping_index_options: None, - inverted_index_options: None, - json_type_hints: vec![], vector_index_options: Some(OptionMap::from([( "expansion_search".to_string(), "0".to_string(), )])), + ..Default::default() }; let result = extensions.build_vector_index_options(); assert!(result.is_err()); @@ -1427,16 +1447,12 @@ AS SELECT number FROM numbers_input where number > 10"#, // Test valid values should succeed let extensions = ColumnExtensions { - fulltext_index_options: None, - vector_options: None, - skipping_index_options: None, - inverted_index_options: None, - json_type_hints: vec![], vector_index_options: Some(OptionMap::from([ ("connectivity".to_string(), "32".to_string()), ("expansion_add".to_string(), "200".to_string()), ("expansion_search".to_string(), "100".to_string()), ])), + ..Default::default() }; let result = extensions.build_vector_index_options(); assert!(result.is_ok()); diff --git a/tests/compatibility/cases/json2_storage_layout_settings/case.toml b/tests/compatibility/cases/json2_storage_layout_settings/case.toml new file mode 100644 index 00000000000..8f1dd3fe64a --- /dev/null +++ b/tests/compatibility/cases/json2_storage_layout_settings/case.toml @@ -0,0 +1,8 @@ +name = "json2_storage_layout_settings" +reason = "Verify JSON2 storage layout settings in Arrow extension metadata survive restarts and upgrades." +introduced_by = "PR #8895" +topologies = ["distributed", "standalone"] +from_range = [">=v1.3.0"] +to_range = [">=v1.3.0"] +features = ["json2", "table", "metadata", "read"] +owner = "query" diff --git a/tests/compatibility/cases/json2_storage_layout_settings/setup.sql b/tests/compatibility/cases/json2_storage_layout_settings/setup.sql new file mode 100644 index 00000000000..2592185ef23 --- /dev/null +++ b/tests/compatibility/cases/json2_storage_layout_settings/setup.sql @@ -0,0 +1,11 @@ +CREATE TABLE t_json2_storage_layout_settings ( + ts TIMESTAMP TIME INDEX, + j JSON2(max_auto_expanded_paths = 0) +) WITH ( + 'append_mode' = 'true' +); + +INSERT INTO t_json2_storage_layout_settings VALUES + ('2026-08-17 00:00:00+0000', '{"a": 1}'); + +ADMIN FLUSH_TABLE('t_json2_storage_layout_settings'); diff --git a/tests/compatibility/cases/json2_storage_layout_settings/verify.result b/tests/compatibility/cases/json2_storage_layout_settings/verify.result new file mode 100644 index 00000000000..452ed7a5516 --- /dev/null +++ b/tests/compatibility/cases/json2_storage_layout_settings/verify.result @@ -0,0 +1,27 @@ +SELECT j.a AS a +FROM t_json2_storage_layout_settings; + ++---+ +| a | ++---+ +| 1 | ++---+ + +SHOW CREATE TABLE t_json2_storage_layout_settings; + ++---------------------------------+----------------------------------------------------------------+ +| Table | Create Table | ++---------------------------------+----------------------------------------------------------------+ +| t_json2_storage_layout_settings | CREATE TABLE IF NOT EXISTS "t_json2_storage_layout_settings" ( | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | "j" JSON2( | +| | max_auto_expanded_paths = 0 | +| | ) NULL, | +| | TIME INDEX ("ts") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | append_mode = 'true' | +| | ) | ++---------------------------------+----------------------------------------------------------------+ diff --git a/tests/compatibility/cases/json2_storage_layout_settings/verify.sql b/tests/compatibility/cases/json2_storage_layout_settings/verify.sql new file mode 100644 index 00000000000..dae3cd8a235 --- /dev/null +++ b/tests/compatibility/cases/json2_storage_layout_settings/verify.sql @@ -0,0 +1,4 @@ +SELECT j.a AS a +FROM t_json2_storage_layout_settings; + +SHOW CREATE TABLE t_json2_storage_layout_settings; From b882e393dfadcf9406c76d1017fe64e0652f0202 Mon Sep 17 00:00:00 2001 From: sun Date: Tue, 18 Aug 2026 03:47:58 +0000 Subject: [PATCH 60/92] feat: update dashboard to v0.13.13 (#8898) --- src/servers/dashboard/VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/servers/dashboard/VERSION b/src/servers/dashboard/VERSION index 3d431cfd5f0..b68b26e6f78 100644 --- a/src/servers/dashboard/VERSION +++ b/src/servers/dashboard/VERSION @@ -1 +1 @@ -v0.13.12 +v0.13.13 From b96ea86a621e7283e77e4e37dd4e3ec8dfc44f73 Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Wed, 19 Aug 2026 02:27:06 +0000 Subject: [PATCH 61/92] refactor(json2): add JSON2 v2 physical layout primitives (#8901) * refactor(json2): add JSON2 v2 physical layout primitives Signed-off-by: luofucong * resolve PR comments Signed-off-by: luofucong * fix ci Signed-off-by: luofucong --------- Signed-off-by: luofucong --- Cargo.lock | 50 +++++ Cargo.toml | 3 + src/datatypes/Cargo.toml | 3 + src/datatypes/src/extension/json.rs | 202 +++++++++++++++++- src/datatypes/src/vectors/json.rs | 1 + src/datatypes/src/vectors/json/array.rs | 248 ++++++++++++++++++++-- src/datatypes/src/vectors/json/variant.rs | 239 +++++++++++++++++++++ 7 files changed, 728 insertions(+), 18 deletions(-) create mode 100644 src/datatypes/src/vectors/json/variant.rs diff --git a/Cargo.lock b/Cargo.lock index 0f473d4e030..16b4c6b1045 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4554,6 +4554,9 @@ dependencies = [ "num", "num-traits", "ordered-float 4.6.0", + "parquet-variant", + "parquet-variant-compute", + "parquet-variant-json", "paste", "regex", "serde", @@ -10259,6 +10262,53 @@ dependencies = [ "zstd", ] +[[package]] +name = "parquet-variant" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74c8db065291f088a2aad8ab831853eae1871c0d311c8d0b83bbc3b7e735d0fc" +dependencies = [ + "arrow 58.3.0", + "arrow-schema 58.3.0", + "chrono", + "half", + "indexmap 2.13.0", + "num-traits", + "simdutf8", + "uuid", +] + +[[package]] +name = "parquet-variant-compute" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a530e8d5b5e14efcb39c9a6ec55432ad11f6afb7dc4455a79be0dc615fe3cc31" +dependencies = [ + "arrow 58.3.0", + "arrow-schema 58.3.0", + "chrono", + "half", + "indexmap 2.13.0", + "parquet-variant", + "parquet-variant-json", + "serde_json", + "uuid", +] + +[[package]] +name = "parquet-variant-json" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00ed89908289f67caa2ca078f9ff9aacd6229a313ec92b12bf4f48f613dc2b97" +dependencies = [ + "arrow-schema 58.3.0", + "base64 0.22.1", + "chrono", + "parquet-variant", + "serde_json", + "uuid", +] + [[package]] name = "parse-size" version = "1.1.0" diff --git a/Cargo.toml b/Cargo.toml index daee1bd4090..777d8b302b4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -195,6 +195,9 @@ otel-arrow-rust = { git = "https://github.com/GreptimeTeam/otel-arrow", rev = "5 ] } parking_lot = "0.12" parquet = { version = "58.3", default-features = false, features = ["arrow", "async", "object_store"] } +parquet-variant = "58.3" +parquet-variant-compute = "58.3" +parquet-variant-json = "58.3" paste = "1.0" pin-project = "1.0" pretty_assertions = "1.4.0" diff --git a/src/datatypes/Cargo.toml b/src/datatypes/Cargo.toml index 3bf2a5d8eec..124b1e0ea64 100644 --- a/src/datatypes/Cargo.toml +++ b/src/datatypes/Cargo.toml @@ -27,6 +27,9 @@ jsonb.workspace = true num = "0.4" num-traits = "0.2" ordered-float.workspace = true +parquet-variant.workspace = true +parquet-variant-compute.workspace = true +parquet-variant-json.workspace = true paste.workspace = true regex.workspace = true serde.workspace = true diff --git a/src/datatypes/src/extension/json.rs b/src/datatypes/src/extension/json.rs index 109aa853b21..ab4a138420e 100644 --- a/src/datatypes/src/extension/json.rs +++ b/src/datatypes/src/extension/json.rs @@ -18,31 +18,148 @@ use std::sync::Arc; use arrow_schema::extension::{ EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType, }; -use arrow_schema::{ArrowError, DataType, Field}; +use arrow_schema::{ArrowError, DataType, Field, FieldRef}; +use parquet_variant_compute::VariantType; use serde::{Deserialize, Serialize}; -use snafu::ResultExt; +use snafu::{ResultExt, ensure}; +use crate::error::InvalidJson2LayoutSnafu; pub use crate::json::JSON2_REMAINDER_FIELD_NAME; use crate::json::JsonSettings; const LEGACY_JSON_STRUCTURE_SETTINGS_KEY: &str = "json_structure_settings"; +const JSON2_LAYOUT_V1: u8 = 1; +const JSON2_LAYOUT_V2: u8 = 2; -#[derive(Debug, Clone, Serialize, Deserialize, Default)] +/// Parsed physical layout of a JSON2 Arrow root field. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Json2PhysicalLayout { + version: u8, +} + +impl Json2PhysicalLayout { + /// Parses the JSON2 layout version from a root field. + /// + /// The version is read from the extension metadata; missing metadata + /// defaults to V1. + /// + /// Supported versions: + /// - V1 + /// - V2 + /// + /// Errors when the field is not a JSON2 extension or the version is not + /// one of the supported versions. + pub fn try_from_root(field: &Field) -> crate::error::Result { + ensure!( + is_json2_extension_type(field), + InvalidJson2LayoutSnafu { + reason: format!("field '{}' is not a JSON2 extension", field.name()), + } + ); + + let version = match field.extension_type_name() { + Some(Json2ExtensionType::NAME) => field + .metadata() + .get(EXTENSION_TYPE_METADATA_KEY) + .map(|x| parse_version(x)) + .transpose()? + .flatten() + .unwrap_or(JSON2_LAYOUT_V1), + _ => JSON2_LAYOUT_V1, + }; + ensure!( + matches!(version, JSON2_LAYOUT_V1 | JSON2_LAYOUT_V2), + InvalidJson2LayoutSnafu { + reason: format!("unsupported JSON2 layout version: {version}"), + } + ); + Ok(Self { version }) + } + + /// Returns whether this is the JSON2 physical layout version 2. + pub fn is_version_2(&self) -> bool { + self.version == JSON2_LAYOUT_V2 + } +} + +fn parse_version(metadata: &str) -> crate::error::Result> { + serde_json::from_str::(metadata) + .map(|x| x.layout_version) + .map_err(|e| { + InvalidJson2LayoutSnafu { + reason: format!(r#"invalid extension metadata: "{metadata}", error: {e}"#), + } + .build() + }) +} + +/// Returns the remainder field of a JSON2 v2 root. +pub(crate) fn json2_remainder_field(field: &Field) -> crate::error::Result> { + let layout = Json2PhysicalLayout::try_from_root(field)?; + if !layout.is_version_2() { + return Ok(None); + } + + let remainder = if let DataType::Struct(fields) = field.data_type() { + fields + .iter() + .find(|x| x.name() == JSON2_REMAINDER_FIELD_NAME) + } else { + None + }; + + if let Some(remainder) = remainder { + let _ = remainder.try_extension_type::().map_err(|e| { + InvalidJson2LayoutSnafu { + reason: e.to_string(), + } + .build() + })?; + } + Ok(remainder) +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct JsonMetadata { /// JSON2 settings stored in Arrow extension metadata. json_settings: JsonSettings, + /// Physical JSON2 layout used by this Arrow field. + #[serde(default, skip_serializing_if = "Option::is_none")] + layout_version: Option, } impl JsonMetadata { - /// Creates JSON2 extension metadata. + /// Creates metadata for the legacy JSON2 layout. pub fn new(json_settings: JsonSettings) -> Self { - Self { json_settings } + Self { + json_settings, + layout_version: None, + } + } + + /// Creates metadata for the new JSON2 physical v2 layout. + pub fn new_v2(json_settings: JsonSettings) -> Self { + Self { + json_settings, + layout_version: Some(JSON2_LAYOUT_V2), + } } /// Returns the JSON2 settings. pub fn json_settings(&self) -> &JsonSettings { &self.json_settings } + + /// Returns whether this metadata describes JSON2 layout version 2. + pub fn is_version_2(&self) -> bool { + self.layout_version == Some(JSON2_LAYOUT_V2) + } +} + +impl Default for JsonMetadata { + fn default() -> Self { + Self::new(JsonSettings::default()) + } } /// Arrow extension type for legacy JSONB columns. @@ -227,6 +344,7 @@ mod tests { use arrow_schema::{Field, Fields}; use super::*; + use crate::vectors::json::variant::variant_field; #[test] fn test_json2_extension_type_detection() { @@ -282,4 +400,78 @@ mod tests { assert!(JsonExtensionType::try_new(&DataType::Null, ()).is_ok()); assert!(JsonExtensionType::try_new(&DataType::Struct(Fields::empty()), ()).is_err()); } + + #[test] + fn test_json_metadata_layout_version_compatibility() -> serde_json::Result<()> { + let legacy: JsonMetadata = serde_json::from_str(r#"{"json_settings":{}}"#)?; + assert!(!legacy.is_version_2()); + + let metadata = JsonMetadata::new_v2(JsonSettings::default()); + assert!(metadata.is_version_2()); + let serialized = serde_json::to_string(&metadata)?; + let deserialized: JsonMetadata = serde_json::from_str(&serialized)?; + assert!(deserialized.is_version_2()); + assert_eq!(deserialized, metadata); + Ok(()) + } + + #[test] + fn test_parse_json2_physical_layout() -> crate::error::Result<()> { + let legacy = Field::new("data", DataType::Struct(Fields::empty()), true) + .with_extension_type(Json2ExtensionType::default()); + assert!(!Json2PhysicalLayout::try_from_root(&legacy)?.is_version_2()); + + let v2 = Field::new( + "data", + DataType::Struct( + vec![ + Arc::new(variant_field(JSON2_REMAINDER_FIELD_NAME, true)), + Arc::new(Field::new("count", DataType::Int64, true)), + ] + .into(), + ), + true, + ) + .with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2( + JsonSettings::default(), + )))); + assert!(Json2PhysicalLayout::try_from_root(&v2)?.is_version_2()); + let remainder = json2_remainder_field(&v2)?.unwrap(); + assert_eq!(JSON2_REMAINDER_FIELD_NAME, remainder.name()); + Ok(()) + } + + #[test] + fn test_reject_invalid_json2_physical_layout() { + let field = Field::new("data", DataType::Struct(Fields::empty()), true); + assert!(Json2PhysicalLayout::try_from_root(&field).is_err()); + + let metadata = + Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(JsonSettings::default()))); + let missing = Field::new("data", DataType::Struct(Fields::empty()), true) + .with_extension_type(metadata.clone()); + assert!(json2_remainder_field(&missing).is_ok_and(|x| x.is_none())); + + let invalid = Field::new( + "data", + DataType::Struct( + vec![Arc::new(Field::new( + JSON2_REMAINDER_FIELD_NAME, + DataType::Binary, + true, + ))] + .into(), + ), + true, + ) + .with_extension_type(metadata); + assert!(json2_remainder_field(&invalid).is_err()); + + let future = Field::new("data", DataType::Struct(Fields::empty()), true) + .with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata { + json_settings: JsonSettings::default(), + layout_version: Some(JSON2_LAYOUT_V2 + 1), + }))); + assert!(Json2PhysicalLayout::try_from_root(&future).is_err()); + } } diff --git a/src/datatypes/src/vectors/json.rs b/src/datatypes/src/vectors/json.rs index 3dcb8d7bc07..c68e6cf3d92 100644 --- a/src/datatypes/src/vectors/json.rs +++ b/src/datatypes/src/vectors/json.rs @@ -14,3 +14,4 @@ pub mod array; pub(crate) mod builder; +pub mod variant; diff --git a/src/datatypes/src/vectors/json/array.rs b/src/datatypes/src/vectors/json/array.rs index afe9f553056..0426334c3dd 100644 --- a/src/datatypes/src/vectors/json/array.rs +++ b/src/datatypes/src/vectors/json/array.rs @@ -17,9 +17,12 @@ use std::sync::Arc; use arrow::compute::{can_cast_types, cast}; use arrow_array::cast::AsArray; -use arrow_array::types::{Float64Type, Int64Type, UInt64Type}; +use arrow_array::types::{ + Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, + UInt32Type, UInt64Type, +}; use arrow_array::{Array, ArrayRef, GenericListArray, ListArray, StructArray, new_null_array}; -use arrow_schema::{DataType, FieldRef}; +use arrow_schema::{DataType, Field, FieldRef}; use common_telemetry::trace; use serde_json::Value; use snafu::{OptionExt, ResultExt}; @@ -29,9 +32,11 @@ use crate::data_type::ConcreteDataType; use crate::error::{ AlignJsonArraySnafu, ArrowComputeSnafu, InvalidJsonSnafu, InvalidJsonbSnafu, Result, }; +use crate::extension::json::{JSON2_REMAINDER_FIELD_NAME, json2_remainder_field}; use crate::json::value::{decode_json_variant, encode_serde_json_as_jsonb}; use crate::prelude::{DataType as _, Value as GreptimeValue}; use crate::value::{ListValue, StructValue}; +use crate::vectors::json::variant::variant_to_json_values; pub struct JsonArray<'a> { inner: &'a ArrayRef, @@ -48,8 +53,15 @@ impl JsonArray<'_> { let value = match array.data_type() { DataType::Null => Value::Null, DataType::Boolean => Value::Bool(array.as_boolean().value(i)), + DataType::Int8 => Value::from(array.as_primitive::().value(i)), + DataType::Int16 => Value::from(array.as_primitive::().value(i)), + DataType::Int32 => Value::from(array.as_primitive::().value(i)), DataType::Int64 => Value::from(array.as_primitive::().value(i)), + DataType::UInt8 => Value::from(array.as_primitive::().value(i)), + DataType::UInt16 => Value::from(array.as_primitive::().value(i)), + DataType::UInt32 => Value::from(array.as_primitive::().value(i)), DataType::UInt64 => Value::from(array.as_primitive::().value(i)), + DataType::Float32 => Value::from(array.as_primitive::().value(i)), DataType::Float64 => Value::from(array.as_primitive::().value(i)), DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { Value::String(string_array_value(array, i).to_string()) @@ -92,6 +104,60 @@ impl JsonArray<'_> { Ok(value) } + /// Projects a physical JSON2 array to a logical query type. + /// + /// TODO(LFC) Supersede `project_to_v2` to `project_to`. + pub fn project_to_v2(&self, field: &Field, target: &DataType) -> Result { + if json2_remainder_field(field)?.is_some() { + project_json_values(self.json2_values()?, target) + } else { + self.project_to(target) + } + } + + fn json2_values(&self) -> Result> { + let structs = self.inner.as_struct_opt().context(AlignJsonArraySnafu { + reason: "JSON2 layout v2 root array must be a struct", + })?; + let remainder = structs.column_by_name(JSON2_REMAINDER_FIELD_NAME); + let mut remainders = if let Some(remainder) = remainder { + variant_to_json_values(remainder)? + } else { + vec![None; structs.len()] + }; + let mut values = Vec::with_capacity(structs.len()); + let mut path = Vec::new(); + + for (i, remainder) in remainders.iter_mut().enumerate() { + if structs.is_null(i) { + values.push(Value::Null); + continue; + } + + let mut object = match remainder.take() { + None => serde_json::Map::new(), + Some(Value::Object(object)) => object, + Some(value) => { + return InvalidJsonSnafu { + value: format!("JSON2 layout v2 remainder must be an object, got {value}"), + } + .fail(); + } + }; + + for (child, column) in structs.fields().iter().zip(structs.columns()) { + if child.name() == JSON2_REMAINDER_FIELD_NAME { + continue; + } + let value = JsonArray::from(column).try_get_value(i)?; + merge_explicit_value(&mut object, child.name().clone(), value, &mut path)?; + } + values.push(Value::Object(object)); + } + + Ok(values) + } + /// Normalizes a JSON2 array to the wider `expect` data type without losing /// information. /// @@ -326,6 +392,34 @@ impl JsonArray<'_> { } } +fn merge_explicit_value( + remainder: &mut serde_json::Map, + key: String, + explicit: Value, + path: &mut Vec, +) -> Result<()> { + let Some(existing) = remainder.get_mut(&key) else { + remainder.insert(key, explicit); + return Ok(()); + }; + path.push(key); + + let (Value::Object(remainder), Value::Object(explicit)) = (existing, explicit) else { + return InvalidJsonSnafu { + value: format!( + "cannot merge '{}' in explicit fields and remainder: not both objects", + path.join("."), + ), + } + .fail(); + }; + for (key, value) in explicit { + merge_explicit_value(remainder, key, value, path)?; + } + path.pop(); + Ok(()) +} + /// Returns whether Arrow can cast between the types without JSON-aware projection. /// Binary and nested types require JSONB decoding or recursive projection. fn can_fast_cast_types(from_type: &DataType, to_type: &DataType) -> bool { @@ -453,13 +547,16 @@ impl<'a> From<&'a ArrayRef> for JsonArray<'a> { mod test { use arrow_array::types::Int64Type; use arrow_array::{ - BinaryArray, BooleanArray, Float64Array, Int32Array, Int64Array, ListArray, StringArray, - UInt64Array, + BinaryArray, BooleanArray, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, + Int64Array, ListArray, StringArray, UInt8Array, UInt16Array, UInt32Array, UInt64Array, }; use arrow_schema::{Field, Fields}; use serde_json::json; use super::*; + use crate::extension::json::{Json2ExtensionType, JsonMetadata}; + use crate::json::JsonSettings; + use crate::vectors::json::variant::{json_values_to_variant, variant_field}; #[test] fn test_try_get_value() -> Result<()> { @@ -474,6 +571,20 @@ mod test { assert_eq!(JsonArray::from(&ints).try_get_value(0)?, json!(-7)); assert_eq!(JsonArray::from(&ints).try_get_value(1)?, Value::Null); + macro_rules! assert_number { + ($array:expr, $expected:expr) => {{ + let array: ArrayRef = Arc::new($array); + assert_eq!(JsonArray::from(&array).try_get_value(0)?, json!($expected)); + }}; + } + assert_number!(Int8Array::from(vec![-8]), -8); + assert_number!(Int16Array::from(vec![-16]), -16); + assert_number!(Int32Array::from(vec![-32]), -32); + assert_number!(UInt8Array::from(vec![8]), 8); + assert_number!(UInt16Array::from(vec![16]), 16); + assert_number!(UInt32Array::from(vec![32]), 32); + assert_number!(Float32Array::from(vec![1.25]), 1.25); + let floats: ArrayRef = Arc::new(Float64Array::from(vec![Some(1.5)])); assert_eq!(JsonArray::from(&floats).try_get_value(0)?, json!(1.5)); @@ -529,15 +640,6 @@ mod test { json!({"flag": null, "items": [2]}) ); - let unsupported: ArrayRef = Arc::new(Int32Array::from(vec![1])); - assert_eq!( - JsonArray::from(&unsupported) - .try_get_value(0) - .unwrap_err() - .to_string(), - "Invalid JSON: unknown JSON type Int32" - ); - Ok(()) } @@ -915,4 +1017,124 @@ mod test { Ok(()) } + + #[test] + fn test_reconstruct_json2_v2_value() -> Result<()> { + let remainders = json_values_to_variant(&[ + Some(json!({"cold": 1, "nested": {"right": true}})), + Some(json!({"!__remainder__!": "user value"})), + ])?; + let remainder = Arc::new(variant_field(JSON2_REMAINDER_FIELD_NAME, true)); + let nested = Arc::new(Field::new_struct( + "nested", + [Arc::new(Field::new("left", DataType::Utf8, true))], + true, + )); + let nested_values: ArrayRef = Arc::new(StructArray::from(vec![( + Arc::new(Field::new("left", DataType::Utf8, true)), + Arc::new(StringArray::from(vec![Some("value"), None])) as ArrayRef, + )])); + let fields = Fields::from(vec![ + remainder, + Arc::new(Field::new("count", DataType::Int64, true)), + nested, + ]); + let array: ArrayRef = Arc::new(StructArray::new( + fields.clone(), + vec![ + remainders, + Arc::new(Int64Array::from(vec![Some(42), None])), + nested_values, + ], + None, + )); + let field = Field::new("data", DataType::Struct(fields), true).with_extension_type( + Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(JsonSettings::default()))), + ); + + assert_eq!( + json!({ + "cold": 1, + "count": 42, + "nested": {"left": "value", "right": true} + }), + JsonArray::from(&array).json2_values()?[0] + ); + assert_eq!( + json!({ + "!__remainder__!": "user value", + "count": null, + "nested": {"left": null} + }), + JsonArray::from(&array).json2_values()?[1] + ); + + let target = DataType::Struct( + vec![ + Arc::new(Field::new("cold", DataType::UInt64, true)), + Arc::new(Field::new("count", DataType::Int64, true)), + ] + .into(), + ); + let projected = JsonArray::from(&array).project_to_v2(&field, &target)?; + assert_eq!( + json!({"cold": 1, "count": 42}), + JsonArray::from(&projected).try_get_value(0)? + ); + assert_eq!( + json!({"cold": null, "count": null}), + JsonArray::from(&projected).try_get_value(1)? + ); + Ok(()) + } + + #[test] + fn test_project_partial_json2_v2_without_remainder() -> Result<()> { + let fields = Fields::from(vec![Arc::new(Field::new("hot", DataType::Int64, true))]); + let array: ArrayRef = Arc::new(StructArray::new( + fields.clone(), + vec![Arc::new(Int64Array::from(vec![1, 2]))], + None, + )); + let field = Field::new("data", DataType::Struct(fields), true).with_extension_type( + Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(JsonSettings::default()))), + ); + + let projected = JsonArray::from(&array).project_to_v2(&field, field.data_type())?; + assert!(Arc::ptr_eq(&array, &projected)); + Ok(()) + } + + #[test] + fn test_reject_conflict_json2_v2_path() -> Result<()> { + let remainders = json_values_to_variant(&[Some(json!({"count": 1}))])?; + let fields = Fields::from(vec![ + Arc::new(variant_field(JSON2_REMAINDER_FIELD_NAME, true)), + Arc::new(Field::new("count", DataType::Int64, true)), + ]); + let array: ArrayRef = Arc::new(StructArray::new( + fields, + vec![remainders, Arc::new(Int64Array::from(vec![2]))], + None, + )); + let error = JsonArray::from(&array).json2_values().unwrap_err(); + assert!( + error.to_string().contains( + "cannot merge 'count' in explicit fields and remainder: not both objects" + ) + ); + + let Value::Object(mut remainder) = json!({"nested": {"count": 1}}) else { + unreachable!(); + }; + let error = merge_explicit_value( + &mut remainder, + "nested".to_string(), + json!({"count": 2}), + &mut Vec::new(), + ) + .unwrap_err(); + assert!(error.to_string().contains("cannot merge 'nested.count'")); + Ok(()) + } } diff --git a/src/datatypes/src/vectors/json/variant.rs b/src/datatypes/src/vectors/json/variant.rs new file mode 100644 index 00000000000..1a4b8a861d5 --- /dev/null +++ b/src/datatypes/src/vectors/json/variant.rs @@ -0,0 +1,239 @@ +// Copyright 2023 Greptime Team +// +// 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. + +use std::sync::Arc; + +use arrow_array::ArrayRef; +#[cfg(test)] +use arrow_schema::ArrowError; +use arrow_schema::{DataType, Field}; +#[cfg(test)] +use parquet_variant::{ObjectFieldBuilder, Variant, VariantBuilderExt, VariantDecimal16}; +#[cfg(test)] +use parquet_variant_compute::VariantArrayBuilder; +use parquet_variant_compute::{VariantArray, VariantType}; +use parquet_variant_json::VariantToJson; +use snafu::ResultExt; + +use crate::error::{ArrowComputeSnafu, Result}; +#[cfg(test)] +use crate::json::value::{JsonNumber, JsonVariant, decode_json_variant}; + +/// Returns the canonical Arrow field for an unshredded Parquet Variant array. +pub fn variant_field(name: impl Into, nullable: bool) -> Field { + Field::new_struct( + name, + [ + Arc::new(Field::new("metadata", DataType::BinaryView, false)), + Arc::new(Field::new("value", DataType::BinaryView, false)), + ], + nullable, + ) + .with_extension_type(VariantType) +} + +#[cfg(test)] +pub(crate) fn json_values_to_variant(values: &[Option]) -> Result { + let mut builder = VariantArrayBuilder::new(values.len()); + for value in values { + match value { + Some(value) => append_json_value(&mut builder, value).context(ArrowComputeSnafu)?, + None => builder.append_null(), + } + } + Ok(ArrayRef::from(builder.build())) +} + +/// Encodes JSON variants as an unshredded Parquet Variant array. +#[cfg(test)] +fn json_variants_to_variant(values: &[Option]) -> Result { + let mut builder = VariantArrayBuilder::new(values.len()); + for value in values { + match value { + Some(value) => append_json_variant(&mut builder, value).context(ArrowComputeSnafu)?, + None => builder.append_null(), + } + } + Ok(ArrayRef::from(builder.build())) +} + +#[cfg(test)] +fn append_json_variant( + builder: &mut impl VariantBuilderExt, + value: &JsonVariant, +) -> std::result::Result<(), ArrowError> { + match value { + JsonVariant::Null => builder.append_value(Variant::Null), + JsonVariant::Bool(value) => builder.append_value(*value), + JsonVariant::Number(JsonNumber::PosInt(value)) => { + if let Ok(value) = i64::try_from(*value) { + builder.append_value(value); + } else { + append_large_u64(builder, *value)?; + } + } + JsonVariant::Number(JsonNumber::NegInt(value)) => builder.append_value(*value), + JsonVariant::Number(JsonNumber::Float(value)) => { + if value.0.is_finite() { + builder.append_value(value.0) + } else { + builder.append_value("NaN") + } + } + JsonVariant::String(value) => builder.append_value(value.as_str()), + JsonVariant::Array(values) => { + let mut list = builder.try_new_list()?; + for value in values { + append_json_variant(&mut list, value)?; + } + list.finish(); + } + JsonVariant::Object(values) => { + let mut object = builder.try_new_object()?; + for (name, value) in values { + append_json_variant(&mut ObjectFieldBuilder::new(name, &mut object), value)?; + } + object.finish(); + } + JsonVariant::Variant(value) => { + let value = decode_json_variant(value) + .map_err(|e| ArrowError::JsonError(format!("Failed to decode JSONB: {e}")))?; + append_json_value(builder, &value)?; + } + } + Ok(()) +} + +#[cfg(test)] +fn append_json_value( + builder: &mut impl VariantBuilderExt, + value: &serde_json::Value, +) -> std::result::Result<(), ArrowError> { + match value { + serde_json::Value::Null => builder.append_value(Variant::Null), + serde_json::Value::Bool(value) => builder.append_value(*value), + serde_json::Value::Number(value) => { + if let Some(value) = value.as_i64() { + builder.append_value(value); + } else if let Some(value) = value.as_u64() { + append_large_u64(builder, value)?; + } else if let Some(value) = value.as_f64() { + builder.append_value(value); + } else { + return Err(ArrowError::InvalidArgumentError(format!( + "Failed to encode JSON number as Variant: {value}" + ))); + } + } + serde_json::Value::String(value) => builder.append_value(value.as_str()), + serde_json::Value::Array(values) => { + let mut list = builder.try_new_list()?; + for value in values { + append_json_value(&mut list, value)?; + } + list.finish(); + } + serde_json::Value::Object(values) => { + let mut object = builder.try_new_object()?; + for (name, value) in values { + append_json_value(&mut ObjectFieldBuilder::new(name, &mut object), value)?; + } + object.finish(); + } + } + Ok(()) +} + +/// Parquet Variant has no unsigned integer primitive. Treat u64 as i64 first, then use Decimal16 +/// to represent large (larger than i64::MAX) u64. +#[cfg(test)] +fn append_large_u64( + builder: &mut impl VariantBuilderExt, + value: u64, +) -> std::result::Result<(), ArrowError> { + let value = VariantDecimal16::try_new(value as i128, 0).map_err(|e| { + ArrowError::InvalidArgumentError(format!( + "Failed to encode JSON large integer as Variant Decimal16: {e}" + )) + })?; + builder.append_value(value); + Ok(()) +} + +/// Decodes an unshredded Parquet Variant array into JSON values. +pub fn variant_to_json_values(array: &ArrayRef) -> Result>> { + let variants = VariantArray::try_new(array.as_ref()).context(ArrowComputeSnafu)?; + (0..variants.len()) + .map(|i| { + if variants.is_null(i) { + Ok(None) + } else { + variants + .try_value(i) + .and_then(|x| x.to_json_value()) + .context(ArrowComputeSnafu) + .map(Some) + } + }) + .collect() +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn test_variant_field_matches_canonical_layout() { + let expected = VariantArrayBuilder::new(0) + .build() + .field("remainder") + .with_nullable(true); + assert_eq!(expected, variant_field("remainder", true)); + } + + #[test] + fn test_variant_json_round_trip() -> Result<()> { + let values = vec![ + None, + Some(serde_json::Value::Null), + Some(json!({})), + Some(json!({"nested": {"items": [1, "two", null]}})), + Some(json!({"unicode": "\u{503c}"})), + Some(json!({"max_u64": u64::MAX})), + ]; + + let array = json_values_to_variant(&values)?; + assert_eq!(values, variant_to_json_values(&array)?); + + let variants = values + .clone() + .into_iter() + .map(|x| x.map(JsonVariant::from)) + .collect::>(); + let array = json_variants_to_variant(&variants)?; + assert_eq!(values, variant_to_json_values(&array)?); + + let variants = [Some(JsonVariant::Variant( + jsonb::parse_value(br#"{"nested": true}"#).unwrap().to_vec(), + ))]; + let array = json_variants_to_variant(&variants)?; + assert_eq!( + vec![Some(json!({"nested": true}))], + variant_to_json_values(&array)? + ); + Ok(()) + } +} From 8dbfbee611cc9d3996830ac8aeebffeb755c196f Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Wed, 19 Aug 2026 03:47:44 +0000 Subject: [PATCH 62/92] refactor: remove open metrics parser (#8905) --- Cargo.lock | 18 ------------------ src/frontend/src/instance/prom_store.rs | 10 +--------- src/servers/Cargo.toml | 1 - src/servers/src/http/prom_store.rs | 9 --------- src/servers/src/prom_store.rs | 6 ------ src/servers/src/query_handler.rs | 3 --- src/servers/tests/http/prom_store_test.rs | 6 +----- 7 files changed, 2 insertions(+), 51 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 16b4c6b1045..d7c9241afca 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -977,12 +977,6 @@ dependencies = [ "syn 2.0.117", ] -[[package]] -name = "auto_ops" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7460f7dd8e100147b82a63afca1a20eb6c231ee36b90ba7272e14951cb58af59" - [[package]] name = "autocfg" version = "1.5.0" @@ -9802,17 +9796,6 @@ dependencies = [ "url", ] -[[package]] -name = "openmetrics-parser" -version = "0.4.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e40a68c62e09c5dfec2f6472af3bd5e8ddf506fcf14c78ece23794ffbb874eca" -dependencies = [ - "auto_ops", - "pest", - "pest_derive", -] - [[package]] name = "opensrv-mysql" version = "0.8.0" @@ -13461,7 +13444,6 @@ dependencies = [ "notify", "object-pool", "once_cell", - "openmetrics-parser", "opensrv-mysql", "opentelemetry-proto", "operator", diff --git a/src/frontend/src/instance/prom_store.rs b/src/frontend/src/instance/prom_store.rs index 428628df123..e0b17b1a89e 100644 --- a/src/frontend/src/instance/prom_store.rs +++ b/src/frontend/src/instance/prom_store.rs @@ -47,7 +47,7 @@ use servers::http::header::{CONTENT_ENCODING_SNAPPY, CONTENT_TYPE_PROTOBUF, coll use servers::http::prom_store::PHYSICAL_TABLE_PARAM; use servers::interceptor::{PromStoreProtocolInterceptor, PromStoreProtocolInterceptorRef}; use servers::pending_rows_batcher::PendingRowsSchemaAlterer; -use servers::prom_store::{self, Metrics}; +use servers::prom_store; use servers::query_handler::{ PromStoreProtocolHandler, PromStoreProtocolHandlerRef, PromStoreResponse, }; @@ -648,10 +648,6 @@ impl PromStoreProtocolHandler for Instance { .fail(), } } - - async fn ingest_metrics(&self, _metrics: Metrics) -> ServerResult<()> { - todo!(); - } } impl Instance { @@ -791,10 +787,6 @@ impl PromStoreProtocolHandler for ExportMetricHandler { ) -> ServerResult { unreachable!(); } - - async fn ingest_metrics(&self, _metrics: Metrics) -> ServerResult<()> { - unreachable!(); - } } #[cfg(test)] diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index d64a1bd5a5d..59d00b90ad7 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -87,7 +87,6 @@ mysql_common = "0.38" notify.workspace = true object-pool = "0.5" once_cell.workspace = true -openmetrics-parser = "0.4" # Wait for https://github.com/databendlabs/opensrv/pull/81 opensrv-mysql = { git = "https://github.com/GreptimeTeam/opensrv", rev = "6c5a451544194b7bb60a8318d155d4f892b49f2c" } opentelemetry-proto.workspace = true diff --git a/src/servers/src/http/prom_store.rs b/src/servers/src/http/prom_store.rs index 1f0d26789eb..cb7b6307c4b 100644 --- a/src/servers/src/http/prom_store.rs +++ b/src/servers/src/http/prom_store.rs @@ -770,7 +770,6 @@ mod tests { use super::*; use crate::prom_remote_write::validation::PromValidationMode; - use crate::prom_store::Metrics; use crate::query_handler::PromStoreProtocolHandler; #[test] @@ -966,10 +965,6 @@ mod tests { ) -> Result { unimplemented!() } - - async fn ingest_metrics(&self, _metrics: Metrics) -> Result<()> { - unimplemented!() - } } #[tokio::test] @@ -1067,9 +1062,5 @@ mod tests { ) -> Result { unimplemented!() } - - async fn ingest_metrics(&self, _metrics: Metrics) -> Result<()> { - unimplemented!() - } } } diff --git a/src/servers/src/prom_store.rs b/src/servers/src/prom_store.rs index c4f2f929cae..980504515a8 100644 --- a/src/servers/src/prom_store.rs +++ b/src/servers/src/prom_store.rs @@ -34,7 +34,6 @@ use datafusion::dataframe::DataFrame; use datafusion::prelude::{Expr, col, lit, regexp_match}; use datafusion_common::ScalarValue; use datafusion_expr::LogicalPlan; -use openmetrics_parser::{MetricsExposition, PrometheusType, PrometheusValue}; use snafu::{OptionExt, ResultExt, ensure}; use snap::raw::{Decoder, Encoder}; @@ -95,11 +94,6 @@ pub fn is_physical_table_selection_label(label: &str) -> bool { label == PHYSICAL_TABLE_LABEL || label == PHYSICAL_TABLE_LABEL_ALT } -/// Metrics for push gateway protocol -pub struct Metrics { - pub exposition: MetricsExposition, -} - /// Get table name from remote query pub fn table_name(q: &Query) -> Result { let mut matchers = q diff --git a/src/servers/src/query_handler.rs b/src/servers/src/query_handler.rs index a6a086e0f23..96f412ac1e1 100644 --- a/src/servers/src/query_handler.rs +++ b/src/servers/src/query_handler.rs @@ -54,7 +54,6 @@ use crate::error::Result; use crate::http::jaeger::QueryTraceParams; use crate::influxdb::InfluxdbRequest; use crate::opentsdb::codec::DataPoint; -use crate::prom_store::Metrics; pub type OpentsdbProtocolHandlerRef = Arc; pub type InfluxdbLineProtocolHandlerRef = Arc; pub type PromStoreProtocolHandlerRef = Arc; @@ -125,8 +124,6 @@ pub trait PromStoreProtocolHandler { /// Handling prometheus remote read requests async fn read(&self, request: ReadRequest, ctx: QueryContextRef) -> Result; - /// Handling push gateway requests - async fn ingest_metrics(&self, metrics: Metrics) -> Result<()>; } #[async_trait] diff --git a/src/servers/tests/http/prom_store_test.rs b/src/servers/tests/http/prom_store_test.rs index 6e8658f663f..8f3e2b0e770 100644 --- a/src/servers/tests/http/prom_store_test.rs +++ b/src/servers/tests/http/prom_store_test.rs @@ -44,7 +44,7 @@ use servers::http::{HttpOptions, HttpServerBuilder}; use servers::prom_remote_write::v2::test_util as remote_write_v2; use servers::prom_remote_write::validation::PromValidationMode; use servers::prom_store; -use servers::prom_store::{Metrics, snappy_compress}; +use servers::prom_store::snappy_compress; use servers::query_handler::sql::SqlQueryHandler; use servers::query_handler::{PromStoreProtocolHandler, PromStoreResponse}; use session::context::QueryContextRef; @@ -146,10 +146,6 @@ impl PromStoreProtocolHandler for DummyInstance { body: response.encode_to_vec(), }) } - - async fn ingest_metrics(&self, _metrics: Metrics) -> Result<()> { - unimplemented!(); - } } #[async_trait] From 3de4feddd5a5c6a282d5b2eae7541937de6bcdf8 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Wed, 19 Aug 2026 04:32:23 +0000 Subject: [PATCH 63/92] feat(otlp): report the cause of rejected trace spans (#8897) * feat(otlp): report the cause of rejected trace spans When trace-v1 ingestion cannot coerce an attribute value, it falls back to single-span writes and rejects the bad span. That behavior is correct, but the OTLP partial-success message only carried `Rejected span : (InvalidArguments)`: the column, the source value, the source type and the target type were all dropped, so locating the bad attribute required adding a detailed exporter on the collector side and replaying traffic. Two places lost the information. `prepare_trace_column_rewrites` built a message without the failing value, and the span rejection path kept only the status code from the error. Coercion errors now name the failing value, e.g. failed to coerce trace column 'span_attributes.http.response.body.size' in table 'opentelemetry_traces' from String("") to Int64 and the rejection detail carries that cause. Values are user data, so a string keeps at most 16 characters, is escaped, and binary payloads report only their length; the cause itself is bounded at 256 characters. Both truncations cut on a char boundary. Failure details now deduplicate: repeats of the same (site, cause) collapse into one entry with an occurrence count, keyed on the untruncated cause so two failures that differ past the display limit stay separate. Only four distinct entries are retained and the rest are counted, which keeps the state bounded no matter how many distinct bad values a request carries. A fully rejected request is logged at warn level and a partial success at debug level, since the latter repeats every export interval; the detail goes out as a Debug field so a newline in an attribute key cannot forge log lines. Rejection semantics are unchanged: partial success, same accepted and rejected counts, same HTTP status mapping. Signed-off-by: Dennis Zhuang * refactor(otlp): compare failure keys directly instead of hashing The dedup identity was a DefaultHasher fingerprint of `(label, key)`, which bought a fixed 8 bytes per entry at the cost of an import, four lines, and a collision argument the reader has to make. Entries are capped at four and a cause runs a couple of hundred characters, so the saving is about a kilobyte per in-flight request while the column name it avoids retaining is already held several times over by the request itself. Compare the strings instead, keeping the untruncated cause as the key so failures differing past the display limit still stay apart. Labels are metric label values and always static, so the entry borrows them. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/frontend/src/instance/otlp/README.md | 12 ++ .../src/instance/otlp/trace_ingest.rs | 172 ++++++++++++++++-- .../src/instance/otlp/trace_ingest/tests.rs | 151 ++++++++++++--- src/frontend/src/instance/otlp/trace_types.rs | 153 ++++++++++++++-- tests-integration/tests/http.rs | 9 + 5 files changed, 436 insertions(+), 61 deletions(-) diff --git a/src/frontend/src/instance/otlp/README.md b/src/frontend/src/instance/otlp/README.md index 8db27384183..4489ab0c3f3 100644 --- a/src/frontend/src/instance/otlp/README.md +++ b/src/frontend/src/instance/otlp/README.md @@ -158,6 +158,18 @@ Accounting follows the main-table write: accepted or rejected span counts; and - failure details are bounded before they are folded into `TraceIngestOutcome`. +A rejection detail carries the failing cause, not just a status code, so an +unusable attribute value names its column, source value, and target type. Repeats +of the same `(site, cause)` collapse into one entry with an occurrence count, and +only a fixed number of distinct entries is retained; failures past that limit are +counted but their causes are dropped, which keeps the state bounded regardless of +how many distinct bad values a request carries. + +A request that rejects every span is logged at warn level. A partial success is +logged at debug level only: it repeats every export interval while the sender +keeps emitting the value, and its detail already reaches the sender through the +OTLP partial-success message. + Finally, the HTTP handler returns: - full success when there are no rejected spans or failure details; diff --git a/src/frontend/src/instance/otlp/trace_ingest.rs b/src/frontend/src/instance/otlp/trace_ingest.rs index f6ce60b2d11..671f6c6f7cf 100644 --- a/src/frontend/src/instance/otlp/trace_ingest.rs +++ b/src/frontend/src/instance/otlp/trace_ingest.rs @@ -25,7 +25,7 @@ use client::Output; use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; use common_meta::rpc::ddl::TriggerReason; -use common_telemetry::warn; +use common_telemetry::{debug, warn}; use datatypes::prelude::ConcreteDataType; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; use pipeline::{GreptimePipelineParams, PipelineWay}; @@ -50,11 +50,15 @@ use crate::instance::otlp::trace_types::{ PendingTraceColumnRewrite, PreparedTraceColumnRewrites, TraceColumnRewriteError, choose_trace_reconcile_decision, enrich_trace_reconcile_error, is_trace_reconcile_candidate_type, prepare_trace_column_rewrites, push_observed_trace_type, + truncate_for_diagnostics, }; use crate::metrics::{OTLP_TRACES_FAILURE_COUNT, OTLP_TRACES_ROWS}; const TRACE_FAILURE_MESSAGE_LIMIT: usize = 4; +/// Maximum characters of a failure cause echoed to the client and the log. +const TRACE_FAILURE_CAUSE_LIMIT: usize = 256; + /// Determines how trace ingestion responds to a failure before a write is dispatched. #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ChunkFailureReaction { @@ -86,7 +90,34 @@ struct TraceChunkIngestContext<'a> { struct TraceIngestState { aux_data: TraceAuxData, outcome: TraceIngestOutcome, - failure_messages: Vec, + failure_messages: TraceFailureMessages, +} + +/// Bounded, deduplicated failure details for one trace request. +/// +/// Occurrences past [`TRACE_FAILURE_MESSAGE_LIMIT`] distinct failures are +/// counted but their keys are dropped: retaining them would let this state grow +/// with the number of distinct bad values in a request. +#[derive(Debug, Default)] +struct TraceFailureMessages { + entries: Vec, + suppressed_occurrences: usize, +} + +#[derive(Debug)] +struct TraceFailureEntry { + label: &'static str, + /// Untruncated: two causes can agree on a truncated prefix and differ + /// exactly where the actionable detail is. + key: String, + message: String, + occurrences: usize, +} + +impl TraceFailureMessages { + fn is_empty(&self) -> bool { + self.entries.is_empty() + } } /// How a v1 chunk should be reconciled when it is written. @@ -863,7 +894,7 @@ impl Instance { let mut ingest_state = TraceIngestState { aux_data: TraceAuxData::default(), outcome: TraceIngestOutcome::default(), - failure_messages: Vec::new(), + failure_messages: TraceFailureMessages::default(), }; let main_result: ServerResult<()> = async { @@ -954,6 +985,7 @@ impl Instance { Self::push_trace_failure_message( &mut ingest_state.failure_messages, "aux_table_update_failed", + err.status_code().as_ref(), format!( "Auxiliary trace tables were not fully updated ({})", err.status_code().as_ref() @@ -968,6 +1000,7 @@ impl Instance { Err(err) => Self::push_trace_failure_message( &mut ingest_state.failure_messages, "aux_table_update_failed", + err.status_code().as_ref(), format!( "Auxiliary trace tables were not fully updated ({})", err.status_code().as_ref() @@ -985,6 +1018,35 @@ impl Instance { ingest_state.failure_messages, ); + if let Some(error_message) = &ingest_state.outcome.error_message { + let accepted_spans = ingest_state.outcome.accepted_spans; + let rejected_spans = ingest_state.outcome.rejected_spans; + // A partial success repeats every export interval while the sender + // keeps emitting the bad value, so only a fully rejected request is + // worth a warning. Either way the detail reaches the sender: as an + // OTLP partial success, or as the status message of a 400. + // + // The detail embeds attribute keys verbatim, so it goes out as a + // Debug field; interpolating it would let a newline in an attribute + // key forge log lines. + if accepted_spans == 0 && rejected_spans > 0 { + warn!( + table_name = ingest_ctx.table_name, + rejected_spans, + error_message = ?error_message, + "OTLP trace ingest rejected every span" + ); + } else { + debug!( + table_name = ingest_ctx.table_name, + accepted_spans, + rejected_spans, + error_message = ?error_message, + "OTLP trace ingest reported failures" + ); + } + } + Ok(ingest_state.outcome) } @@ -1019,6 +1081,7 @@ impl Instance { Self::push_trace_failure_message( &mut ingest_state.failure_messages, ChunkFailureReaction::Propagate.as_metric_label(), + err.status_code().as_ref(), format!( "Propagating chunk write failure ({})", err.status_code().as_ref() @@ -1069,13 +1132,18 @@ impl Instance { // recover valid data. let span_count = chunks.iter().map(|chunk| chunk.rows.len()).sum::(); ingest_state.outcome.rejected_spans += span_count; + let (cause, shown_cause) = Self::trace_failure_cause(&err); Self::push_trace_failure_message( &mut ingest_state.failure_messages, ChunkFailureReaction::DiscardChunk.as_metric_label(), + // Merging discards of different sizes would make the reported + // count times its occurrences wrong. + &format!("{span_count}:{cause}"), format!( - "Discarded {} spans after pre-write request failure ({})", + "Discarded {} spans after pre-write request failure ({}): {}", span_count, - err.status_code().as_ref() + err.status_code().as_ref(), + shown_cause ), ); return Ok(()); @@ -1110,6 +1178,7 @@ impl Instance { Self::push_trace_failure_message( &mut ingest_state.failure_messages, ChunkFailureReaction::RetryPerSpan.as_metric_label(), + "incompatible_binary_and_json", "Chunk fallback triggered by incompatible binary and JSON values".to_string(), ); return self @@ -1136,6 +1205,7 @@ impl Instance { Self::push_trace_failure_message( &mut ingest_state.failure_messages, ChunkFailureReaction::RetryPerSpan.as_metric_label(), + err.status_code().as_ref(), format!("Chunk fallback triggered by {}", err.status_code().as_ref()), ); return self @@ -1144,13 +1214,19 @@ impl Instance { } ChunkFailureReaction::DiscardChunk => { ingest_state.outcome.rejected_spans += span_count; + let (cause, shown_cause) = Self::trace_failure_cause(&err); Self::push_trace_failure_message( &mut ingest_state.failure_messages, ChunkFailureReaction::DiscardChunk.as_metric_label(), + // Chunks discarded for one cause differ in size; merging + // them would make the reported count times its + // occurrences wrong. + &format!("{span_count}:{cause}"), format!( - "Discarded {} spans after pre-write chunk failure ({})", + "Discarded {} spans after pre-write chunk failure ({}): {}", span_count, - err.status_code().as_ref() + err.status_code().as_ref(), + shown_cause ), ); return Ok(()); @@ -1159,6 +1235,7 @@ impl Instance { Self::push_trace_failure_message( &mut ingest_state.failure_messages, ChunkFailureReaction::Propagate.as_metric_label(), + err.status_code().as_ref(), format!( "Propagating pre-write chunk failure ({})", err.status_code().as_ref() @@ -1182,6 +1259,7 @@ impl Instance { Self::push_trace_failure_message( &mut ingest_state.failure_messages, ChunkFailureReaction::Propagate.as_metric_label(), + err.status_code().as_ref(), format!( "Propagating chunk write failure ({})", err.status_code().as_ref() @@ -1221,6 +1299,7 @@ impl Instance { Self::push_trace_failure_message( &mut ingest_state.failure_messages, ChunkFailureReaction::Propagate.as_metric_label(), + err.status_code().as_ref(), format!( "Propagating pre-write span failure for {}:{} ({})", span.trace_id, @@ -1232,14 +1311,19 @@ impl Instance { } ingest_state.outcome.rejected_spans += 1; + // Dedup on the cause, not the span id: one bad column rejects + // every span and would otherwise fill the entry list. + let (cause, shown_cause) = Self::trace_failure_cause(&err); Self::push_trace_failure_message( &mut ingest_state.failure_messages, "span_rejected", + &cause, format!( - "Rejected span {}:{} ({})", + "Rejected span {}:{} ({}): {}", span.trace_id, span.span_id, - err.status_code().as_ref() + err.status_code().as_ref(), + shown_cause ), ); continue; @@ -1256,6 +1340,7 @@ impl Instance { Self::push_trace_failure_message( &mut ingest_state.failure_messages, ChunkFailureReaction::Propagate.as_metric_label(), + err.status_code().as_ref(), format!( "Propagating span write failure for {}:{} ({})", span.trace_id, @@ -1619,24 +1704,52 @@ impl Instance { outcome.write_cost += cost; } - fn push_trace_failure_message(messages: &mut Vec, label: &str, message: String) { + /// Returns the full cause of a pre-write failure and its display form. + /// + /// `output_msg` masks internal errors and unwraps the root cause. Dedup must + /// key on the full text: two causes can agree on a truncated prefix and + /// differ exactly where the actionable detail is. + fn trace_failure_cause(err: &error::Error) -> (String, String) { + let cause = err.output_msg(); + let display = truncate_for_diagnostics(&cause, TRACE_FAILURE_CAUSE_LIMIT); + (cause, display) + } + + /// Records one failure, merging repeats of `(label, key)` into a count. + fn push_trace_failure_message( + messages: &mut TraceFailureMessages, + label: &'static str, + key: &str, + message: String, + ) { OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).inc(); - if messages.len() < TRACE_FAILURE_MESSAGE_LIMIT { - messages.push(message); - } else if messages.len() == TRACE_FAILURE_MESSAGE_LIMIT { - tracing::debug!( - label, - limit = TRACE_FAILURE_MESSAGE_LIMIT, - "Trace ingest failure message limit reached; suppressing additional failure details" - ); + if let Some(entry) = messages + .entries + .iter_mut() + .find(|entry| entry.label == label && entry.key == key) + { + entry.occurrences += 1; + return; } + + if messages.entries.len() >= TRACE_FAILURE_MESSAGE_LIMIT { + messages.suppressed_occurrences += 1; + return; + } + + messages.entries.push(TraceFailureEntry { + label, + key: key.to_string(), + message, + occurrences: 1, + }); } fn finish_trace_failure_message( accepted_spans: usize, rejected_spans: usize, - messages: Vec, + messages: TraceFailureMessages, ) -> Option { if rejected_spans == 0 && messages.is_empty() { return None; @@ -1648,8 +1761,27 @@ impl Instance { ); if !messages.is_empty() { + let details = messages + .entries + .into_iter() + .map(|entry| { + if entry.occurrences > 1 { + format!("{} (x{})", entry.message, entry.occurrences) + } else { + entry.message + } + }) + .collect::>() + .join("; "); summary.push_str(": "); - summary.push_str(&messages.join("; ")); + summary.push_str(&details); + } + + if messages.suppressed_occurrences > 0 { + summary.push_str(&format!( + "; {} additional failures suppressed", + messages.suppressed_occurrences + )); } Some(summary) diff --git a/src/frontend/src/instance/otlp/trace_ingest/tests.rs b/src/frontend/src/instance/otlp/trace_ingest/tests.rs index b3a0570db51..ec5bad89754 100644 --- a/src/frontend/src/instance/otlp/trace_ingest/tests.rs +++ b/src/frontend/src/instance/otlp/trace_ingest/tests.rs @@ -34,8 +34,8 @@ use servers::otlp::trace::span::{SpanEvents, SpanLinks, TraceSpan}; use servers::otlp::trace::v1::{TraceBinaryType, TraceRetryColumn}; use super::{ - ChunkFailureReaction, Instance, TraceChunkRetry, TraceChunkSchemaState, TraceRequestSchema, - TraceRequestSchemaPlan, TraceSpanMetadata, TraceTablePreAlter, chunk_owned, + ChunkFailureReaction, Instance, TraceChunkRetry, TraceChunkSchemaState, TraceFailureMessages, + TraceRequestSchema, TraceRequestSchemaPlan, TraceSpanMetadata, TraceTablePreAlter, chunk_owned, wrap_trace_alter_failure, }; use crate::metrics::OTLP_TRACES_FAILURE_COUNT; @@ -102,22 +102,28 @@ fn test_classify_trace_prewrite_failure() { #[test] fn test_finish_trace_failure_message() { - let message = Instance::finish_trace_failure_message( - 3, - 2, - vec!["Rejected span trace:span (InvalidArguments)".to_string()], - ) - .unwrap(); + let mut messages = TraceFailureMessages::default(); + Instance::push_trace_failure_message( + &mut messages, + "span_rejected", + "coercion", + "Rejected span trace:span (InvalidArguments)".to_string(), + ); + + let message = Instance::finish_trace_failure_message(3, 2, messages).unwrap(); assert!(message.contains("Accepted 3 spans, rejected 2 spans")); assert!(message.contains("Rejected span trace:span")); - assert_eq!(Instance::finish_trace_failure_message(2, 0, vec![]), None); + assert_eq!( + Instance::finish_trace_failure_message(2, 0, TraceFailureMessages::default()), + None + ); } #[test] fn test_finish_trace_failure_message_without_detail_messages() { assert_eq!( - Instance::finish_trace_failure_message(0, 2, vec![]), + Instance::finish_trace_failure_message(0, 2, TraceFailureMessages::default()), Some("Accepted 0 spans, rejected 2 spans".to_string()) ); } @@ -126,39 +132,138 @@ fn test_finish_trace_failure_message_without_detail_messages() { fn test_push_trace_failure_message_increments_labeled_counter() { let label = "retry_per_span_counter_test"; let initial = OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).get(); - let mut messages = Vec::new(); + let mut messages = TraceFailureMessages::default(); Instance::push_trace_failure_message( &mut messages, label, + "InvalidArguments", "Chunk fallback triggered by InvalidArguments".to_string(), ); - assert_eq!(messages.len(), 1); + assert_eq!(messages.entries.len(), 1); assert_eq!( OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).get(), initial + 1 ); } +/// One bad column rejects every span, so repeats must collapse instead of +/// filling the bounded entry list with the same cause. +#[test] +fn test_push_trace_failure_message_collapses_repeated_cause() { + let label = "span_rejected_dedup_test"; + let initial = OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).get(); + let mut messages = TraceFailureMessages::default(); + + for idx in 0..3 { + Instance::push_trace_failure_message( + &mut messages, + label, + "failed to coerce", + format!("Rejected span trace:span-{idx} (InvalidArguments): failed to coerce"), + ); + } + + assert_eq!(messages.entries.len(), 1); + assert_eq!(messages.entries[0].occurrences, 3); + assert_eq!(messages.suppressed_occurrences, 0); + // Every occurrence is still metered even though the details collapse. + assert_eq!( + OTLP_TRACES_FAILURE_COUNT.with_label_values(&[label]).get(), + initial + 3 + ); + + let summary = Instance::finish_trace_failure_message(0, 3, messages).unwrap(); + assert!( + summary.contains("Rejected span trace:span-0 (InvalidArguments): failed to coerce (x3)"), + "unexpected summary: {summary}" + ); +} + +/// Dedup identity comes from the full cause, so two failures that only differ +/// past the display limit must stay separate instead of one silently vanishing. +#[test] +fn test_trace_failure_cause_keeps_identity_beyond_the_display_limit() { + let shared_prefix = "x".repeat(400); + let (first_cause, first_shown) = Instance::trace_failure_cause( + &servers::error::InvalidParameterSnafu { + reason: format!("{shared_prefix}-int64"), + } + .build(), + ); + let (second_cause, second_shown) = Instance::trace_failure_cause( + &servers::error::InvalidParameterSnafu { + reason: format!("{shared_prefix}-float64"), + } + .build(), + ); + + assert_eq!(first_shown, second_shown); + assert_ne!(first_cause, second_cause); + + let mut messages = TraceFailureMessages::default(); + Instance::push_trace_failure_message(&mut messages, "span_rejected", &first_cause, first_shown); + Instance::push_trace_failure_message( + &mut messages, + "span_rejected", + &second_cause, + second_shown, + ); + + assert_eq!(messages.entries.len(), 2); +} + +/// The same cause reported by two different sites must stay distinguishable. +#[test] +fn test_push_trace_failure_message_separates_labels_sharing_a_key() { + let mut messages = TraceFailureMessages::default(); + + Instance::push_trace_failure_message( + &mut messages, + "span_rejected", + "failed to coerce", + "Rejected span trace:span (InvalidArguments): failed to coerce".to_string(), + ); + Instance::push_trace_failure_message( + &mut messages, + "discard_chunk", + "failed to coerce", + "Discarded 7 spans after pre-write chunk failure (InvalidArguments): failed to coerce" + .to_string(), + ); + + assert_eq!(messages.entries.len(), 2); +} + #[test] fn test_push_trace_failure_message_caps_recorded_messages() { let label = "retry_per_span_limit_test"; - let mut messages = Vec::new(); + let mut messages = TraceFailureMessages::default(); - for idx in 0..=4 { - Instance::push_trace_failure_message(&mut messages, label, format!("failure-{idx}")); + for idx in 0..=5 { + Instance::push_trace_failure_message( + &mut messages, + label, + &format!("cause-{idx}"), + format!("failure-{idx}"), + ); } - assert_eq!(messages.len(), 4); assert_eq!( - messages, - vec![ - "failure-0".to_string(), - "failure-1".to_string(), - "failure-2".to_string(), - "failure-3".to_string() - ] + messages + .entries + .iter() + .map(|entry| entry.message.as_str()) + .collect::>(), + vec!["failure-0", "failure-1", "failure-2", "failure-3"] + ); + assert_eq!(messages.suppressed_occurrences, 2); + + let summary = Instance::finish_trace_failure_message(0, 6, messages).unwrap(); + assert!( + summary.ends_with("; 2 additional failures suppressed"), + "unexpected summary: {summary}" ); } diff --git a/src/frontend/src/instance/otlp/trace_types.rs b/src/frontend/src/instance/otlp/trace_types.rs index 0e2df25053e..603cc66c0d1 100644 --- a/src/frontend/src/instance/otlp/trace_types.rs +++ b/src/frontend/src/instance/otlp/trace_types.rs @@ -22,6 +22,39 @@ use servers::otlp::coerce::{ use crate::instance::otlp::trace_semconv::trace_semconv_fixed_type; +/// Attribute values are user data echoed back in the OTLP partial-success +/// message and the server log, so diagnostics keep at most this many characters. +const TRACE_VALUE_DIAGNOSTIC_LIMIT: usize = 16; + +/// Truncates to `limit` characters, marking the cut with `...`. +/// +/// Diagnostics carry user-controlled text; slicing by byte offset would panic in +/// the middle of a multi-byte sequence. +pub(super) fn truncate_for_diagnostics(text: &str, limit: usize) -> String { + match text.char_indices().nth(limit) { + Some((offset, _)) => format!("{}...", &text[..offset]), + None => text.to_string(), + } +} + +/// Renders a failing trace value as `Type(value)`, e.g. `String("")`. +fn describe_trace_value(value: &ValueData, request_type: ColumnDataType) -> String { + let payload = match value { + ValueData::StringValue(string_value) => format!( + "{:?}", + truncate_for_diagnostics(string_value, TRACE_VALUE_DIAGNOSTIC_LIMIT) + ), + ValueData::BoolValue(bool_value) => bool_value.to_string(), + ValueData::I64Value(int_value) => int_value.to_string(), + ValueData::F64Value(float_value) => float_value.to_string(), + ValueData::BinaryValue(bytes) => format!("{} bytes", bytes.len()), + // Other value kinds never reach trace coercion, so report the type alone. + _ => return format!("{request_type:?}"), + }; + + format!("{request_type:?}({payload})") +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(super) enum TraceReconcileDecision { UseExisting(ColumnDataType), @@ -191,8 +224,10 @@ pub(super) fn prepare_trace_column_rewrites( let Some(value) = row.values.get(pending_rewrite.col_idx) else { continue; }; - let Some(request_type) = value.value_data.as_ref().and_then(trace_value_datatype) - else { + let Some(request_value) = value.value_data.as_ref() else { + continue; + }; + let Some(request_type) = trace_value_datatype(request_value) else { continue; }; if request_type == pending_rewrite.target_type { @@ -201,20 +236,18 @@ pub(super) fn prepare_trace_column_rewrites( let value_data = coerce_value_data(&value.value_data, pending_rewrite.target_type, request_type) - .map_err(|_| { - TraceColumnRewriteError { - error: error::InvalidParameterSnafu { - reason: format!( - "failed to coerce trace column '{}' in table '{}' from {:?} to {:?}", - pending_rewrite.column_name, - table_name, - request_type, - pending_rewrite.target_type - ), - } - .build(), - column_name: pending_rewrite.column_name.clone(), - } + .map_err(|_| TraceColumnRewriteError { + error: error::InvalidParameterSnafu { + reason: format!( + "failed to coerce trace column '{}' in table '{}' from {} to {:?}", + pending_rewrite.column_name, + table_name, + describe_trace_value(request_value, request_type), + pending_rewrite.target_type + ), + } + .build(), + column_name: pending_rewrite.column_name.clone(), })?; values.push(PreparedTraceValueRewrite { row_idx, @@ -301,8 +334,8 @@ mod tests { use super::{ PendingTraceColumnRewrite, TraceReconcileDecision, choose_trace_reconcile_decision, - enrich_trace_reconcile_error, is_trace_reconcile_candidate_type, - prepare_trace_column_rewrites, push_observed_trace_type, + describe_trace_value, enrich_trace_reconcile_error, is_trace_reconcile_candidate_type, + prepare_trace_column_rewrites, push_observed_trace_type, truncate_for_diagnostics, }; #[test] @@ -458,6 +491,90 @@ mod tests { .unwrap_err(); assert_eq!(err.error.status_code(), StatusCode::InvalidArguments); assert_eq!(err.column_name, "span_attributes.attr_int"); + assert!( + err.error.to_string().contains( + "failed to coerce trace column 'span_attributes.attr_int' in table \ + 'trace_type_atomicity' from String(\"not_a_number\") to Int64" + ), + "unexpected error message: {}", + err.error + ); + } + + /// The PHP instrumentation case: an empty string must be distinguishable + /// from any other unparsable value in the reported diagnostics. + #[test] + fn test_prepare_trace_column_rewrites_reports_empty_string_value() { + let rows = vec![Row { + values: vec![Value { + value_data: Some(ValueData::StringValue(String::new())), + }], + }]; + let pending_rewrites = vec![PendingTraceColumnRewrite { + col_idx: 0, + target_type: ColumnDataType::Int64, + column_name: "span_attributes.http.response.body.size".to_string(), + }]; + + let err = prepare_trace_column_rewrites(&rows, pending_rewrites, "opentelemetry_traces") + .unwrap_err(); + assert!( + err.error.to_string().contains( + "'span_attributes.http.response.body.size' in table 'opentelemetry_traces' \ + from String(\"\") to Int64" + ), + "unexpected error message: {}", + err.error + ); + } + + #[test] + fn test_describe_trace_value_bounds_and_escapes_strings() { + assert_eq!( + describe_trace_value( + &ValueData::StringValue(String::new()), + ColumnDataType::String + ), + r#"String("")"# + ); + assert_eq!( + describe_trace_value( + &ValueData::StringValue("a\tb\"c".to_string()), + ColumnDataType::String + ), + r#"String("a\tb\"c")"# + ); + assert_eq!( + describe_trace_value( + &ValueData::StringValue("0123456789abcdefghij".to_string()), + ColumnDataType::String + ), + r#"String("0123456789abcdef...")"# + ); + } + + #[test] + fn test_describe_trace_value_omits_binary_content() { + assert_eq!( + describe_trace_value( + &ValueData::BinaryValue(vec![1_u8, 2, 3]), + ColumnDataType::Binary + ), + "Binary(3 bytes)" + ); + assert_eq!( + describe_trace_value(&ValueData::F64Value(1.5), ColumnDataType::Float64), + "Float64(1.5)" + ); + } + + /// Truncation runs over user-supplied text, so it must not split a + /// multi-byte character. + #[test] + fn test_truncate_for_diagnostics_cuts_on_char_boundary() { + assert_eq!(truncate_for_diagnostics("日本語テキスト", 3), "日本語..."); + assert_eq!(truncate_for_diagnostics("short", 16), "short"); + assert_eq!(truncate_for_diagnostics("exact", 5), "exact"); } #[test] diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 079439bf78c..90cbb8c93fe 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -7782,6 +7782,15 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ), "unexpected partial success body: {body:?}" ); + // The rejection must name the column, the failing value and the target type, + // otherwise locating the bad attribute needs a collector-side capture. + assert!( + partial_success.error_message.contains( + "failed to coerce trace column 'span_attributes.attr_int' in table \ + 'trace_type_abort' from String(\"not_a_number\") to Int64" + ), + "unexpected partial success body: {body:?}" + ); validate_data( "otlp_traces_v1_type_abort_rows", From 2182dccd9b76efc87e2e3ad7e7d355e62d9b1d03 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Wed, 19 Aug 2026 10:04:50 +0000 Subject: [PATCH 64/92] fix: cap default runtime sizes to a minimum of 2 threads (#8908) * fix: cap default runtime sizes to a minimum of 2 threads RuntimeOptions derived its default sizes directly from num_cpus. On single-core machines every runtime (global, compact, query, ingest) ended up with one worker thread, which can easily deadlock async code (e.g. block_on combined with spawn). Clamp all CPU-derived runtime sizes to at least 2 threads. Signed-off-by: Lei, HUANG * fix: init logging before runtimes so runtime options are logged The global runtimes were initialized before the global logging subscriber, so the "Creating runtime ..." info logs that carry the runtime sizes were silently dropped. Initialize logging first in all node start paths; common-telemetry has no dependency on common-runtime, so the reorder is safe. Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- config/config.md | 10 ++--- config/datanode.example.toml | 4 +- config/frontend.example.toml | 2 +- config/metasrv.example.toml | 2 +- config/standalone.example.toml | 2 +- src/cmd/src/datanode/builder.rs | 6 +-- src/cmd/src/flownode.rs | 4 +- src/cmd/src/frontend.rs | 4 +- src/cmd/src/metasrv.rs | 4 +- src/cmd/src/standalone.rs | 4 +- src/common/runtime/src/global.rs | 67 +++++++++++++++++++++++++++----- 11 files changed, 78 insertions(+), 31 deletions(-) diff --git a/config/config.md b/config/config.md index 123eaf90a20..1b18565cc08 100644 --- a/config/config.md +++ b/config/config.md @@ -25,7 +25,7 @@ | `runtime` | -- | -- | The runtime options. | | `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. | | `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute compact operations. | -| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 1). | +| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | | `http` | -- | -- | The HTTP server options. | | `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. | | `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When Prometheus pending-row batching is enabled, a nonzero timeout less than or equal to the
`prom_store.pending_rows_flush_interval` plus 1 second is adjusted to that value. | @@ -255,7 +255,7 @@ | `runtime` | -- | -- | The runtime options. | | `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. | | `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute compact operations. | -| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 1). | +| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | | `http` | -- | -- | The HTTP server options. | | `http.addr` | String | `127.0.0.1:4000` | The address to bind the HTTP server. | | `http.timeout` | String | `0s` | HTTP request timeout. Set to 0 to disable timeout.
When Prometheus pending-row batching is enabled, a nonzero timeout less than or equal to the
`prom_store.pending_rows_flush_interval` plus 1 second is adjusted to that value. | @@ -398,7 +398,7 @@ | `runtime` | -- | -- | The runtime options. | | `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. | | `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute compact operations. | -| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 1). | +| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | | `backend_tls` | -- | -- | TLS configuration for kv store backend (applicable for etcd, PostgreSQL, and MySQL backends)
When using etcd, PostgreSQL, or MySQL as metadata store, you can configure TLS here

Note: if TLS is configured in both this section and the `store_addrs` connection string, the
settings here will override the TLS settings in `store_addrs`. | | `backend_tls.mode` | String | `prefer` | TLS mode, refer to https://www.postgresql.org/docs/current/libpq-ssl.html
- "disable" - No TLS
- "prefer" (default) - Try TLS, fallback to plain
- "require" - Require TLS
- "verify_ca" - Require TLS and verify CA
- "verify_full" - Require TLS and verify hostname | | `backend_tls.cert_path` | String | `""` | Path to client certificate file (for client authentication)
Like "/path/to/client.crt" | @@ -509,8 +509,8 @@ | `runtime` | -- | -- | The runtime options. | | `runtime.global_rt_size` | Integer | `8` | The number of threads to execute the runtime for global read operations. | | `runtime.compact_rt_size` | Integer | `4` | The number of threads to execute compact operations. | -| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 1). | -| `runtime.query_rt_size` | Integer | `7` | The number of threads to execute datanode query operations.
Defaults to max(num_cpus - 1, 1). | +| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.
Defaults to max(num_cpus / 2, 2). | +| `runtime.query_rt_size` | Integer | `7` | The number of threads to execute datanode query operations.
Defaults to max(num_cpus - 1, 2). | | `runtime.ingest_rt_size` | Integer | `8` | The number of threads to execute datanode ingestion operations. | | `meta_client` | -- | -- | The metasrv client options. | | `meta_client.metasrv_addrs` | Array | -- | The addresses of the metasrv. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 1f1377d91cc..5966405452a 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -84,10 +84,10 @@ watch = false ## The number of threads to execute compact operations. #+ compact_rt_size = 4 ## The maximum number of blocking threads for compact operations. -## Defaults to max(num_cpus / 2, 1). +## Defaults to max(num_cpus / 2, 2). #+ compact_rt_max_blocking_threads = 4 ## The number of threads to execute datanode query operations. -## Defaults to max(num_cpus - 1, 1). +## Defaults to max(num_cpus - 1, 2). #+ query_rt_size = 7 ## The number of threads to execute datanode ingestion operations. #+ ingest_rt_size = 8 diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 260984a4ab3..ec1f9c89b42 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -44,7 +44,7 @@ default_column_prefix = "greptime" ## The number of threads to execute compact operations. #+ compact_rt_size = 4 ## The maximum number of blocking threads for compact operations. -## Defaults to max(num_cpus / 2, 1). +## Defaults to max(num_cpus / 2, 2). #+ compact_rt_max_blocking_threads = 4 ## The HTTP server options. diff --git a/config/metasrv.example.toml b/config/metasrv.example.toml index f607211d636..6017bfcc49a 100644 --- a/config/metasrv.example.toml +++ b/config/metasrv.example.toml @@ -92,7 +92,7 @@ node_max_idle_time = "24hours" ## The number of threads to execute compact operations. #+ compact_rt_size = 4 ## The maximum number of blocking threads for compact operations. -## Defaults to max(num_cpus / 2, 1). +## Defaults to max(num_cpus / 2, 2). #+ compact_rt_max_blocking_threads = 4 ## TLS configuration for kv store backend (applicable for etcd, PostgreSQL, and MySQL backends) diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 068025e2a28..5ada605b92b 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -57,7 +57,7 @@ max_concurrent_queries = 0 ## The number of threads to execute compact operations. #+ compact_rt_size = 4 ## The maximum number of blocking threads for compact operations. -## Defaults to max(num_cpus / 2, 1). +## Defaults to max(num_cpus / 2, 2). #+ compact_rt_max_blocking_threads = 4 ## The HTTP server options. diff --git a/src/cmd/src/datanode/builder.rs b/src/cmd/src/datanode/builder.rs index afabcb85f10..3dfa03e52e5 100644 --- a/src/cmd/src/datanode/builder.rs +++ b/src/cmd/src/datanode/builder.rs @@ -61,9 +61,6 @@ impl InstanceBuilder { } async fn init(opts: &mut DatanodeOptions, plugins: &mut Plugins) -> Result> { - common_runtime::init_global_runtimes(&opts.runtime); - common_runtime::init_datanode_runtimes(&opts.runtime); - let dn_opts = &mut opts.component; let guard = common_telemetry::init_global_logging( APP_NAME, @@ -73,6 +70,9 @@ impl InstanceBuilder { None, ); + common_runtime::init_global_runtimes(&opts.runtime); + common_runtime::init_datanode_runtimes(&opts.runtime); + crate::options::flush_dropped_plugin_warnings(); log_versions(verbose_version(), short_version(), APP_NAME); maybe_activate_heap_profile(&dn_opts.memory); diff --git a/src/cmd/src/flownode.rs b/src/cmd/src/flownode.rs index ee071875c3a..eccfaa63b07 100644 --- a/src/cmd/src/flownode.rs +++ b/src/cmd/src/flownode.rs @@ -248,8 +248,6 @@ impl StartCommand { } async fn build(&self, opts: FlownodeOptions) -> Result { - common_runtime::init_global_runtimes(&opts.runtime); - let guard = common_telemetry::init_global_logging( APP_NAME, &opts.component.logging, @@ -258,6 +256,8 @@ impl StartCommand { None, ); + common_runtime::init_global_runtimes(&opts.runtime); + crate::options::flush_dropped_plugin_warnings(); log_versions(verbose_version(), short_version(), APP_NAME); maybe_activate_heap_profile(&opts.component.memory); diff --git a/src/cmd/src/frontend.rs b/src/cmd/src/frontend.rs index 9d5d0bab1ee..e3737088e8d 100644 --- a/src/cmd/src/frontend.rs +++ b/src/cmd/src/frontend.rs @@ -330,8 +330,6 @@ impl StartCommand { } async fn build(&self, opts: FrontendOptions) -> Result { - common_runtime::init_global_runtimes(&opts.runtime); - let guard = common_telemetry::init_global_logging( APP_NAME, &opts.component.logging, @@ -340,6 +338,8 @@ impl StartCommand { Some(&opts.component.slow_query), ); + common_runtime::init_global_runtimes(&opts.runtime); + crate::options::flush_dropped_plugin_warnings(); log_versions(verbose_version(), short_version(), APP_NAME); maybe_activate_heap_profile(&opts.component.memory); diff --git a/src/cmd/src/metasrv.rs b/src/cmd/src/metasrv.rs index 791efeb6e86..369f5c428d1 100644 --- a/src/cmd/src/metasrv.rs +++ b/src/cmd/src/metasrv.rs @@ -318,8 +318,6 @@ impl StartCommand { } pub async fn build(&self, opts: MetasrvOptions) -> Result { - common_runtime::init_global_runtimes(&opts.runtime); - let guard = common_telemetry::init_global_logging( APP_NAME, &opts.component.logging, @@ -328,6 +326,8 @@ impl StartCommand { None, ); + common_runtime::init_global_runtimes(&opts.runtime); + crate::options::flush_dropped_plugin_warnings(); log_versions(verbose_version(), short_version(), APP_NAME); maybe_activate_heap_profile(&opts.component.memory); diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 9f8155b9cc2..2ffcc5e729c 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -395,8 +395,6 @@ impl StartCommand { #[allow(clippy::diverging_sub_expression)] /// Build GreptimeDB instance with the loaded options. pub async fn build(&self, opts: GreptimeOptions) -> Result { - common_runtime::init_global_runtimes(&opts.runtime); - let guard = common_telemetry::init_global_logging( APP_NAME, &opts.component.logging, @@ -405,6 +403,8 @@ impl StartCommand { Some(&opts.component.slow_query), ); + common_runtime::init_global_runtimes(&opts.runtime); + crate::options::flush_dropped_plugin_warnings(); log_versions(verbose_version(), short_version(), APP_NAME); maybe_activate_heap_profile(&opts.component.memory); diff --git a/src/common/runtime/src/global.rs b/src/common/runtime/src/global.rs index e28eaf61d15..827832fca22 100644 --- a/src/common/runtime/src/global.rs +++ b/src/common/runtime/src/global.rs @@ -27,6 +27,9 @@ use crate::{Builder, JoinHandle, Runtime}; const GLOBAL_WORKERS: usize = 8; const COMPACT_WORKERS: usize = 4; const HB_WORKERS: usize = 2; +/// The minimum number of worker threads for runtimes sized by CPU count. +/// A single-threaded runtime can easily deadlock in async code. +const MIN_RUNTIME_THREADS: usize = 2; /// The options for the global runtimes. #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] @@ -44,19 +47,25 @@ pub struct RuntimeOptions { pub ingest_rt_size: usize, } -impl Default for RuntimeOptions { - fn default() -> Self { - let cpus = num_cpus::get(); +impl RuntimeOptions { + fn with_num_cpus(cpus: usize) -> Self { + let cpus = usize::max(cpus, MIN_RUNTIME_THREADS); Self { global_rt_size: cpus, - compact_rt_size: usize::max(cpus / 2, 1), - compact_rt_max_blocking_threads: usize::max(cpus / 2, 1), - query_rt_size: usize::max(cpus.saturating_sub(1), 1), + compact_rt_size: usize::max(cpus / 2, MIN_RUNTIME_THREADS), + compact_rt_max_blocking_threads: usize::max(cpus / 2, MIN_RUNTIME_THREADS), + query_rt_size: usize::max(cpus.saturating_sub(1), MIN_RUNTIME_THREADS), ingest_rt_size: cpus, } } } +impl Default for RuntimeOptions { + fn default() -> Self { + Self::with_num_cpus(num_cpus::get()) + } +} + pub fn create_runtime(runtime_name: &str, thread_name: &str, worker_threads: usize) -> Runtime { info!( "Creating runtime with runtime_name: {runtime_name}, thread_name: {thread_name}, work_threads: {worker_threads}." @@ -287,18 +296,56 @@ mod tests { #[test] fn test_datanode_runtime_options_default() { let options = RuntimeOptions::default(); - let cpus = num_cpus::get(); + let cpus = usize::max(num_cpus::get(), MIN_RUNTIME_THREADS); assert_eq!(cpus, options.global_rt_size); - assert_eq!(usize::max(cpus / 2, 1), options.compact_rt_size); assert_eq!( - usize::max(cpus / 2, 1), + usize::max(cpus / 2, MIN_RUNTIME_THREADS), + options.compact_rt_size + ); + assert_eq!( + usize::max(cpus / 2, MIN_RUNTIME_THREADS), options.compact_rt_max_blocking_threads ); - assert_eq!(usize::max(cpus.saturating_sub(1), 1), options.query_rt_size); + assert_eq!( + usize::max(cpus.saturating_sub(1), MIN_RUNTIME_THREADS), + options.query_rt_size + ); assert_eq!(cpus, options.ingest_rt_size); } + #[test] + fn test_runtime_options_min_threads() { + for cpus in [0, 1, 2] { + let options = RuntimeOptions::with_num_cpus(cpus); + assert!( + options.global_rt_size >= MIN_RUNTIME_THREADS, + "global_rt_size {} < {MIN_RUNTIME_THREADS} with {cpus} cpus", + options.global_rt_size + ); + assert!( + options.compact_rt_size >= MIN_RUNTIME_THREADS, + "compact_rt_size {} < {MIN_RUNTIME_THREADS} with {cpus} cpus", + options.compact_rt_size + ); + assert!( + options.compact_rt_max_blocking_threads >= MIN_RUNTIME_THREADS, + "compact_rt_max_blocking_threads {} < {MIN_RUNTIME_THREADS} with {cpus} cpus", + options.compact_rt_max_blocking_threads + ); + assert!( + options.query_rt_size >= MIN_RUNTIME_THREADS, + "query_rt_size {} < {MIN_RUNTIME_THREADS} with {cpus} cpus", + options.query_rt_size + ); + assert!( + options.ingest_rt_size >= MIN_RUNTIME_THREADS, + "ingest_rt_size {} < {MIN_RUNTIME_THREADS} with {cpus} cpus", + options.ingest_rt_size + ); + } + } + #[test] fn test_datanode_runtimes_fallback_to_global_runtime() { let runtimes = GlobalRuntimes::new( From 8fac712870d5f95aef6ccdec8e9ae45f872f4afb Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 19 Aug 2026 12:43:42 +0000 Subject: [PATCH 65/92] fix(meta): avoid blocking runtime on stats cache lock (#8910) Signed-off-by: WenyXu --- .../src/handler/collect_stats_handler.rs | 293 +++++++++++++++++- 1 file changed, 289 insertions(+), 4 deletions(-) diff --git a/src/meta-srv/src/handler/collect_stats_handler.rs b/src/meta-srv/src/handler/collect_stats_handler.rs index d444498798f..e3ae7ebe915 100644 --- a/src/meta-srv/src/handler/collect_stats_handler.rs +++ b/src/meta-srv/src/handler/collect_stats_handler.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::cmp::Ordering; +use std::sync::Arc; use api::v1::meta::{HeartbeatRequest, Role}; use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue, Stat}; @@ -24,6 +25,7 @@ use common_meta::rpc::store::PutRequest; use common_telemetry::{error, info, warn}; use dashmap::DashMap; use snafu::ResultExt; +use tokio::sync::Mutex; use crate::error::{self, Result}; use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler}; @@ -70,7 +72,7 @@ impl EpochStats { const DEFAULT_FLUSH_STATS_FACTOR: usize = 3; pub struct CollectStatsHandler { - stats_cache: DashMap, + stats_cache: DashMap>>, flush_stats_factor: usize, } @@ -106,10 +108,16 @@ impl HeartbeatHandler for CollectStatsHandler { }; let key = current_stat.stat_key(); - let mut entry = self.stats_cache.entry(key).or_default(); + let state = { + let entry = self + .stats_cache + .entry(key) + .or_insert_with(|| Arc::new(Mutex::new(EpochStats::default()))); + Arc::clone(entry.value()) + }; + let mut epoch_stats = state.lock().await; let key: Vec = key.into(); - let epoch_stats = entry.value_mut(); let refresh = if let Some(epoch) = epoch_stats.epoch() { match current_stat.node_epoch.cmp(&epoch) { @@ -170,7 +178,7 @@ impl HeartbeatHandler for CollectStatsHandler { } } -async fn rewrite_node_address(ctx: &mut Context, stat: &Stat) { +async fn rewrite_node_address(ctx: &Context, stat: &Stat) { let peer = Peer { id: stat.id, addr: stat.addr.clone(), @@ -217,10 +225,188 @@ async fn rewrite_node_address(ctx: &mut Context, stat: &Stat) { #[cfg(test)] mod tests { + use std::any::Any; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering as AtomicOrdering}; + use std::sync::{Arc, Mutex as StdMutex, mpsc}; + use std::thread; + use std::time::Duration; + use common_meta::datanode::DatanodeStatKey; + use common_meta::error::{Error as MetaError, Result as MetaResult}; + use common_meta::kv_backend::{KvBackend, KvBackendRef, ResettableKvBackend, TxnService}; + use common_meta::rpc::store::{ + BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse, + BatchPutRequest, BatchPutResponse, DeleteRangeRequest, DeleteRangeResponse, PutResponse, + RangeRequest, RangeResponse, + }; + use tokio::sync::Semaphore; + use tokio::time::{sleep, timeout}; use super::*; use crate::handler::test_utils::TestEnv; + use crate::service::store::cached_kv::LeaderCachedKvBackend; + + struct ControlledKvBackend { + recorded_puts: StdMutex>, + put_entered: Semaphore, + put_release: Semaphore, + block_next_put: AtomicBool, + delay_puts: AtomicBool, + active_puts: AtomicUsize, + max_active_puts: AtomicUsize, + } + + impl ControlledKvBackend { + fn new() -> Self { + Self { + recorded_puts: StdMutex::new(Vec::new()), + put_entered: Semaphore::new(0), + put_release: Semaphore::new(0), + block_next_put: AtomicBool::new(false), + delay_puts: AtomicBool::new(false), + active_puts: AtomicUsize::new(0), + max_active_puts: AtomicUsize::new(0), + } + } + + fn block_next_put(&self) { + self.block_next_put.store(true, AtomicOrdering::Relaxed); + } + + async fn wait_for_blocked_put(&self) { + self.put_entered.acquire().await.unwrap().forget(); + } + + fn release_one_put(&self) { + self.put_release.add_permits(1); + } + + fn set_delay_puts(&self, delay: bool) { + self.delay_puts.store(delay, AtomicOrdering::Relaxed); + } + + fn clear_recorded_puts(&self) { + self.recorded_puts.lock().unwrap().clear(); + } + + fn recorded_puts(&self) -> Vec { + self.recorded_puts.lock().unwrap().clone() + } + + fn max_active_puts(&self) -> usize { + self.max_active_puts.load(AtomicOrdering::Relaxed) + } + + fn start_put(&self, req: &PutRequest) -> PutGuard<'_> { + self.recorded_puts.lock().unwrap().push(req.clone()); + let active = self.active_puts.fetch_add(1, AtomicOrdering::Relaxed) + 1; + self.max_active_puts + .fetch_max(active, AtomicOrdering::Relaxed); + PutGuard { backend: self } + } + } + + struct PutGuard<'a> { + backend: &'a ControlledKvBackend, + } + + impl Drop for PutGuard<'_> { + fn drop(&mut self) { + self.backend + .active_puts + .fetch_sub(1, AtomicOrdering::Relaxed); + } + } + + #[async_trait::async_trait] + impl TxnService for ControlledKvBackend { + type Error = MetaError; + } + + #[async_trait::async_trait] + impl KvBackend for ControlledKvBackend { + fn name(&self) -> &str { + "controlled" + } + + fn as_any(&self) -> &dyn Any { + self + } + + async fn range(&self, _req: RangeRequest) -> MetaResult { + unimplemented!() + } + + async fn put(&self, req: PutRequest) -> MetaResult { + let _guard = self.start_put(&req); + if self.block_next_put.swap(false, AtomicOrdering::Relaxed) { + self.put_entered.add_permits(1); + self.put_release.acquire().await.unwrap().forget(); + } + if self.delay_puts.load(AtomicOrdering::Relaxed) { + sleep(Duration::from_millis(10)).await; + } + Ok(PutResponse::default()) + } + + async fn batch_put(&self, _req: BatchPutRequest) -> MetaResult { + unimplemented!() + } + + async fn batch_get(&self, _req: BatchGetRequest) -> MetaResult { + unimplemented!() + } + + async fn delete_range(&self, _req: DeleteRangeRequest) -> MetaResult { + unimplemented!() + } + + async fn batch_delete(&self, _req: BatchDeleteRequest) -> MetaResult { + unimplemented!() + } + } + + impl ResettableKvBackend for ControlledKvBackend { + fn reset(&self) { + self.clear_recorded_puts(); + } + + fn as_kv_backend_ref(self: Arc) -> KvBackendRef { + self + } + } + + fn stat(node_id: u64, epoch: u64, marker: u64, addr: &str) -> Stat { + Stat { + timestamp_millis: marker as i64, + id: node_id, + addr: addr.to_string(), + region_num: marker, + node_epoch: epoch, + ..Default::default() + } + } + + async fn handle_stat( + handler: Arc, + mut ctx: Context, + stat: Stat, + ) -> Result { + let mut acc = HeartbeatAccumulator { + stat: Some(stat), + ..Default::default() + }; + handler + .handle(&HeartbeatRequest::default(), &mut ctx, &mut acc) + .await + } + + fn use_controlled_address_backend(ctx: &mut Context) -> Arc { + let backend = Arc::new(ControlledKvBackend::new()); + ctx.leader_cached_kv_backend = + Arc::new(LeaderCachedKvBackend::with_always_leader(backend.clone())); + backend + } #[tokio::test] async fn test_handle_datanode_stats() { @@ -250,6 +436,105 @@ mod tests { assert_eq!(handler.flush_stats_factor, val.stats.len()); } + #[test] + fn test_same_datanode_wait_keeps_current_thread_runtime_responsive() { + let (backend_tx, backend_rx) = mpsc::sync_channel(1); + let (timer_tx, timer_rx) = mpsc::sync_channel(1); + let (done_tx, done_rx) = mpsc::sync_channel(1); + + let worker = thread::spawn(move || { + tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap() + .block_on(async move { + let env = TestEnv::new(); + let mut ctx = env.ctx(); + let address_backend = use_controlled_address_backend(&mut ctx); + address_backend.block_next_put(); + backend_tx.send(address_backend.clone()).unwrap(); + + let handler = Arc::new(CollectStatsHandler::default()); + let first = tokio::spawn(handle_stat( + handler.clone(), + ctx.clone(), + stat(101, 1, 1, "dn-101-v1"), + )); + address_backend.wait_for_blocked_put().await; + + let second = + tokio::spawn(handle_stat(handler, ctx, stat(101, 1, 2, "dn-101-v1"))); + tokio::spawn(async move { + sleep(Duration::from_millis(10)).await; + timer_tx.send(()).unwrap(); + }); + tokio::task::yield_now().await; + + first.await.unwrap().unwrap(); + second.await.unwrap().unwrap(); + done_tx.send(()).unwrap(); + }); + }); + + let address_backend = backend_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + let timer_result = timer_rx.recv_timeout(Duration::from_secs(1)); + address_backend.release_one_put(); + timer_result.expect("the current-thread runtime must remain responsive"); + done_rx + .recv_timeout(Duration::from_secs(1)) + .expect("both heartbeat handlers must complete after releasing the write"); + worker.join().unwrap(); + } + + #[tokio::test] + async fn test_concurrent_flush_persists_every_stat_once() { + let env = TestEnv::new(); + let mut ctx = env.ctx(); + let stats_backend = Arc::new(ControlledKvBackend::new()); + ctx.in_memory = stats_backend.clone(); + stats_backend.set_delay_puts(true); + + let flush_stats_factor = 3; + let handler = Arc::new(CollectStatsHandler::new(Some(flush_stats_factor))); + handle_stat(handler.clone(), ctx.clone(), stat(101, 1, 0, "dn-101")) + .await + .unwrap(); + stats_backend.clear_recorded_puts(); + + let mut tasks = Vec::with_capacity(2 * flush_stats_factor); + for marker in 1..=(2 * flush_stats_factor) { + tasks.push(tokio::spawn(handle_stat( + handler.clone(), + ctx.clone(), + stat(101, 1, marker as u64, "dn-101"), + ))); + } + for task in tasks { + timeout(Duration::from_secs(1), task) + .await + .unwrap() + .unwrap() + .unwrap(); + } + + let puts = stats_backend.recorded_puts(); + assert_eq!(2, puts.len()); + let mut markers = puts + .into_iter() + .flat_map(|put| { + let value: DatanodeStatValue = put.value.try_into().unwrap(); + value + .stats + .into_iter() + .map(|stat| stat.region_num) + .collect::>() + }) + .collect::>(); + markers.sort_unstable(); + assert_eq!((1..=6).collect::>(), markers); + assert_eq!(1, stats_backend.max_active_puts()); + } + async fn handle_request_many_times( mut ctx: Context, handler: &CollectStatsHandler, From 76f08d2b3f9f5905c9169244dbed8b8bd2178973 Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Fri, 21 Aug 2026 03:28:39 +0000 Subject: [PATCH 66/92] refactor(json2): add bounded auto-expansion to the JSON2 vector builder (#8909) * refactor(json2): add bounded auto-expansion to the JSON2 vector builder Signed-off-by: luofucong * resolve PR comments Signed-off-by: luofucong * fix ci Signed-off-by: luofucong --------- Signed-off-by: luofucong --- src/datatypes/src/vectors/json/builder.rs | 841 ++++++++++++++++++++-- src/datatypes/src/vectors/json/variant.rs | 11 +- 2 files changed, 774 insertions(+), 78 deletions(-) diff --git a/src/datatypes/src/vectors/json/builder.rs b/src/datatypes/src/vectors/json/builder.rs index e6ceb6d6a5f..125dbdad58a 100644 --- a/src/datatypes/src/vectors/json/builder.rs +++ b/src/datatypes/src/vectors/json/builder.rs @@ -13,72 +13,488 @@ // limitations under the License. use std::any::Any; +use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; +use arrow_array::cast::AsArray; +use arrow_array::{Array, ArrayRef, StructArray}; use arrow_schema::DataType; +use parquet_variant_compute::VariantArrayBuilder; +use snafu::{ResultExt, ensure}; use crate::data_type::ConcreteDataType; -use crate::error::{Result, TryFromValueSnafu, UnexpectedSnafu, UnsupportedOperationSnafu}; +use crate::error::{ + ArrowComputeSnafu, Result, TryFromValueSnafu, UnexpectedSnafu, UnsupportedOperationSnafu, +}; +use crate::extension::json::JSON2_REMAINDER_FIELD_NAME; use crate::json::value::{JsonNumber, JsonVariant, encode_json_variant}; +use crate::json::{JSON2_MAX_STRUCTURED_DEPTH, JsonSettings}; use crate::prelude::{ValueRef, Vector, VectorRef}; use crate::types::StructType; use crate::types::json_type::{JsonNativeType, is_include}; use crate::value::{ListValue, StructValue, StructValueRef, Value}; -use crate::vectors::{MutableVector, StructVectorBuilder}; +use crate::vectors::json::variant::{append_json_variant, variant_field}; +use crate::vectors::{Helper, MutableVector, StructVectorBuilder}; +type JsonObjectValue = BTreeMap; + +/// Builds JSON2 vectors from object values. +/// +/// Legacy mode merges all observed paths into the explicit Struct schema. +/// Auto-expanding mode always materializes type-hinted paths, selects up to +/// `max_auto_expanded_paths` compatible unhinted leaf paths by frequency, and stores +/// conflicting or unselected paths in the Variant remainder field. #[derive(Clone)] pub(crate) struct JsonVectorBuilder { - merged_type: JsonNativeType, - values: Vec, + state: JsonVectorBuilderState, +} + +#[derive(Clone)] +enum JsonVectorBuilderState { + Legacy { + merged_type: JsonNativeType, + values: Vec, + }, + AutoExpanding { + /// Paths declared by type hints and always stored as dedicated Struct fields. + explicit_type: JsonNativeType, + /// Maximum number of additional paths selected from buffered values. + max_auto_expanded_paths: u32, + /// Buffered values used to infer auto-expanded paths before building the vector. + values: Vec, + }, +} + +impl JsonVectorBuilderState { + fn native_type(&self) -> JsonNativeType { + match self { + Self::Legacy { merged_type, .. } => merged_type.clone(), + Self::AutoExpanding { + explicit_type, + max_auto_expanded_paths, + values, + } => infer_expanded_type(explicit_type, *max_auto_expanded_paths, values), + } + } + + fn len(&self) -> usize { + match self { + Self::Legacy { values, .. } | Self::AutoExpanding { values, .. } => values.len(), + } + } + + fn try_build(&mut self) -> Result { + match self { + Self::Legacy { + merged_type, + values, + } => build_legacy(values, merged_type), + Self::AutoExpanding { + explicit_type, + max_auto_expanded_paths, + values, + } => { + let expanded_type = + infer_expanded_type(explicit_type, *max_auto_expanded_paths, values); + build_with_remainder(values, &expanded_type) + } + } + } + + fn try_push_value_ref(&mut self, value: &ValueRef) -> Result<()> { + let ValueRef::Json(value) = value else { + return TryFromValueSnafu { + reason: format!("expected JSON value, got {value:?}"), + } + .fail(); + }; + ensure!( + value.is_object() || value.is_null(), + TryFromValueSnafu { + reason: format!("expected JSON object value, got {value:?}"), + } + ); + match self { + Self::Legacy { + merged_type, + values, + } => { + let json_type = value.json_type(); + if !is_include(merged_type, json_type.as_ref()) { + merged_type.merge(json_type.as_ref()); + } + values.push(JsonVariant::from(value.variant())); + } + Self::AutoExpanding { values, .. } => { + values.push(JsonVariant::from(value.variant())); + } + } + Ok(()) + } + + fn push_null(&mut self) { + match self { + Self::Legacy { values, .. } | Self::AutoExpanding { values, .. } => { + values.push(JsonVariant::Null) + } + } + } } impl JsonVectorBuilder { + /// Creates a builder that merges all observed paths into the explicit schema. pub(crate) fn new(initial_native_type: JsonNativeType, capacity: usize) -> Self { debug_assert!(matches!( initial_native_type, JsonNativeType::Object(_) | JsonNativeType::Null )); Self { - merged_type: initial_native_type, - values: Vec::with_capacity(capacity), + state: JsonVectorBuilderState::Legacy { + merged_type: initial_native_type, + values: Vec::with_capacity(capacity), + }, + } + } + + /// Creates a builder bounded by the JSON settings and their type hints. + #[cfg_attr(not(test), expect(dead_code))] + pub(crate) fn with_settings(settings: &JsonSettings, capacity: usize) -> Self { + let mut explicit_type = JsonNativeType::Object(Default::default()); + for hint in settings.type_hints() { + insert_dynamic_type(&mut explicit_type, &hint.path, (&hint.data_type).into()); + } + Self { + state: JsonVectorBuilderState::AutoExpanding { + explicit_type, + max_auto_expanded_paths: settings.max_auto_expanded_paths().unwrap_or(u32::MAX), + values: Vec::with_capacity(capacity), + }, } } fn try_build(&mut self) -> Result { - let DataType::Struct(fields) = self.merged_type.as_arrow_type() else { - return UnexpectedSnafu { - reason: "merged JSON2 type must map to Arrow Struct in JsonVectorBuilder", - } - .fail(); - }; - // TODO(LFC): Direct use Arrow's Struct datatype here. - let struct_type = StructType::from(&fields); - - let mut builder = - StructVectorBuilder::with_type_and_capacity(struct_type.clone(), self.values.len()); - for value in std::mem::take(&mut self.values) { - if matches!(&value, JsonVariant::Null) { - builder.push_null(); - continue; - } - let value = json_variant_into_struct_value(value, struct_type.clone())?; - builder.push_struct_value_ref(StructValueRef::Ref(&value))?; - } - Ok(builder.to_vector()) + self.state.try_build() } } -fn json_variant_into_struct_value( - value: JsonVariant, - struct_type: StructType, -) -> Result { - let JsonVariant::Object(object) = value else { - return TryFromValueSnafu { - reason: format!("expected json object value, got {value:?}"), +fn build_legacy(values: &mut Vec, merged_type: &JsonNativeType) -> Result { + build_explicit(values, merged_type, false, |value| match value { + JsonVariant::Null => Ok(None), + JsonVariant::Object(value) => Ok(Some(value)), + _ => TryFromValueSnafu { + reason: "expected json object value".to_string(), + } + .fail(), + }) +} + +fn build_with_remainder( + values: &mut Vec, + expanded_type: &JsonNativeType, +) -> Result { + let mut remainder = VariantArrayBuilder::new(values.len()); + let explicit = build_explicit(values, expanded_type, true, |value| { + if matches!(value, JsonVariant::Null) { + remainder.append_null(); + return Ok(None); + } + let (value, rest) = split_to_explicit(value, expanded_type)?; + append_json_variant(&mut remainder, &JsonVariant::Object(rest)) + .context(ArrowComputeSnafu)?; + Ok(Some(value)) + })?; + finish_vector(explicit, ArrayRef::from(remainder.build())) +} + +fn build_explicit( + values: &mut Vec, + explicit_type: &JsonNativeType, + // Temporary compatibility switch for the legacy storage layout. Once JSON2 fully switches to + // the v2 storage layout, empty objects should always be preserved instead of treated as null. + preserve_empty_structs: bool, + mut project: impl FnMut(JsonVariant) -> Result>, +) -> Result { + let DataType::Struct(fields) = explicit_type.as_arrow_type() else { + return UnexpectedSnafu { + reason: "merged JSON2 type must map to Arrow Struct in JsonVectorBuilder", } .fail(); }; + // TODO(LFC): Direct use Arrow's Struct datatype here. + let struct_type = StructType::from(&fields); + let mut builder = + StructVectorBuilder::with_type_and_capacity(struct_type.clone(), values.len()); + for value in std::mem::take(values) { + let Some(value) = project(value)? else { + builder.push_null(); + continue; + }; + let value = + json_variant_into_struct_value(value, struct_type.clone(), preserve_empty_structs)?; + builder.push_struct_value_ref(StructValueRef::Ref(&value))?; + } + Ok(builder.to_vector()) +} + +fn finish_vector(explicit: VectorRef, remainder: ArrayRef) -> Result { + let explicit = explicit.to_arrow_array(); + let explicit = explicit.as_struct(); + let mut children = explicit + .fields() + .iter() + .cloned() + .zip(explicit.columns().iter().cloned()) + .chain(std::iter::once(( + Arc::new(variant_field(JSON2_REMAINDER_FIELD_NAME, true)), + remainder, + ))) + .collect::>(); + children.sort_unstable_by(|(x, _), (y, _)| x.name().cmp(y.name())); + let (fields, columns): (Vec<_>, Vec) = children.into_iter().unzip(); + let array: ArrayRef = Arc::new(StructArray::new( + fields.into(), + columns, + explicit.nulls().cloned(), + )); + Helper::try_into_vector(array) +} + +fn infer_expanded_type( + explicit_type: &JsonNativeType, + max_auto_expanded_paths: u32, + values: &[JsonVariant], +) -> JsonNativeType { + if max_auto_expanded_paths == 0 { + return explicit_type.clone(); + } + + let mut stats = HashMap::new(); + let mut path = Vec::new(); + init_explicit_path_stats(explicit_type, &mut path, &mut stats); + for value in values { + count_dynamic_paths(value, &mut path, &mut stats); + } + let mut candidates = stats + .iter() + // Explicit paths are already in the output schema and do not consume the dynamic + // expansion budget. + .filter(|(_, stat)| !stat.is_explicit && stat.is_leaf) + // A leaf is eligible only when both itself and every object prefix have one stable + // role and type across all observed values. + .filter(|(path, _)| { + !(1..=path.len()) + .any(|len| stats.get(&path[..len]).is_some_and(|stats| stats.conflicts)) + }) + .collect::>(); + candidates.sort_unstable_by(|(x_path, x), (y_path, y)| { + y.seen_count + .cmp(&x.seen_count) + .then_with(|| x_path.cmp(y_path)) + }); + + let mut expanded_type = explicit_type.clone(); + for (path, candidate) in candidates + .into_iter() + .take(max_auto_expanded_paths as usize) + { + insert_dynamic_type( + &mut expanded_type, + path, + candidate.expected_leaf_type.clone(), + ); + } + expanded_type +} + +/// Aggregated observations for one JSON path. +/// +/// Schema inference first seeds the map with explicit paths, then walks all input values once. +/// Objects, including empty objects, are non-leaf paths; every other non-null value is a leaf. +/// The first dynamic observation fixes the path role and exact leaf type. A later role or type +/// mismatch sets [`PathStats::conflicts`] permanently. Missing paths and null values do not affect +/// the statistics. +/// +/// After collection, dynamic leaves are ranked by [`PathStats::seen_count`]. A candidate is +/// rejected when it or any parent path conflicts, so candidate selection never rescans the input +/// values. +struct PathStats { + /// Whether the path came from a type hint and is already part of the output schema. + is_explicit: bool, + /// Whether the path is a non-object value rather than an object prefix. + is_leaf: bool, + /// Exact type required for a leaf; unused non-leaf paths keep [`JsonNativeType::Null`]. + expected_leaf_type: JsonNativeType, + /// Number of compatible observations used to rank dynamic leaves. + seen_count: usize, + /// Whether the path has ever had inconsistent roles or leaf types. + conflicts: bool, +} + +/// Seeds path statistics from the configured explicit JSON shape. +fn init_explicit_path_stats<'a>( + explicit_type: &'a JsonNativeType, + path: &mut Vec<&'a str>, + stats: &mut HashMap, PathStats>, +) { + let JsonNativeType::Object(fields) = explicit_type else { + return; + }; + for (name, data_type) in fields { + path.push(name); + let is_leaf = !matches!(data_type, JsonNativeType::Object(_)); + let expected_leaf_type = if is_leaf { + data_type.clone() + } else { + JsonNativeType::default() + }; + stats.insert( + path.clone(), + PathStats { + is_explicit: true, + is_leaf, + expected_leaf_type, + seen_count: 0, + conflicts: false, + }, + ); + init_explicit_path_stats(data_type, path, stats); + path.pop(); + } +} + +/// Collects dynamic path statistics while traversing each input value once. +fn count_dynamic_paths<'a>( + value: &'a JsonVariant, + path: &mut Vec<&'a str>, + stats: &mut HashMap, PathStats>, +) { + if matches!(value, JsonVariant::Null) || path.len() > JSON2_MAX_STRUCTURED_DEPTH { + return; + } + + if !path.is_empty() { + let is_leaf = !matches!(value, JsonVariant::Object(_)); + let conflicts = if let Some(stats) = stats.get_mut(path.as_slice()) { + if !stats.conflicts { + let role_conflict = stats.is_leaf != is_leaf; + let type_conflict = || match (&stats.expected_leaf_type, value) { + // If both objects, they are compatible. + (JsonNativeType::Null | JsonNativeType::Object(_), JsonVariant::Object(_)) => { + false + } + _ => stats.expected_leaf_type != value.native_type(), + }; + if role_conflict || type_conflict() { + stats.conflicts = true; + } else { + stats.seen_count += 1; + } + } + stats.conflicts + } else { + let expected_leaf_type = if is_leaf { + value.native_type() + } else { + JsonNativeType::default() + }; + stats.insert( + path.clone(), + PathStats { + is_explicit: false, + is_leaf, + expected_leaf_type, + seen_count: 1, + conflicts: false, + }, + ); + false + }; + if conflicts { + return; + } + } + + if let JsonVariant::Object(object) = value + && !object.is_empty() + { + for (name, value) in object { + path.push(name); + count_dynamic_paths(value, path, stats); + path.pop(); + } + } +} + +fn insert_dynamic_type>( + explicit_type: &mut JsonNativeType, + path: &[S], + data_type: JsonNativeType, +) { + let JsonNativeType::Object(fields) = explicit_type else { + return; + }; + let Some((name, path)) = path.split_first() else { + return; + }; + let name = name.as_ref().to_string(); + if path.is_empty() { + fields.insert(name, data_type); + return; + } + insert_dynamic_type( + fields + .entry(name) + .or_insert_with(|| JsonNativeType::Object(Default::default())), + path, + data_type, + ) +} + +fn split_to_explicit( + value: JsonVariant, + explicit_type: &JsonNativeType, +) -> Result<(JsonObjectValue, JsonObjectValue)> { + let JsonVariant::Object(mut remainder) = value else { + return TryFromValueSnafu { + reason: "expected json object value".to_string(), + } + .fail(); + }; + let JsonNativeType::Object(fields) = explicit_type else { + return UnexpectedSnafu { + reason: "JSON2 explicit type must be an object", + } + .fail(); + }; + let mut explicit = JsonObjectValue::new(); + + for (name, data_type) in fields { + let Some(value) = remainder.remove(name) else { + continue; + }; + if matches!(data_type, JsonNativeType::Object(_)) { + if matches!(value, JsonVariant::Null) { + explicit.insert(name.clone(), value); + continue; + } + let (child_explicit, child_remainder) = split_to_explicit(value, data_type)?; + explicit.insert(name.clone(), JsonVariant::Object(child_explicit)); + if !child_remainder.is_empty() { + remainder.insert(name.clone(), JsonVariant::Object(child_remainder)); + } + } else { + explicit.insert(name.clone(), value); + } + } + Ok((explicit, remainder)) +} + +fn json_variant_into_struct_value( + object: JsonObjectValue, + struct_type: StructType, + preserve_empty_structs: bool, +) -> Result { let mut entries = object.into_iter(); let mut entry = entries.next(); let mut values = Vec::with_capacity(struct_type.fields().len()); @@ -86,7 +502,7 @@ fn json_variant_into_struct_value( let value = match entry.take() { Some((name, value)) if name == field.name() => { entry = entries.next(); - json_variant_into_value(value, field.data_type())? + json_variant_into_value(value, field.data_type(), preserve_empty_structs)? } Some((name, _)) if name.as_str() < field.name() => { return TryFromValueSnafu { @@ -111,10 +527,19 @@ fn json_variant_into_struct_value( Ok(StructValue::new(values, struct_type)) } -fn json_variant_into_value(value: JsonVariant, expected_type: &ConcreteDataType) -> Result { +fn json_variant_into_value( + value: JsonVariant, + expected_type: &ConcreteDataType, + preserve_empty_structs: bool, +) -> Result { let value = match (value, expected_type) { (JsonVariant::Null, _) | (_, ConcreteDataType::Null(_)) => Value::Null, - (JsonVariant::Object(object), _) if object.is_empty() => Value::Null, + (JsonVariant::Object(object), _) if object.is_empty() && !preserve_empty_structs => { + Value::Null + } + (JsonVariant::Object(object), ConcreteDataType::Struct(struct_type)) => Value::Struct( + json_variant_into_struct_value(object, struct_type.clone(), preserve_empty_structs)?, + ), (JsonVariant::Bool(x), ConcreteDataType::Boolean(_)) => Value::Boolean(x), (JsonVariant::Number(x), ConcreteDataType::UInt64(_)) => { let Some(x) = x.as_u64() else { @@ -153,13 +578,10 @@ fn json_variant_into_value(value: JsonVariant, expected_type: &ConcreteDataType) let item_type = list_type.item_type().clone(); let values = array .into_iter() - .map(|v| json_variant_into_value(v, &item_type)) + .map(|v| json_variant_into_value(v, &item_type, preserve_empty_structs)) .collect::>>()?; Value::List(ListValue::new(values, Arc::new(item_type))) } - (value @ JsonVariant::Object(_), ConcreteDataType::Struct(struct_type)) => { - Value::Struct(json_variant_into_struct_value(value, struct_type.clone())?) - } (value, ConcreteDataType::Binary(_)) => Value::from(encode_json_variant(value)?), (value, expected_type) => { return TryFromValueSnafu { @@ -173,11 +595,11 @@ fn json_variant_into_value(value: JsonVariant, expected_type: &ConcreteDataType) impl MutableVector for JsonVectorBuilder { fn data_type(&self) -> ConcreteDataType { - ConcreteDataType::json2(self.merged_type.clone()) + ConcreteDataType::json2(self.state.native_type()) } fn len(&self) -> usize { - self.values.len() + self.state.len() } fn as_any(&self) -> &dyn Any { @@ -197,30 +619,11 @@ impl MutableVector for JsonVectorBuilder { } fn try_push_value_ref(&mut self, value: &ValueRef) -> Result<()> { - let ValueRef::Json(value) = value else { - return TryFromValueSnafu { - reason: format!("expected json value, got {value:?}"), - } - .fail(); - }; - let json_type = value.json_type(); - let json_type = json_type.as_ref(); - if !matches!(json_type, JsonNativeType::Object(_) | JsonNativeType::Null) { - return TryFromValueSnafu { - reason: format!("expected json object value, got {value:?}"), - } - .fail(); - } - if !is_include(&self.merged_type, json_type) { - self.merged_type.merge(json_type); - } - - self.values.push(JsonVariant::from(value.variant())); - Ok(()) + self.state.try_push_value_ref(value) } fn push_null(&mut self) { - self.values.push(JsonVariant::Null) + self.state.push_null() } fn extend_slice_of(&mut self, _: &dyn Vector, _: usize, _: usize) -> Result<()> { @@ -236,13 +639,21 @@ impl MutableVector for JsonVectorBuilder { mod tests { use std::sync::Arc; + use arrow_array::cast::AsArray; + use arrow_schema::Field; use common_base::bytes::Bytes; + use serde_json::json; use super::*; use crate::data_type::ConcreteDataType; + use crate::extension::json::{Json2ExtensionType, JsonMetadata}; + use crate::json::JsonTypeHint; + use crate::json::value::decode_json_variant; use crate::types::StructField; use crate::types::json_type::JsonObjectType; use crate::value::{ListValue, StructValue, Value, ValueRef}; + use crate::vectors::json::array::JsonArray; + use crate::vectors::json::variant::variant_to_json_values; #[test] fn test_json_vector_builder() -> Result<()> { @@ -346,7 +757,7 @@ mod tests { let err = object_builder .try_push_value_ref(&boolean.as_value_ref()) .unwrap_err(); - assert!(err.to_string().contains("expected json object value")); + assert!(err.to_string().contains("expected JSON object value")); object_builder.try_push_value_ref(&object.as_value_ref())?; // Non-JSON values should be rejected at push time. @@ -355,20 +766,311 @@ mod tests { let err = invalid_builder .try_push_value_ref(&ValueRef::Boolean(true)) .unwrap_err(); - assert!(err.to_string().contains("expected json value")); + assert!(err.to_string().contains("expected JSON value")); + + Ok(()) + } + + #[test] + fn test_zero_budget_builder_uses_fixed_schema_and_remainder() -> Result<()> { + let settings = JsonSettings::try_new( + vec![ + JsonTypeHint { + path: vec!["kind".to_string()], + data_type: ConcreteDataType::string_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }, + JsonTypeHint { + path: vec!["commit".to_string(), "operation".to_string()], + data_type: ConcreteDataType::string_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }, + JsonTypeHint { + path: vec!["time_us".to_string()], + data_type: ConcreteDataType::int64_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }, + ], + Some(0), + )?; + let mut builder = JsonVectorBuilder::with_settings(&settings, 2); + let values = [ + json!({ + "kind": "record", + "commit": {"operation": "create", "collection": "post"}, + "extra": 1, + "time_us": 1 + }), + json!({"kind": "other", "dynamic": true, "time_us": 2}), + ]; + for value in values.clone() { + let value = settings.encode(value)?; + builder.try_push_value_ref(&value.as_value_ref())?; + } + let array = builder.to_vector().to_arrow_array(); + assert_eq!(0, builder.len()); + let structs = array.as_struct(); + assert_eq!( + vec![JSON2_REMAINDER_FIELD_NAME, "commit", "kind", "time_us"], + structs + .fields() + .iter() + .map(|x| x.name().as_str()) + .collect::>() + ); + assert_eq!( + vec![ + Some(json!({"commit": {"collection": "post"}, "extra": 1})), + Some(json!({"dynamic": true})), + ], + variant_to_json_values(structs.column_by_name(JSON2_REMAINDER_FIELD_NAME).unwrap())? + ); + + let field = Field::new("data", array.data_type().clone(), true).with_extension_type( + Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(settings))), + ); + let reconstructed = JsonArray::from(&array).project_to_v2(&field, &DataType::Binary)?; + let reconstructed = reconstructed.as_binary::(); + assert_eq!( + values[0], + decode_json_variant(reconstructed.value(0)).unwrap() + ); + assert_eq!( + json!({ + "kind": "other", + "commit": {"operation": null}, + "dynamic": true, + "time_us": 2 + }), + decode_json_variant(reconstructed.value(1)).unwrap() + ); + + let settings = JsonSettings::try_new(vec![], Some(0))?; + let mut builder = JsonVectorBuilder::with_settings(&settings, 3); + for value in [json!({}), json!({"x": 1})] { + let value = settings.encode(value)?; + builder.try_push_value_ref(&value.as_value_ref())?; + } + builder.push_null(); + let array = builder.to_vector().to_arrow_array(); + let structs = array.as_struct(); + assert_eq!(1, structs.num_columns()); + assert_eq!( + vec![Some(json!({})), Some(json!({"x": 1})), None], + variant_to_json_values(structs.column(0))? + ); + + Ok(()) + } + + #[test] + fn test_finite_budget_selects_dynamic_paths() -> Result<()> { + let settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["hint".to_string()], + data_type: ConcreteDataType::string_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }], + Some(2), + )?; + let values = [ + json!({ + "hint": "first", + "conflict": 1, + "popular": {"nested": 1}, + "tie_a": "a", + "tie_b": true, + "rare": 1 + }), + json!({ + "hint": "second", + "conflict": "string", + "popular": {"nested": 2}, + "tie_a": "b", + "tie_b": false + }), + json!({"hint": "third", "popular": "scalar"}), + ]; + let mut builder = JsonVectorBuilder::with_settings(&settings, values.len()); + for value in values.clone() { + let value = settings.encode(value)?; + builder.try_push_value_ref(&value.as_value_ref())?; + } + + let array = builder.to_vector().to_arrow_array(); + let structs = array.as_struct(); + assert_eq!( + vec![JSON2_REMAINDER_FIELD_NAME, "hint", "tie_a", "tie_b"], + structs + .fields() + .iter() + .map(|x| x.name().as_str()) + .collect::>() + ); + assert!(structs.column_by_name("popular").is_none()); + assert_eq!( + vec![ + Some(json!({"conflict": 1, "popular": {"nested": 1}, "rare": 1})), + Some(json!({"conflict": "string", "popular": {"nested": 2}})), + Some(json!({"popular": "scalar"})), + ], + variant_to_json_values(structs.column_by_name(JSON2_REMAINDER_FIELD_NAME).unwrap())? + ); + + let field = Field::new("data", array.data_type().clone(), true).with_extension_type( + Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(settings))), + ); + let reconstructed = JsonArray::from(&array).project_to_v2(&field, &DataType::Binary)?; + let reconstructed = reconstructed.as_binary::(); + assert_eq!( + values[0], + decode_json_variant(reconstructed.value(0)).unwrap() + ); + assert_eq!( + values[1], + decode_json_variant(reconstructed.value(1)).unwrap() + ); + assert_eq!( + json!({ + "hint": "third", + "popular": "scalar", + "tie_a": null, + "tie_b": null + }), + decode_json_variant(reconstructed.value(2)).unwrap() + ); + + Ok(()) + } + + #[test] + fn test_reconstruct_nested_remainder_only_value() -> Result<()> { + let settings = JsonSettings::try_new(vec![], Some(1))?; + let values = [ + json!({"a": {"hot": 1}}), + json!({"a": {"hot": 2}}), + json!({"a": {"cold": 3}}), + ]; + let mut builder = JsonVectorBuilder::with_settings(&settings, values.len()); + for value in values.clone() { + let value = settings.encode(value)?; + builder.try_push_value_ref(&value.as_value_ref())?; + } + + let array = builder.to_vector().to_arrow_array(); + let field = Field::new("data", array.data_type().clone(), true).with_extension_type( + Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(settings))), + ); + let reconstructed = JsonArray::from(&array).project_to_v2(&field, &DataType::Binary)?; + let reconstructed = reconstructed.as_binary::(); + assert_eq!( + vec![ + values[0].clone(), + values[1].clone(), + json!({"a": {"cold": 3, "hot": null}}), + ], + (0..reconstructed.len()) + .map(|i| decode_json_variant(reconstructed.value(i)).unwrap()) + .collect::>() + ); + + Ok(()) + } + + #[test] + fn test_dynamic_paths_require_the_same_leaf_type() -> Result<()> { + let settings = JsonSettings::try_new( + vec![JsonTypeHint { + path: vec!["nested".to_string(), "hinted".to_string()], + data_type: ConcreteDataType::string_datatype(), + nullable: true, + default_constraint: None, + inverted_index: false, + }], + Some(8), + )?; + let values = [ + json!({ + "branch": {}, + "different": [1], + "empty": {}, + "nested": {}, + "reverse": "scalar", + "same": [1] + }), + json!({ + "branch": {"leaf": 1}, + "different": ["x"], + "empty": {}, + "nested": {"hinted": "x", "leaf": 1}, + "same": [2] + }), + json!({ + "branch": {}, + "different": [2], + "empty": {}, + "nested": {}, + "reverse": {"leaf": 1}, + "same": [3] + }), + ]; + let mut builder = JsonVectorBuilder::with_settings(&settings, values.len()); + for value in values { + let value = settings.encode(value)?; + builder.try_push_value_ref(&value.as_value_ref())?; + } + + let JsonNativeType::Object(fields) = builder.state.native_type() else { + unreachable!(); + }; + assert!(!fields.contains_key("different")); + assert!(!fields.contains_key("empty")); + assert!(!fields.contains_key("reverse")); + assert!(matches!(fields.get("same"), Some(JsonNativeType::Array(_)))); + assert!(matches!( + fields.get("branch"), + Some(JsonNativeType::Object(fields)) if fields.contains_key("leaf") + )); + assert!(matches!( + fields.get("nested"), + Some(JsonNativeType::Object(fields)) + if fields.contains_key("hinted") && fields.contains_key("leaf") + )); Ok(()) } #[test] fn test_json_variant_into_struct_value() -> Result<()> { + let struct_type = StructType::new(Arc::new(vec![StructField::new( + "value".to_string(), + ConcreteDataType::string_datatype(), + true, + )])); assert_eq!( json_variant_into_value( JsonVariant::Object(Default::default()), - &ConcreteDataType::string_datatype(), + &ConcreteDataType::struct_datatype(struct_type.clone()), + false, )?, Value::Null ); + assert_eq!( + json_variant_into_value( + JsonVariant::Object(Default::default()), + &ConcreteDataType::struct_datatype(struct_type.clone()), + true, + )?, + Value::Struct(StructValue::new(vec![Value::Null], struct_type)) + ); let item_type = ConcreteDataType::struct_datatype(StructType::new(Arc::new(vec![StructField::new( @@ -394,22 +1096,23 @@ mod tests { true, ), ])); - let variant = JsonVariant::from([ + let variant = JsonObjectValue::from([ ( - "items", + "items".to_string(), JsonVariant::Array(vec![ JsonVariant::from([("id", JsonVariant::from(1i64))]), JsonVariant::from([("id", JsonVariant::from(2i64))]), ]), ), ( - "meta", + "meta".to_string(), JsonVariant::from([("name", JsonVariant::from("foo"))]), ), ]); let value = Value::Struct(json_variant_into_struct_value( variant, struct_type.clone(), + true, )?); assert_eq!( diff --git a/src/datatypes/src/vectors/json/variant.rs b/src/datatypes/src/vectors/json/variant.rs index 1a4b8a861d5..75d87c99861 100644 --- a/src/datatypes/src/vectors/json/variant.rs +++ b/src/datatypes/src/vectors/json/variant.rs @@ -15,10 +15,7 @@ use std::sync::Arc; use arrow_array::ArrayRef; -#[cfg(test)] -use arrow_schema::ArrowError; -use arrow_schema::{DataType, Field}; -#[cfg(test)] +use arrow_schema::{ArrowError, DataType, Field}; use parquet_variant::{ObjectFieldBuilder, Variant, VariantBuilderExt, VariantDecimal16}; #[cfg(test)] use parquet_variant_compute::VariantArrayBuilder; @@ -27,7 +24,6 @@ use parquet_variant_json::VariantToJson; use snafu::ResultExt; use crate::error::{ArrowComputeSnafu, Result}; -#[cfg(test)] use crate::json::value::{JsonNumber, JsonVariant, decode_json_variant}; /// Returns the canonical Arrow field for an unshredded Parquet Variant array. @@ -68,8 +64,7 @@ fn json_variants_to_variant(values: &[Option]) -> Result Ok(ArrayRef::from(builder.build())) } -#[cfg(test)] -fn append_json_variant( +pub(super) fn append_json_variant( builder: &mut impl VariantBuilderExt, value: &JsonVariant, ) -> std::result::Result<(), ArrowError> { @@ -115,7 +110,6 @@ fn append_json_variant( Ok(()) } -#[cfg(test)] fn append_json_value( builder: &mut impl VariantBuilderExt, value: &serde_json::Value, @@ -157,7 +151,6 @@ fn append_json_value( /// Parquet Variant has no unsigned integer primitive. Treat u64 as i64 first, then use Decimal16 /// to represent large (larger than i64::MAX) u64. -#[cfg(test)] fn append_large_u64( builder: &mut impl VariantBuilderExt, value: u64, From 8d887ddd00d67699d0fb873cda11eb50b5771555 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 21 Aug 2026 03:50:30 +0000 Subject: [PATCH 67/92] fix(query): restore columnar group-by for dictionary-encoded tags (#8902) * fix(query): restore columnar group-by for dictionary-encoded tags Bump the DataFusion fork to be93ffd85 (feat/dict-group-column-53), which backports apache/datafusion #23187: DictionaryGroupValuesColumn lets dictionary-encoded group keys use the columnar GroupValuesColumn fast path (hash distinct dictionary values once per batch, resolve rows by key index) instead of falling back to row-based GroupValuesRows. This fixes the TSBS double-groupby regression introduced by #8541 (preserve dictionary-encoded query labels): v1.2.0-beta.1 scan output changed tag columns to Dictionary(UInt32, Utf8), which DataFusion 53.1.0 did not support in GroupValuesColumn's supported_type allow-list, so GROUP BY queries silently dropped to the ~60% slower row path (time_calculating_group_ids +57%, peak_mem +50%, end-to-end +38%). Adds an end-to-end integration test (dict_groupby_sst) that flushes a flat-format SST with dictionary-encoded hostname, runs the tsbs-style double-groupby query, and asserts correct results with no CastExec inserted before the aggregate. Signed-off-by: discord9 * chore(deps): pin merged dictionary group-by support Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> Co-authored-by: discord9 --- Cargo.lock | 68 +++++----- Cargo.toml | 30 ++--- tests-integration/tests/dict_groupby_sst.rs | 139 ++++++++++++++++++++ tests-integration/tests/main.rs | 1 + 4 files changed, 189 insertions(+), 49 deletions(-) create mode 100644 tests-integration/tests/dict_groupby_sst.rs diff --git a/Cargo.lock b/Cargo.lock index d7c9241afca..a389391ea90 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3738,7 +3738,7 @@ checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" [[package]] name = "datafusion" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "arrow-schema 58.3.0", @@ -3792,7 +3792,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3816,7 +3816,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3838,7 +3838,7 @@ dependencies = [ [[package]] name = "datafusion-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -3862,7 +3862,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "futures", "log", @@ -3872,7 +3872,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "async-compression", @@ -3906,7 +3906,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "arrow-ipc 58.3.0", @@ -3929,7 +3929,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3951,7 +3951,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3974,7 +3974,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "async-trait", @@ -4003,12 +4003,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" [[package]] name = "datafusion-execution" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "arrow-buffer 58.3.0", @@ -4030,7 +4030,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "async-trait", @@ -4052,7 +4052,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4064,7 +4064,7 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "arrow-buffer 58.3.0", @@ -4095,7 +4095,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4116,7 +4116,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4128,7 +4128,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "arrow-ord 58.3.0", @@ -4152,7 +4152,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "async-trait", @@ -4167,7 +4167,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4184,7 +4184,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -4193,7 +4193,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "datafusion-doc", "quote", @@ -4203,7 +4203,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "chrono", @@ -4252,7 +4252,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4275,7 +4275,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4289,7 +4289,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4305,7 +4305,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4323,7 +4323,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4354,7 +4354,7 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "chrono", @@ -4381,7 +4381,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4391,7 +4391,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4407,7 +4407,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "async-trait", "datafusion-common", @@ -4420,7 +4420,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "arrow 58.3.0", "bigdecimal 0.4.8", @@ -4438,7 +4438,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=d705d4433d2fcbe2bead54e3f1d6cc59727857e1#d705d4433d2fcbe2bead54e3f1d6cc59727857e1" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=4c8a6bf28347b346348136247eb67dfb933f1147#4c8a6bf28347b346348136247eb67dfb933f1147" dependencies = [ "async-recursion", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 777d8b302b4..e6b849a6cd0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -351,21 +351,21 @@ git = "https://github.com/GreptimeTeam/greptime-meter.git" rev = "5618e779cf2bb4755b499c630fba4c35e91898cb" [patch.crates-io] -datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } -datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "d705d4433d2fcbe2bead54e3f1d6cc59727857e1" } +datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } +datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4c8a6bf28347b346348136247eb67dfb933f1147" } sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "2aefa08a8d69c96eec2d6d6703598a009bba6e4c" } # on branch v0.61.x [profile.release] diff --git a/tests-integration/tests/dict_groupby_sst.rs b/tests-integration/tests/dict_groupby_sst.rs new file mode 100644 index 00000000000..862b77ef108 --- /dev/null +++ b/tests-integration/tests/dict_groupby_sst.rs @@ -0,0 +1,139 @@ +// Copyright 2023 Greptime Team +// +// 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. + +use common_query::Output; +use frontend::instance::Instance; +use servers::query_handler::sql::SqlQueryHandler; +use session::context::QueryContext; +use tests_integration::standalone::GreptimeDbStandaloneBuilder; + +async fn execute_sql(instance: &Instance, sql: &str) -> Output { + SqlQueryHandler::do_query(instance, sql, QueryContext::arc()) + .await + .remove(0) + .unwrap() +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_dict_groupby_sst() { + common_telemetry::init_default_ut_logging(); + + let standalone = GreptimeDbStandaloneBuilder::new("dict_groupby_sst") + .build() + .await; + let frontend = standalone.fe_instance(); + + execute_sql( + frontend, + r#" + CREATE TABLE cpu ( + hostname STRING PRIMARY KEY, + usage_user DOUBLE, + usage_system DOUBLE, + usage_idle DOUBLE, + greptime_timestamp TIMESTAMP TIME INDEX + ) WITH (append_mode = 'true', sst_format = 'flat') + "#, + ) + .await; + + let mut rows = String::new(); + for index in 0..1200 { + if index != 0 { + rows.push(','); + } + let hostname = match index % 3 { + 0 => "host_a", + 1 => "host_b", + _ => "host_c", + }; + let usage_user = match index % 3 { + 0 => 10.0, + 1 => 20.0, + _ => 30.0, + }; + let hour = 11 + index / 600; + let minute = (index % 600) / 60; + let second = index % 60; + rows.push_str(&format!( + "('{hostname}', {usage_user}, 1.0, 0.0, '2023-06-12T{hour:02}:{minute:02}:{second:02}Z')" + )); + } + execute_sql(frontend, &format!("INSERT INTO cpu VALUES {rows}")).await; + + execute_sql(frontend, "ADMIN FLUSH_TABLE('cpu')").await; + + let count = execute_sql(frontend, "SELECT count(*) FROM cpu") + .await + .data + .pretty_print() + .await; + assert!( + count.contains("| 1200 |"), + "unexpected row count:\n{count}" + ); + + let query = r#" + SELECT hostname, avg(usage_user), date_trunc('hour', greptime_timestamp) + FROM cpu + WHERE greptime_timestamp >= '2023-06-12T00:00:00Z' + AND greptime_timestamp < '2023-06-14T00:00:00Z' + GROUP BY date_trunc('hour', greptime_timestamp), hostname + ORDER BY hostname, date_trunc('hour', greptime_timestamp) + "#; + let groups = execute_sql(frontend, query).await.data.pretty_print().await; + assert_eq!( + 6, + groups + .lines() + .filter(|line| line.starts_with("| host_")) + .count(), + "expected one group per hostname and hour:\n{groups}" + ); + for (hostname, average) in [("host_a", "10.0"), ("host_b", "20.0"), ("host_c", "30.0")] { + assert_eq!( + 2, + groups.matches(&format!("| {hostname} ")).count(), + "expected two hourly groups for {hostname}:\n{groups}" + ); + assert!( + groups.contains(average), + "expected average {average} for {hostname}:\n{groups}" + ); + } + + let explain = execute_sql( + frontend, + r#" + EXPLAIN ANALYZE VERBOSE + SELECT avg(usage_user), date_trunc('hour', greptime_timestamp) + FROM cpu + WHERE greptime_timestamp >= '2023-06-12T00:00:00Z' + AND greptime_timestamp < '2023-06-14T00:00:00Z' + GROUP BY date_trunc('hour', greptime_timestamp), hostname + "#, + ) + .await + .data + .pretty_print() + .await; + assert!( + explain.contains("AggregateExec"), + "expected aggregate in physical plan:\n{explain}" + ); + assert!( + !explain.contains("CastExec") && !explain.contains("CAST(hostname"), + "dictionary hostname must not be cast before aggregation:\n{explain}" + ); +} diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index 3112cb72c32..965a3a67462 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -16,6 +16,7 @@ mod admin_function_event; mod database_ddl_event; +mod dict_groupby_sst; mod event_recorder_test_util; #[macro_use] mod grpc; From 7fd0a7bb9877b218943429296ad442efcf309f7c Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Fri, 21 Aug 2026 12:14:47 +0000 Subject: [PATCH 68/92] refactor(json2): optimize JSON2 building without auto-expanded paths (#8928) * refactor(json2): optimize JSON2 building without auto-expanded paths Signed-off-by: luofucong * resolve PR comments Signed-off-by: luofucong * avoid panicking memtable write Signed-off-by: luofucong --------- Signed-off-by: luofucong --- src/datatypes/src/error.rs | 8 + src/datatypes/src/json.rs | 25 +++ src/datatypes/src/vectors/json/builder.rs | 254 ++++++++++++++++++++-- src/datatypes/src/vectors/json/variant.rs | 48 +++- 4 files changed, 321 insertions(+), 14 deletions(-) diff --git a/src/datatypes/src/error.rs b/src/datatypes/src/error.rs index b34a5ffa9bf..1def0d6be8d 100644 --- a/src/datatypes/src/error.rs +++ b/src/datatypes/src/error.rs @@ -77,6 +77,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Unimplemented: {feat}"))] + Unimplemented { + feat: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to parse version in schema meta, value: {}", value))] ParseSchemaVersion { value: String, @@ -312,6 +319,7 @@ impl ErrorExt for Error { use Error::*; match self { UnsupportedOperation { .. } + | Unimplemented { .. } | UnsupportedArrowType { .. } | UnsupportedJsonType { .. } | UnsupportedDefaultExpr { .. } => StatusCode::Unsupported, diff --git a/src/datatypes/src/json.rs b/src/datatypes/src/json.rs index ec9ca4ac789..79fc10b12e7 100644 --- a/src/datatypes/src/json.rs +++ b/src/datatypes/src/json.rs @@ -140,6 +140,16 @@ impl JsonSettings { /// Encode a serde_json::Value into a Value::Json using current settings. pub fn encode(&self, json: Json) -> Result { + if let Json::Object(object) = &json + && object.contains_key(JSON2_REMAINDER_FIELD_NAME) + { + return error::InvalidJsonSnafu { + value: format!( + "root object cannot contain reserved field '{JSON2_REMAINDER_FIELD_NAME}'" + ), + } + .fail(); + } let mut context = JsonContext { path: Vec::new(), settings: self, @@ -777,6 +787,21 @@ mod tests { } } + #[test] + fn test_encode_rejects_reserved_remainder_field() -> Result<()> { + let settings = JsonSettings::default(); + let err = settings + .encode(json!({"!__remainder__!": "user-value"})) + .unwrap_err(); + assert!( + err.to_string() + .contains("root object cannot contain reserved field") + ); + + settings.encode(json!({"nested": {"!__remainder__!": "user-value"}}))?; + Ok(()) + } + #[test] fn test_encode_json_object() { let json = json!({ diff --git a/src/datatypes/src/vectors/json/builder.rs b/src/datatypes/src/vectors/json/builder.rs index 125dbdad58a..ab3fd6d0436 100644 --- a/src/datatypes/src/vectors/json/builder.rs +++ b/src/datatypes/src/vectors/json/builder.rs @@ -24,17 +24,18 @@ use snafu::{ResultExt, ensure}; use crate::data_type::ConcreteDataType; use crate::error::{ - ArrowComputeSnafu, Result, TryFromValueSnafu, UnexpectedSnafu, UnsupportedOperationSnafu, + ArrowComputeSnafu, Result, TryFromValueSnafu, UnexpectedSnafu, UnimplementedSnafu, + UnsupportedOperationSnafu, }; use crate::extension::json::JSON2_REMAINDER_FIELD_NAME; -use crate::json::value::{JsonNumber, JsonVariant, encode_json_variant}; +use crate::json::value::{JsonNumber, JsonVariant, JsonVariantRef, encode_json_variant}; use crate::json::{JSON2_MAX_STRUCTURED_DEPTH, JsonSettings}; use crate::prelude::{ValueRef, Vector, VectorRef}; use crate::types::StructType; use crate::types::json_type::{JsonNativeType, is_include}; use crate::value::{ListValue, StructValue, StructValueRef, Value}; -use crate::vectors::json::variant::{append_json_variant, variant_field}; -use crate::vectors::{Helper, MutableVector, StructVectorBuilder}; +use crate::vectors::json::variant::{append_json_variant, append_json_variant_ref, variant_field}; +use crate::vectors::{Helper, MutableVector, NullVector, StructVectorBuilder}; type JsonObjectValue = BTreeMap; @@ -44,17 +45,25 @@ type JsonObjectValue = BTreeMap; /// Auto-expanding mode always materializes type-hinted paths, selects up to /// `max_auto_expanded_paths` compatible unhinted leaf paths by frequency, and stores /// conflicting or unselected paths in the Variant remainder field. -#[derive(Clone)] pub(crate) struct JsonVectorBuilder { state: JsonVectorBuilderState, } -#[derive(Clone)] enum JsonVectorBuilderState { Legacy { merged_type: JsonNativeType, values: Vec, }, + ExplicitOnly { + /// Paths declared by type hints and stored as dedicated Struct fields. + explicit_type: JsonNativeType, + /// Concrete Struct type used to append explicit values without buffering rows. + struct_type: StructType, + /// Builder for values selected by the explicit type hints. + explicit: StructVectorBuilder, + /// Builder for all values outside the explicit type hints. + remainder: VariantArrayBuilder, + }, AutoExpanding { /// Paths declared by type hints and always stored as dedicated Struct fields. explicit_type: JsonNativeType, @@ -69,6 +78,7 @@ impl JsonVectorBuilderState { fn native_type(&self) -> JsonNativeType { match self { Self::Legacy { merged_type, .. } => merged_type.clone(), + Self::ExplicitOnly { explicit_type, .. } => explicit_type.clone(), Self::AutoExpanding { explicit_type, max_auto_expanded_paths, @@ -80,6 +90,7 @@ impl JsonVectorBuilderState { fn len(&self) -> usize { match self { Self::Legacy { values, .. } | Self::AutoExpanding { values, .. } => values.len(), + Self::ExplicitOnly { explicit, .. } => explicit.len(), } } @@ -89,6 +100,14 @@ impl JsonVectorBuilderState { merged_type, values, } => build_legacy(values, merged_type), + Self::ExplicitOnly { + explicit, + remainder, + .. + } => { + let remainder = std::mem::replace(remainder, VariantArrayBuilder::new(0)).build(); + finish_vector(explicit.to_vector(), ArrayRef::from(remainder)) + } Self::AutoExpanding { explicit_type, max_auto_expanded_paths, @@ -101,6 +120,38 @@ impl JsonVectorBuilderState { } } + fn try_build_cloned(&self) -> Result { + let mut state = match self { + Self::Legacy { + merged_type, + values, + } => Self::Legacy { + merged_type: merged_type.clone(), + values: values.clone(), + }, + Self::AutoExpanding { + explicit_type, + max_auto_expanded_paths, + values, + } => Self::AutoExpanding { + explicit_type: explicit_type.clone(), + max_auto_expanded_paths: *max_auto_expanded_paths, + values: values.clone(), + }, + // Only TimeSeriesMemtable requires a non-consuming snapshot, while JSON2 targets + // BulkMemtable. We've tried our best to support it above, but if this match arm does + // not, it's OK. The only reason it doesn't is because of `VariantArrayBuilder`. We'll + // track the upstream and see. + Self::ExplicitOnly { .. } => { + return UnimplementedSnafu { + feat: "no auto expanded JSON2 array builder", + } + .fail(); + } + }; + state.try_build() + } + fn try_push_value_ref(&mut self, value: &ValueRef) -> Result<()> { let ValueRef::Json(value) = value else { return TryFromValueSnafu { @@ -125,6 +176,22 @@ impl JsonVectorBuilderState { } values.push(JsonVariant::from(value.variant())); } + Self::ExplicitOnly { + explicit_type, + struct_type, + explicit, + remainder, + } => { + if value.is_null() { + explicit.push_null(); + remainder.append_null(); + } else { + let (value, rest) = + split_to_explicit_ref(value.variant(), explicit_type, struct_type)?; + explicit.push_struct_value_ref(value)?; + append_json_variant_ref(remainder, &rest).context(ArrowComputeSnafu)?; + } + } Self::AutoExpanding { values, .. } => { values.push(JsonVariant::from(value.variant())); } @@ -137,6 +204,14 @@ impl JsonVectorBuilderState { Self::Legacy { values, .. } | Self::AutoExpanding { values, .. } => { values.push(JsonVariant::Null) } + Self::ExplicitOnly { + explicit, + remainder, + .. + } => { + explicit.push_null(); + remainder.append_null(); + } } } } @@ -163,13 +238,28 @@ impl JsonVectorBuilder { for hint in settings.type_hints() { insert_dynamic_type(&mut explicit_type, &hint.path, (&hint.data_type).into()); } - Self { - state: JsonVectorBuilderState::AutoExpanding { + let state = if settings.max_auto_expanded_paths() == Some(0) { + let DataType::Struct(fields) = explicit_type.as_arrow_type() else { + unreachable!("JSON2 explicit type must map to Arrow Struct") + }; + let struct_type = StructType::from(&fields); + JsonVectorBuilderState::ExplicitOnly { + explicit_type, + explicit: StructVectorBuilder::with_type_and_capacity( + struct_type.clone(), + capacity, + ), + struct_type, + remainder: VariantArrayBuilder::new(capacity), + } + } else { + JsonVectorBuilderState::AutoExpanding { explicit_type, max_auto_expanded_paths: settings.max_auto_expanded_paths().unwrap_or(u32::MAX), values: Vec::with_capacity(capacity), - }, - } + } + }; + Self { state } } fn try_build(&mut self) -> Result { @@ -451,6 +541,132 @@ fn insert_dynamic_type>( ) } +fn split_to_explicit_ref<'a>( + value: &JsonVariantRef<'a>, + explicit_type: &JsonNativeType, + struct_type: &StructType, +) -> Result<(StructValueRef<'a>, JsonVariantRef<'a>)> { + let JsonVariantRef::Object(object) = value else { + return TryFromValueSnafu { + reason: "expected json object value".to_string(), + } + .fail(); + }; + let explicit = json_object_ref_into_struct_value_ref(object, struct_type)?; + let remainder = remainder_ref(object, explicit_type)?; + Ok((explicit, JsonVariantRef::Object(remainder))) +} + +fn json_object_ref_into_struct_value_ref<'a>( + object: &BTreeMap<&'a str, JsonVariantRef<'a>>, + struct_type: &StructType, +) -> Result> { + let mut values = Vec::with_capacity(struct_type.fields().len()); + for field in struct_type.fields().iter() { + let value = match object.get(field.name()) { + Some(value) => json_variant_ref_into_value_ref(value, field.data_type())?, + None => ValueRef::Null, + }; + values.push(value); + } + Ok(StructValueRef::RefList { + val: values, + fields: struct_type.clone(), + }) +} + +fn json_variant_ref_into_value_ref<'a>( + value: &JsonVariantRef<'a>, + expected_type: &ConcreteDataType, +) -> Result> { + let value = match (value, expected_type) { + (JsonVariantRef::Null, _) | (_, ConcreteDataType::Null(_)) => ValueRef::Null, + (JsonVariantRef::Object(object), ConcreteDataType::Struct(struct_type)) => { + ValueRef::Struct(json_object_ref_into_struct_value_ref(object, struct_type)?) + } + (JsonVariantRef::Bool(x), ConcreteDataType::Boolean(_)) => ValueRef::Boolean(*x), + (JsonVariantRef::Number(x), ConcreteDataType::UInt64(_)) => { + let Some(x) = x.as_u64() else { + return TryFromValueSnafu { + reason: format!("unable to convert {x:?} to UInt64"), + } + .fail(); + }; + ValueRef::UInt64(x) + } + (JsonVariantRef::Number(x), ConcreteDataType::Int64(_)) => { + let x = match x { + JsonNumber::PosInt(x) => i64::try_from(*x).ok(), + JsonNumber::NegInt(x) => Some(*x), + JsonNumber::Float(_) => None, + }; + let Some(x) = x else { + return TryFromValueSnafu { + reason: format!("unable to convert {x:?} to Int64"), + } + .fail(); + }; + ValueRef::Int64(x) + } + (JsonVariantRef::Number(JsonNumber::PosInt(x)), ConcreteDataType::Float64(_)) => { + ValueRef::Float64((*x as f64).into()) + } + (JsonVariantRef::Number(JsonNumber::NegInt(x)), ConcreteDataType::Float64(_)) => { + ValueRef::Float64((*x as f64).into()) + } + (JsonVariantRef::Number(JsonNumber::Float(x)), ConcreteDataType::Float64(_)) => { + ValueRef::Float64(*x) + } + (JsonVariantRef::String(x), ConcreteDataType::String(_)) => ValueRef::String(x), + (value, expected_type) => { + return TryFromValueSnafu { + reason: format!("unable to convert json value {value:?} to {expected_type}"), + } + .fail(); + } + }; + Ok(value) +} + +fn remainder_ref<'a>( + object: &BTreeMap<&'a str, JsonVariantRef<'a>>, + explicit_type: &JsonNativeType, +) -> Result>> { + let JsonNativeType::Object(fields) = explicit_type else { + return UnexpectedSnafu { + reason: "JSON2 explicit type must be an object", + } + .fail(); + }; + let mut remainder = BTreeMap::new(); + for (&name, value) in object { + match fields.get(name) { + Some(data_type @ JsonNativeType::Object(_)) => match value { + JsonVariantRef::Null => {} + JsonVariantRef::Object(object) => { + let child = remainder_ref(object, data_type)?; + if !child.is_empty() { + remainder.insert(name, JsonVariantRef::Object(child)); + } + } + _ => { + return TryFromValueSnafu { + reason: "expected json object value".to_string(), + } + .fail(); + } + }, + // A non-object entry in the explicit type tree is an explicit leaf and is + // already written to the Struct builder. So here does nothing. + Some(_) => {} + None => { + remainder.insert(name, value.clone()); + } + } + } + Ok(remainder) +} + fn split_to_explicit( value: JsonVariant, explicit_type: &JsonNativeType, @@ -611,11 +827,19 @@ impl MutableVector for JsonVectorBuilder { } fn to_vector(&mut self) -> VectorRef { - self.try_build().unwrap_or_else(|e| panic!("{:?}", e)) + self.try_build().unwrap_or_else(|e| { + // Just try to avoid panicking here. + common_telemetry::error!(e; "Unable to build JSON2 vector"); + Arc::new(NullVector::new(self.len())) + }) } fn to_vector_cloned(&self) -> VectorRef { - self.clone().to_vector() + self.state.try_build_cloned().unwrap_or_else(|e| { + // Just try to avoid panicking here. + common_telemetry::error!(e; "Unable to build JSON2 vector"); + Arc::new(NullVector::new(self.len())) + }) } fn try_push_value_ref(&mut self, value: &ValueRef) -> Result<()> { @@ -772,7 +996,7 @@ mod tests { } #[test] - fn test_zero_budget_builder_uses_fixed_schema_and_remainder() -> Result<()> { + fn test_zero_budget_builder_uses_explicit_only_schema_and_remainder() -> Result<()> { let settings = JsonSettings::try_new( vec![ JsonTypeHint { @@ -800,6 +1024,10 @@ mod tests { Some(0), )?; let mut builder = JsonVectorBuilder::with_settings(&settings, 2); + assert!(matches!( + &builder.state, + JsonVectorBuilderState::ExplicitOnly { .. } + )); let values = [ json!({ "kind": "record", diff --git a/src/datatypes/src/vectors/json/variant.rs b/src/datatypes/src/vectors/json/variant.rs index 75d87c99861..9f90f0aaacc 100644 --- a/src/datatypes/src/vectors/json/variant.rs +++ b/src/datatypes/src/vectors/json/variant.rs @@ -24,7 +24,7 @@ use parquet_variant_json::VariantToJson; use snafu::ResultExt; use crate::error::{ArrowComputeSnafu, Result}; -use crate::json::value::{JsonNumber, JsonVariant, decode_json_variant}; +use crate::json::value::{JsonNumber, JsonVariant, JsonVariantRef, decode_json_variant}; /// Returns the canonical Arrow field for an unshredded Parquet Variant array. pub fn variant_field(name: impl Into, nullable: bool) -> Field { @@ -110,6 +110,52 @@ pub(super) fn append_json_variant( Ok(()) } +pub(super) fn append_json_variant_ref( + builder: &mut impl VariantBuilderExt, + value: &JsonVariantRef<'_>, +) -> std::result::Result<(), ArrowError> { + match value { + JsonVariantRef::Null => builder.append_value(Variant::Null), + JsonVariantRef::Bool(value) => builder.append_value(*value), + JsonVariantRef::Number(JsonNumber::PosInt(value)) => { + if let Ok(value) = i64::try_from(*value) { + builder.append_value(value); + } else { + append_large_u64(builder, *value)?; + } + } + JsonVariantRef::Number(JsonNumber::NegInt(value)) => builder.append_value(*value), + JsonVariantRef::Number(JsonNumber::Float(value)) => { + if value.0.is_finite() { + builder.append_value(value.0) + } else { + builder.append_value("NaN") + } + } + JsonVariantRef::String(value) => builder.append_value(*value), + JsonVariantRef::Array(values) => { + let mut list = builder.try_new_list()?; + for value in values { + append_json_variant_ref(&mut list, value)?; + } + list.finish(); + } + JsonVariantRef::Object(values) => { + let mut object = builder.try_new_object()?; + for (name, value) in values { + append_json_variant_ref(&mut ObjectFieldBuilder::new(name, &mut object), value)?; + } + object.finish(); + } + JsonVariantRef::Variant(value) => { + let value = decode_json_variant(value) + .map_err(|e| ArrowError::JsonError(format!("Failed to decode JSONB: {e}")))?; + append_json_value(builder, &value)?; + } + } + Ok(()) +} + fn append_json_value( builder: &mut impl VariantBuilderExt, value: &serde_json::Value, From a502dfdefd59cf0ff8f3b2fdf3d1a8c79236bf95 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Mon, 24 Aug 2026 06:09:44 +0000 Subject: [PATCH 69/92] fix(mito2): fence checkpoints during region transitions (#8847) * fix: fence checkpoints during region transitions Signed-off-by: WenyXu * test(datanode): fix transient downgrade setup Signed-off-by: WenyXu * test(mito2): fix checkpoint lifecycle test setup Signed-off-by: WenyXu * test(mito2): cover cancelled downgrade waiter retry Signed-off-by: WenyXu * fix(mito2): fence direct follower transitions Signed-off-by: WenyXu * test: trim checkpoint transition coverage Signed-off-by: WenyXu * refactor(mito2): clarify checkpoint task lifecycle Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- .../src/heartbeat/handler/open_region.rs | 61 +++- .../src/heartbeat/handler/sync_region.rs | 50 +++ src/datanode/src/region_server.rs | 13 +- .../region_migration/open_candidate_region.rs | 8 +- .../repartition/group/sync_region.rs | 49 ++- src/meta-srv/src/procedure/test_util.rs | 19 +- src/mito2/src/engine/set_role_state_test.rs | 328 +++++++++++++++++- src/mito2/src/engine/staging_test.rs | 99 +++++- src/mito2/src/error.rs | 22 +- src/mito2/src/manifest/checkpointer.rs | 84 +++-- src/mito2/src/manifest/manager.rs | 34 +- src/mito2/src/manifest/storage/delta.rs | 18 +- src/mito2/src/manifest/tests/checkpoint.rs | 198 ++++++++++- src/mito2/src/region.rs | 25 +- src/mito2/src/test_util.rs | 133 ++++++- src/mito2/src/worker/handle_enter_staging.rs | 4 + 16 files changed, 1083 insertions(+), 62 deletions(-) diff --git a/src/datanode/src/heartbeat/handler/open_region.rs b/src/datanode/src/heartbeat/handler/open_region.rs index 89ce857ce4d..608fab0d90e 100644 --- a/src/datanode/src/heartbeat/handler/open_region.rs +++ b/src/datanode/src/heartbeat/handler/open_region.rs @@ -98,6 +98,8 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; + use common_error::ext::{BoxedError, RetryHint}; + use common_error::status_code::StatusCode; use common_meta::RegionIdent; use common_meta::heartbeat::handler::{HandleControl, HeartbeatResponseHandler}; use common_meta::heartbeat::mailbox::MessageMeta; @@ -105,14 +107,21 @@ mod tests { use common_meta::kv_backend::memory::MemoryKvBackend; use mito2::config::MitoConfig; use mito2::engine::MITO_ENGINE_NAME; + use mito2::error::ManifestDeltaNotFoundSnafu; use mito2::test_util::{CreateRequestBuilder, TestEnv}; + use object_store::{Error as ObjectStoreError, ErrorKind}; + use snafu::IntoError; use store_api::path_utils::table_dir; use store_api::region_request::{RegionCloseRequest, RegionRequest, RegionRequirements}; use store_api::storage::RegionId; - use crate::heartbeat::handler::RegionHeartbeatResponseHandler; + use super::OpenRegionsHandler; + use crate::error::{self, HandleRegionRequestSnafu}; use crate::heartbeat::handler::tests::HeartbeatResponseTestEnv; - use crate::tests::mock_region_server; + use crate::heartbeat::handler::{ + HandlerContext, InstructionHandler, RegionHeartbeatResponseHandler, + }; + use crate::tests::{MockRegionEngine, mock_region_server}; fn open_regions_instruction( region_ids: impl IntoIterator, @@ -198,4 +207,52 @@ mod tests { assert!(engine.is_region_exists(region_id)); assert!(engine.is_region_exists(region_id1)); } + + #[tokio::test] + async fn test_open_regions_preserves_manifest_delta_not_found_retry_hint() { + let retryable_region = RegionId::new(1024, 1); + let non_retryable_region = RegionId::new(1024, 2); + let (engine, _) = MockRegionEngine::with_mock_fn( + MITO_ENGINE_NAME, + Box::new(move |region_id, _request| { + if region_id == retryable_region { + let manifest_error = ManifestDeltaNotFoundSnafu { + version: 1_u64, + path: "manifest/00000000000000000001.json", + } + .into_error(ObjectStoreError::new( + ErrorKind::NotFound, + "mock listed manifest delta not found", + )); + return Err(HandleRegionRequestSnafu { region_id } + .into_error(BoxedError::new(manifest_error))); + } + + error::RegionNotFoundSnafu { region_id }.fail() + }), + ); + let mut region_server = mock_region_server(); + region_server.register_engine(engine); + let ctx = HandlerContext::new_for_test(region_server, Arc::new(MemoryKvBackend::new())); + let Instruction::OpenRegions(open_regions) = + open_regions_instruction([retryable_region, non_retryable_region], "test") + else { + unreachable!() + }; + + // Serial execution makes the first error selection deterministic. + let reply = OpenRegionsHandler { + open_region_parallelism: 1, + } + .handle(&ctx, open_regions) + .await + .unwrap() + .expect_open_regions_reply(); + + assert!(!reply.result); + let error = reply.error.unwrap(); + assert_eq!(StatusCode::StorageUnavailable, error.code); + assert_eq!(RetryHint::Retryable, error.retry_hint); + assert!(error.message.contains("00000000000000000001.json")); + } } diff --git a/src/datanode/src/heartbeat/handler/sync_region.rs b/src/datanode/src/heartbeat/handler/sync_region.rs index 48687541137..f8685f28470 100644 --- a/src/datanode/src/heartbeat/handler/sync_region.rs +++ b/src/datanode/src/heartbeat/handler/sync_region.rs @@ -103,11 +103,18 @@ impl SyncRegionHandler { mod tests { use std::sync::Arc; + use common_error::ext::{BoxedError, RetryHint}; + use common_error::status_code::StatusCode; use common_meta::kv_backend::memory::MemoryKvBackend; + use mito2::engine::MITO_ENGINE_NAME; + use mito2::error::ManifestDeltaNotFoundSnafu; + use object_store::{Error as ObjectStoreError, ErrorKind}; + use snafu::IntoError; use store_api::metric_engine_consts::METRIC_ENGINE_NAME; use store_api::region_engine::{RegionRole, SyncRegionFromRequest}; use store_api::storage::RegionId; + use crate::error::HandleRegionRequestSnafu; use crate::heartbeat::handler::sync_region::SyncRegionHandler; use crate::heartbeat::handler::{HandlerContext, InstructionHandler}; use crate::tests::{MockRegionEngine, mock_region_server}; @@ -201,4 +208,47 @@ mod tests { assert!(reply[0].ready); assert!(reply[0].error.is_none()); } + + #[tokio::test] + async fn test_handle_sync_region_preserves_manifest_delta_not_found_retry_hint() { + let mock_region_server = mock_region_server(); + let region_id = RegionId::new(1024, 1); + let (mock_engine, _) = MockRegionEngine::with_custom_apply_fn(MITO_ENGINE_NAME, |engine| { + engine.mock_role = Some(Some(RegionRole::Leader)); + engine.handle_sync_region_mock_fn = Some(Box::new(|region_id, _request| { + let manifest_error = ManifestDeltaNotFoundSnafu { + version: 1_u64, + path: "manifest/00000000000000000001.json", + } + .into_error(ObjectStoreError::new( + ErrorKind::NotFound, + "mock listed manifest delta not found", + )); + Err(HandleRegionRequestSnafu { region_id } + .into_error(BoxedError::new(manifest_error))) + })); + }); + mock_region_server.register_test_region(region_id, mock_engine); + + let handler_context = + HandlerContext::new_for_test(mock_region_server, Arc::new(MemoryKvBackend::new())); + let sync_region = common_meta::instruction::SyncRegion { + region_id, + request: SyncRegionFromRequest::from_manifest(Default::default()), + }; + + let reply = SyncRegionHandler + .handle(&handler_context, vec![sync_region]) + .await + .unwrap() + .expect_sync_regions_reply(); + + assert_eq!(1, reply.len()); + assert!(reply[0].exists); + assert!(!reply[0].ready); + let error = reply[0].error.as_ref().unwrap(); + assert_eq!(StatusCode::StorageUnavailable, error.code); + assert_eq!(RetryHint::Retryable, error.retry_hint); + assert!(error.message.contains("00000000000000000001.json")); + } } diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 55dc6a670a6..57facfd53ba 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -1330,11 +1330,8 @@ impl RegionServerInner { } if !errors.is_empty() { - return error::UnexpectedSnafu { - // Returns the first error. - violated: format!("Failed to open batch regions: {:?}", errors[0]), - } - .fail(); + // Preserve the first region error so callers can honor its status code and retry hint. + return Err(errors.swap_remove(0)).context(HandleBatchOpenRequestSnafu); } Ok(open_regions) @@ -1979,7 +1976,7 @@ mod tests { use std::sync::Arc; use api::v1::{Rows, SemanticType}; - use common_error::ext::ErrorExt; + use common_error::ext::{ErrorExt, RetryHint}; use common_recordbatch::RecordBatches; use common_recordbatch::adapter::{RecordBatchMetrics, RegionWatermarkEntry}; use datatypes::prelude::{ConcreteDataType, VectorRef}; @@ -2514,7 +2511,9 @@ mod tests { ) .await .unwrap_err(); - assert_eq!(err.status_code(), StatusCode::Unexpected); + assert_matches!(&err, error::Error::HandleBatchOpenRequest { .. }); + assert_eq!(err.status_code(), StatusCode::RegionNotFound); + assert_eq!(err.retry_hint(), RetryHint::NonRetryable); } struct CurrentEngineTest { diff --git a/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs b/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs index 0cb1131e545..bc21746a046 100644 --- a/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs +++ b/src/meta-srv/src/procedure/region_migration/open_candidate_region.rs @@ -487,10 +487,14 @@ mod tests { .await; send_mock_reply(mailbox, rx, |id| { - Ok(new_open_region_reply( + Ok(new_open_region_reply_with_error( id, false, - Some("test mocked".to_string()), + Some(InstructionError { + code: StatusCode::StorageUnavailable, + message: "test mocked".to_string(), + retry_hint: RetryHint::Retryable, + }), )) }); diff --git a/src/meta-srv/src/procedure/repartition/group/sync_region.rs b/src/meta-srv/src/procedure/repartition/group/sync_region.rs index 8ce5ea3dec3..dd55db73bcc 100644 --- a/src/meta-srv/src/procedure/repartition/group/sync_region.rs +++ b/src/meta-srv/src/procedure/repartition/group/sync_region.rs @@ -350,6 +350,9 @@ impl SyncRegion { mod tests { use std::assert_matches; + use common_error::ext::RetryHint; + use common_error::status_code::StatusCode; + use common_meta::instruction::InstructionError; use common_meta::peer::Peer; use common_meta::rpc::router::{Region, RegionRoute}; use store_api::region_engine::SyncRegionFromRequest; @@ -359,7 +362,9 @@ mod tests { use crate::procedure::repartition::group::GroupPrepareResult; use crate::procedure::repartition::group::sync_region::SyncRegion; use crate::procedure::repartition::test_util::{TestingEnv, new_persistent_context}; - use crate::procedure::test_util::{new_sync_region_reply, send_mock_reply}; + use crate::procedure::test_util::{ + new_sync_region_reply, new_sync_region_reply_with_error, send_mock_reply, + }; use crate::service::mailbox::Channel; #[test] @@ -455,4 +460,46 @@ mod tests { let err = sync_region.sync_regions(&mut ctx).await.unwrap_err(); assert_matches!(err, Error::RetryLater { .. }); } + + #[tokio::test] + async fn test_sync_regions_retryable_instruction_error() { + let mut env = TestingEnv::new(); + let table_id = 1024; + let region_id = RegionId::new(table_id, 3); + let mut persistent_context = new_persistent_context(table_id, vec![], vec![]); + persistent_context.group_prepare_result = Some(test_prepare_result(table_id)); + + let (tx, rx) = tokio::sync::mpsc::channel(1); + env.mailbox_ctx + .insert_heartbeat_response_receiver(Channel::Datanode(1), tx) + .await; + send_mock_reply(env.mailbox_ctx.mailbox().clone(), rx, move |id| { + Ok(new_sync_region_reply_with_error( + id, + region_id, + false, + true, + Some(InstructionError { + code: StatusCode::StorageUnavailable, + message: "manifest delta disappeared".to_string(), + retry_hint: RetryHint::Retryable, + }), + )) + }); + + let mut ctx = env.create_context(persistent_context); + let sync_region = SyncRegion { + region_routes: vec![RegionRoute { + region: Region { + id: region_id, + ..Default::default() + }, + leader_peer: Some(Peer::empty(1)), + ..Default::default() + }], + }; + + let err = sync_region.sync_regions(&mut ctx).await.unwrap_err(); + assert_matches!(err, Error::RetryLater { .. }); + } } diff --git a/src/meta-srv/src/procedure/test_util.rs b/src/meta-srv/src/procedure/test_util.rs index ea8c3d4aec4..4208c6e52e7 100644 --- a/src/meta-srv/src/procedure/test_util.rs +++ b/src/meta-srv/src/procedure/test_util.rs @@ -288,6 +288,23 @@ pub fn new_sync_region_reply( ready: bool, exists: bool, error: Option, +) -> MailboxMessage { + new_sync_region_reply_with_error( + id, + region_id, + ready, + exists, + legacy_instruction_error(error), + ) +} + +/// Generates a [InstructionReply::SyncRegions] reply with a structured error. +pub fn new_sync_region_reply_with_error( + id: u64, + region_id: RegionId, + ready: bool, + exists: bool, + error: Option, ) -> MailboxMessage { MailboxMessage { id, @@ -301,7 +318,7 @@ pub fn new_sync_region_reply( region_id, ready, exists, - error: legacy_instruction_error(error), + error, }, ]))) .unwrap(), diff --git a/src/mito2/src/engine/set_role_state_test.rs b/src/mito2/src/engine/set_role_state_test.rs index 40e03b063a6..1d0f3c188b8 100644 --- a/src/mito2/src/engine/set_role_state_test.rs +++ b/src/mito2/src/engine/set_role_state_test.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::time::Duration; + use api::v1::Rows; use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; @@ -20,12 +22,16 @@ use store_api::region_engine::{ SettableRegionRoleState, }; use store_api::region_request::{ - EnterStagingRequest, RegionPutRequest, RegionRequest, StagingPartitionDirective, + EnterStagingRequest, RegionFlushRequest, RegionPutRequest, RegionRequest, + StagingPartitionDirective, }; use store_api::storage::RegionId; use crate::config::MitoConfig; -use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema}; +use crate::region::{RegionLeaderState, RegionRoleState}; +use crate::test_util::{ + CheckpointTaskBlocker, CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema, +}; /// Helper function to assert a successful response with expected entry id fn assert_success_response(response: &SetRegionRoleStateResponse, expected_entry_id: u64) { @@ -212,6 +218,324 @@ async fn test_write_downgrading_region_with_format(flat_format: bool) { assert_eq!(err.status_code(), StatusCode::RegionNotReady) } +#[tokio::test(flavor = "multi_thread")] +async fn test_downgrading_waits_for_checkpoint_and_stops_new_checkpoints() { + let (blocker, mock_layer) = CheckpointTaskBlocker::block_cleanup(); + let mut env = TestEnv::new().await.with_mock_layer(mock_layer); + let engine = env + .create_engine(MitoConfig { + manifest_checkpoint_distance: 1, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas.clone(), + rows: build_rows(0, 1), + }, + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked()) + .await + .expect("checkpoint cleanup did not start"); + + // Leave one memtable for the final flush after entering Downgrading. + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas.clone(), + rows: build_rows(1, 2), + }, + ) + .await; + + let region = engine.get_region(region_id).unwrap(); + let wait_started = region + .manifest_ctx + .manifest_manager + .read() + .await + .checkpointer() + .pending_checkpoint_wait_started(); + // Register before starting the transition so the notification cannot be lost. + let wait_started = wait_started.notified(); + let cloned_engine = engine.clone(); + let downgrade = tokio::spawn(async move { + cloned_engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::DowngradingLeader) + .await + }); + tokio::time::timeout(Duration::from_secs(5), wait_started) + .await + .expect("downgrade did not start waiting for the checkpoint"); + assert_eq!( + RegionRoleState::Leader(RegionLeaderState::Downgrading), + region.state() + ); + assert!( + !downgrade.is_finished(), + "downgrade returned before checkpoint cleanup finished" + ); + + blocker.release(); + assert_success_response(&downgrade.await.unwrap().unwrap(), 2); + + // Final flush publishes a normal delta, but Downgrading must not start a + // checkpoint after the barrier. + blocker.arm_next_close(); + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + assert!( + !region + .manifest_ctx + .manifest_manager + .read() + .await + .checkpointer() + .is_doing_checkpoint(), + "final flush started a checkpoint while Downgrading" + ); + assert_eq!(2, region.manifest_ctx.manifest().await.manifest_version); + assert_eq!( + 1, + region + .manifest_ctx + .manifest_manager + .read() + .await + .checkpointer() + .last_checkpoint_version() + ); + + // Leaving Downgrading removes the scheduling restriction. The next normal + // manifest update can checkpoint the accumulated deltas. + engine + .set_region_role(region_id, RegionRole::Leader) + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas, + rows: build_rows(2, 3), + }, + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked()) + .await + .expect("checkpoint scheduling did not resume after leaving Downgrading"); + blocker.release(); + region + .manifest_ctx + .manifest_manager + .write() + .await + .wait_for_pending_checkpoint() + .await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_direct_follower_waits_for_pending_checkpoint() { + let (blocker, mock_layer) = CheckpointTaskBlocker::block_cleanup(); + let mut env = TestEnv::new().await.with_mock_layer(mock_layer); + let engine = env + .create_engine(MitoConfig { + manifest_checkpoint_distance: 1, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas, + rows: build_rows(0, 1), + }, + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked()) + .await + .expect("checkpoint cleanup did not start"); + + let region = engine.get_region(region_id).unwrap(); + let wait_started = region + .manifest_ctx + .manifest_manager + .read() + .await + .checkpointer() + .pending_checkpoint_wait_started(); + // The no-flush downgrade path requests Follower directly. It must still + // wait for checkpoint cleanup before replying to the caller. + let wait_started = wait_started.notified(); + let cloned_engine = engine.clone(); + let set_follower = tokio::spawn(async move { + cloned_engine + .set_region_role_state_gracefully(region_id, SettableRegionRoleState::Follower) + .await + }); + tokio::time::timeout(Duration::from_secs(5), wait_started) + .await + .expect("follower transition did not start waiting for the checkpoint"); + assert_eq!(RegionRoleState::Follower, region.state()); + assert!( + !set_follower.is_finished(), + "follower transition returned before checkpoint cleanup finished" + ); + + blocker.release(); + assert_success_response(&set_follower.await.unwrap().unwrap(), 1); +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_retried_downgrade_waits_after_first_request_is_cancelled() { + let (blocker, mock_layer) = CheckpointTaskBlocker::block_cleanup(); + let mut env = TestEnv::new().await.with_mock_layer(mock_layer); + let engine = env + .create_engine(MitoConfig { + manifest_checkpoint_distance: 1, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas, + rows: build_rows(0, 1), + }, + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked()) + .await + .expect("checkpoint cleanup did not start"); + + let region = engine.get_region(region_id).unwrap(); + let first_wait_started = region + .manifest_ctx + .manifest_manager + .read() + .await + .checkpointer() + .pending_checkpoint_wait_started(); + // Register before starting the transition so the notification cannot be lost. + let first_wait_started = first_wait_started.notified(); + // Call the region transition directly so aborting this task cancels the + // actual checkpoint waiter. The engine API submits the same transition to + // a detached worker task, so aborting its caller only drops the reply receiver. + let first_region = region.clone(); + let first_downgrade = tokio::spawn(async move { + first_region + .set_role_state_gracefully(SettableRegionRoleState::DowngradingLeader) + .await + }); + tokio::time::timeout(Duration::from_secs(5), first_wait_started) + .await + .expect("first downgrade did not start waiting for the checkpoint"); + assert_eq!( + RegionRoleState::Leader(RegionLeaderState::Downgrading), + region.state() + ); + assert!( + !first_downgrade.is_finished(), + "first downgrade did not wait for checkpoint cleanup" + ); + + // Cancelling the first caller must not remove the checkpoint handle from + // the manifest manager. A retried migration request observes Downgrading + // and must wait for the same checkpoint task before it can proceed. + first_downgrade.abort(); + assert!(first_downgrade.await.unwrap_err().is_cancelled()); + assert_eq!( + RegionRoleState::Leader(RegionLeaderState::Downgrading), + region.state() + ); + + let second_wait_started = region + .manifest_ctx + .manifest_manager + .read() + .await + .checkpointer() + .pending_checkpoint_wait_started(); + // Register before spawning the retry so the notification cannot be lost. + let second_wait_started = second_wait_started.notified(); + let retry_region = region.clone(); + let retry_downgrade = tokio::spawn(async move { + retry_region + .set_role_state_gracefully(SettableRegionRoleState::DowngradingLeader) + .await + }); + tokio::time::timeout(Duration::from_secs(5), second_wait_started) + .await + .expect("retried downgrade did not start waiting for the checkpoint"); + assert!( + !retry_downgrade.is_finished(), + "retried downgrade did not wait for checkpoint cleanup" + ); + + blocker.release(); + retry_downgrade.await.unwrap().unwrap(); +} + #[tokio::test] async fn test_unified_state_transitions() { test_unified_state_transitions_with_format(false).await; diff --git a/src/mito2/src/engine/staging_test.rs b/src/mito2/src/engine/staging_test.rs index 71abd67eb22..12263976aa2 100644 --- a/src/mito2/src/engine/staging_test.rs +++ b/src/mito2/src/engine/staging_test.rs @@ -48,7 +48,9 @@ use crate::manifest::action::{ use crate::region::{RegionLeaderState, RegionRoleState, parse_partition_expr}; use crate::request::WorkerRequest; use crate::sst::FormatType; -use crate::test_util::{CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema}; +use crate::test_util::{ + CheckpointTaskBlocker, CreateRequestBuilder, TestEnv, build_rows, put_rows, rows_schema, +}; fn range_expr(col_name: &str, start: i64, end: i64) -> PartitionExpr { col(col_name) @@ -1101,6 +1103,101 @@ async fn test_write_stall_on_enter_staging_with_format(flat_format: bool) { assert_eq!(expected, batches.pretty_print().unwrap()); } +#[tokio::test(flavor = "multi_thread")] +async fn test_enter_staging_waits_for_pending_checkpoint() { + for block_last_checkpoint_write in [false, true] { + let partition_directive = + StagingPartitionDirective::UpdatePartitionExpr(default_partition_expr()); + let (blocker, mock_layer) = if block_last_checkpoint_write { + CheckpointTaskBlocker::block_last_checkpoint_write() + } else { + CheckpointTaskBlocker::block_cleanup() + }; + let mut env = TestEnv::new().await.with_mock_layer(mock_layer); + let engine = env + .create_engine(MitoConfig { + manifest_checkpoint_distance: 1, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let column_schemas = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + put_rows( + &engine, + region_id, + Rows { + schema: column_schemas, + rows: build_rows(0, 1), + }, + ) + .await; + engine + .handle_request( + region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked()) + .await + .expect("checkpoint task did not reach the configured block point"); + + let region = engine.get_region(region_id).unwrap(); + let wait_started = region + .manifest_ctx + .manifest_manager + .read() + .await + .checkpointer() + .pending_checkpoint_wait_started(); + // Register before starting the transition so the notification cannot be lost. + let wait_started = wait_started.notified(); + let cloned_engine = engine.clone(); + let enter_staging = tokio::spawn(async move { + cloned_engine + .handle_request( + region_id, + RegionRequest::EnterStaging(EnterStagingRequest { + partition_directive, + }), + ) + .await + }); + tokio::time::timeout(Duration::from_secs(5), wait_started) + .await + .expect("EnterStaging did not start waiting for the checkpoint"); + assert_eq!( + RegionRoleState::Leader(RegionLeaderState::EnteringStaging), + region.state() + ); + assert!( + !enter_staging.is_finished(), + "EnterStaging returned before the checkpoint task finished" + ); + + blocker.release(); + enter_staging.await.unwrap().unwrap(); + assert_eq!( + RegionRoleState::Leader(RegionLeaderState::Staging), + region.state() + ); + assert!( + !region + .manifest_ctx + .manifest_manager + .read() + .await + .checkpointer() + .is_doing_checkpoint() + ); + } +} + #[tokio::test] async fn test_enter_staging_clean_staging_manifest_error() { common_telemetry::init_default_ut_logging(); diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index a484d43a318..7708bc3a36e 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -67,6 +67,20 @@ pub enum Error { error: object_store::Error, }, + #[snafu(display( + "Manifest delta {} disappeared after it was listed, path: {}", + version, + path + ))] + ManifestDeltaNotFound { + version: ManifestVersion, + path: String, + #[snafu(source)] + error: object_store::Error, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Fail to compress object by {}, path: {}", compress_type, path))] CompressObject { compress_type: CompressionType, @@ -1359,6 +1373,7 @@ impl Error { pub(crate) fn is_object_not_found(&self) -> bool { match self { Error::OpenDal { error, .. } => error.kind() == ErrorKind::NotFound, + Error::ManifestDeltaNotFound { .. } => true, _ => false, } } @@ -1400,7 +1415,9 @@ impl ErrorExt for Error { match self { DataTypeMismatch { source, .. } => source.status_code(), - OpenDal { .. } | ReadParquet { .. } => StatusCode::StorageUnavailable, + OpenDal { .. } | ManifestDeltaNotFound { .. } | ReadParquet { .. } => { + StatusCode::StorageUnavailable + } WriteWal { source, .. } | ReadWal { source, .. } | DeleteWal { source, .. } => { source.status_code() } @@ -1601,7 +1618,8 @@ impl ErrorExt for Error { | RegionStopped { .. } | RegionBusy { .. } | ManualCompactionAlreadyRunning { .. } - | FlushableRegionState { .. } => RetryHint::Retryable, + | FlushableRegionState { .. } + | ManifestDeltaNotFound { .. } => RetryHint::Retryable, OpenDal { error, .. } | DeleteSsts { error, .. } diff --git a/src/mito2/src/manifest/checkpointer.rs b/src/mito2/src/manifest/checkpointer.rs index 9dd35e189a6..232421a2621 100644 --- a/src/mito2/src/manifest/checkpointer.rs +++ b/src/mito2/src/manifest/checkpointer.rs @@ -14,11 +14,14 @@ use std::fmt::Debug; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU64, Ordering}; +use common_runtime::JoinHandle; use common_telemetry::{error, info, warn}; use store_api::storage::RegionId; use store_api::{MIN_VERSION, ManifestVersion}; +#[cfg(test)] +use tokio::sync::Notify; use crate::error::Result; use crate::manifest::action::{RegionCheckpoint, RegionManifest}; @@ -31,6 +34,9 @@ use crate::metrics::MANIFEST_OP_ELAPSED; pub(crate) struct Checkpointer { manifest_options: RegionManifestOptions, inner: Arc, + checkpoint_task: Option>, + #[cfg(test)] + pending_checkpoint_wait_started: Arc, } #[derive(Debug)] @@ -38,15 +44,10 @@ struct Inner { region_id: RegionId, manifest_store: ManifestObjectStore, last_checkpoint_version: AtomicU64, - is_doing_checkpoint: AtomicBool, } impl Inner { async fn do_checkpoint(&self, checkpoint: RegionCheckpoint) { - let _guard = scopeguard::guard(&self.is_doing_checkpoint, |x| { - x.store(false, Ordering::Relaxed); - }); - let _t = MANIFEST_OP_ELAPSED .with_label_values(&["checkpoint"]) .start_timer(); @@ -90,14 +91,6 @@ impl Inner { fn region_id(&self) -> RegionId { self.region_id } - - fn is_doing_checkpoint(&self) -> bool { - self.is_doing_checkpoint.load(Ordering::Relaxed) - } - - fn set_doing_checkpoint(&self) { - self.is_doing_checkpoint.store(true, Ordering::Relaxed); - } } impl Checkpointer { @@ -113,8 +106,10 @@ impl Checkpointer { region_id, manifest_store, last_checkpoint_version: AtomicU64::new(last_checkpoint_version), - is_doing_checkpoint: AtomicBool::new(false), }), + checkpoint_task: None, + #[cfg(test)] + pending_checkpoint_wait_started: Arc::new(Notify::new()), } } @@ -140,7 +135,19 @@ impl Checkpointer { /// Check if it's needed to do checkpoint for the region by the checkpoint distance. /// If needed, and there's no currently running checkpoint task, it will start a new checkpoint /// task running in the background. - pub(crate) fn maybe_do_checkpoint(&self, manifest: &RegionManifest) { + pub(crate) async fn maybe_do_checkpoint(&mut self, manifest: &RegionManifest) { + if self + .checkpoint_task + .as_ref() + .is_some_and(|handle| !handle.is_finished()) + { + return; + } + + // Reap a completed task before checking whether to start the next one. + // This keeps the handle as the single source of truth for task state. + self.wait_for_pending_checkpoint().await; + if self.manifest_options.checkpoint_distance == 0 { return; } @@ -152,12 +159,6 @@ impl Checkpointer { return; } - // We can simply check whether there's a running checkpoint task like this, all because of - // the caller of this function is ran single threaded, inside the lock of RegionManifestManager. - if self.inner.is_doing_checkpoint() { - return; - } - let start_version = if last_checkpoint_version == 0 { // Checkpoint version can't be zero by implementation. // So last checkpoint version is zero means no last checkpoint. @@ -181,17 +182,44 @@ impl Checkpointer { self.do_checkpoint(checkpoint); } - fn do_checkpoint(&self, checkpoint: RegionCheckpoint) { - self.inner.set_doing_checkpoint(); - + fn do_checkpoint(&mut self, checkpoint: RegionCheckpoint) { let inner = self.inner.clone(); - common_runtime::spawn_global(async move { + self.checkpoint_task = Some(common_runtime::spawn_global(async move { inner.do_checkpoint(checkpoint).await; - }); + })); + } + + /// Waits for the current checkpoint task without removing its handle first. + /// + /// Keeping the handle in `self` while awaiting is important. If the caller is + /// cancelled, another lifecycle transition can still wait for the same task. + pub(crate) async fn wait_for_pending_checkpoint(&mut self) { + let Some(handle) = self.checkpoint_task.as_mut() else { + return; + }; + + #[cfg(test)] + self.pending_checkpoint_wait_started.notify_one(); + + let result = (&mut *handle).await; + // There is no cancellation point between observing completion and + // clearing the handle. + self.checkpoint_task = None; + + if let Err(e) = result { + warn!(e; "Failed to join checkpoint task for region {}", self.inner.region_id()); + } } #[cfg(test)] pub(crate) fn is_doing_checkpoint(&self) -> bool { - self.inner.is_doing_checkpoint() + self.checkpoint_task + .as_ref() + .is_some_and(|handle| !handle.is_finished()) + } + + #[cfg(test)] + pub(crate) fn pending_checkpoint_wait_started(&self) -> Arc { + self.pending_checkpoint_wait_started.clone() } } diff --git a/src/mito2/src/manifest/manager.rs b/src/mito2/src/manifest/manager.rs index fcc48395fd9..c6a0d3f3228 100644 --- a/src/mito2/src/manifest/manager.rs +++ b/src/mito2/src/manifest/manager.rs @@ -550,6 +550,28 @@ impl RegionManifestManager { &mut self, action_list: RegionMetaActionList, is_staging: bool, + ) -> Result { + self.update_inner(action_list, is_staging, !is_staging) + .await + } + + /// Updates the normal manifest without starting a checkpoint. + /// + /// This is only used while a leader is downgrading. The final flush still + /// publishes its manifest edit, but must not start cleanup after the + /// downgrade checkpoint barrier. + pub(crate) async fn update_normal_without_checkpoint( + &mut self, + action_list: RegionMetaActionList, + ) -> Result { + self.update_inner(action_list, false, false).await + } + + async fn update_inner( + &mut self, + action_list: RegionMetaActionList, + is_staging: bool, + allow_checkpoint: bool, ) -> Result { let _t = MANIFEST_OP_ELAPSED .with_label_values(&["update"]) @@ -616,13 +638,21 @@ impl RegionManifestManager { .checkpointer .update_manifest_removed_files(new_manifest)?; self.manifest = Arc::new(updated_manifest); - self.checkpointer - .maybe_do_checkpoint(self.manifest.as_ref()); + if allow_checkpoint { + self.checkpointer + .maybe_do_checkpoint(self.manifest.as_ref()) + .await; + } } Ok(version) } + /// Waits for an in-flight checkpoint to finish, including its cleanup. + pub(crate) async fn wait_for_pending_checkpoint(&mut self) { + self.checkpointer.wait_for_pending_checkpoint().await; + } + /// Clear deleted files from manifest's `removed_files` field without update version. Notice if datanode exit before checkpoint then new manifest by open region may still contain these deleted files, which is acceptable for gc process. pub fn clear_deleted_files(&mut self, deleted_files: Vec) { let mut manifest = (*self.manifest()).clone(); diff --git a/src/mito2/src/manifest/storage/delta.rs b/src/mito2/src/manifest/storage/delta.rs index 594b56ddaec..8f9962a3c69 100644 --- a/src/mito2/src/manifest/storage/delta.rs +++ b/src/mito2/src/manifest/storage/delta.rs @@ -26,7 +26,8 @@ use tokio::sync::Semaphore; use crate::cache::manifest_cache::ManifestCache; use crate::error::{ - CompressObjectSnafu, DecompressObjectSnafu, InvalidScanIndexSnafu, OpenDalSnafu, Result, + CompressObjectSnafu, DecompressObjectSnafu, InvalidScanIndexSnafu, ManifestDeltaNotFoundSnafu, + OpenDalSnafu, Result, }; use crate::manifest::storage::size_tracker::Tracker; use crate::manifest::storage::utils::{ @@ -212,11 +213,16 @@ impl DeltaStorage { // Fetch from remote object store let compress_type = file_compress_type(entry.name()); - let bytes = self - .object_store - .read(entry.path()) - .await - .context(OpenDalSnafu)?; + let bytes = match self.object_store.read(entry.path()).await { + Ok(bytes) => bytes, + Err(error) if error.kind() == ErrorKind::NotFound => { + return Err(error).context(ManifestDeltaNotFoundSnafu { + version: *v, + path: entry.path(), + }); + } + Err(error) => return Err(error).context(OpenDalSnafu), + }; let data = compress_type .decode(bytes) .await diff --git a/src/mito2/src/manifest/tests/checkpoint.rs b/src/mito2/src/manifest/tests/checkpoint.rs index 07f879701ad..86679630ac7 100644 --- a/src/mito2/src/manifest/tests/checkpoint.rs +++ b/src/mito2/src/manifest/tests/checkpoint.rs @@ -18,13 +18,16 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; use common_datasource::compression::CompressionType; +use common_error::ext::{ErrorExt, RetryHint}; +use common_error::status_code::StatusCode; use object_store::layers::mock::{ - Error as MockError, ErrorKind, MockLayerBuilder, OpDelete, Result as MockResult, oio, + Buffer, Error as MockError, ErrorKind, MockLayer, MockLayerBuilder, OpDelete, + Result as MockResult, oio, }; use store_api::storage::{FileId, RegionId}; use strum::IntoEnumIterator; -use crate::error::Error::ChecksumMismatch; +use crate::error::Error::{ChecksumMismatch, ManifestDeltaNotFound}; use crate::manifest::action::{ RegionCheckpoint, RegionEdit, RegionMetaAction, RegionMetaActionList, }; @@ -33,7 +36,7 @@ use crate::manifest::storage::checkpoint::CheckpointMetadata; use crate::manifest::storage::is_delta_file; use crate::manifest::tests::utils::basic_region_metadata; use crate::sst::file::FileMeta; -use crate::test_util::TestEnv; +use crate::test_util::{CheckpointTaskBlocker, TestEnv}; async fn build_manager( checkpoint_distance: u64, @@ -85,6 +88,31 @@ fn nop_action() -> RegionMetaActionList { })]) } +struct NotFoundReader; + +impl oio::Read for NotFoundReader { + async fn read(&mut self) -> MockResult { + Err(MockError::new( + ErrorKind::NotFound, + "mock listed manifest delta not found", + )) + } +} + +fn fail_manifest_delta_reads_layer() -> MockLayer { + MockLayerBuilder::default() + .reader_factory(Arc::new(|path, _args, inner| { + let file_name = path.rsplit('/').next().unwrap_or(path); + if is_delta_file(file_name) { + Box::new(NotFoundReader) + } else { + inner + } + })) + .build() + .unwrap() +} + #[tokio::test] async fn manager_without_checkpoint() { let (_env, mut manager) = build_manager(0, CompressionType::Uncompressed).await; @@ -589,3 +617,167 @@ async fn checkpoint_advances_and_recovery_works_when_delete_fails() { .expect("manifest should be recoverable"); assert_eq!(reopened.manifest().manifest_version, 10); } + +#[tokio::test] +async fn open_preserves_listed_delta_not_found_retry_hint() { + let env = TestEnv::new() + .await + .with_mock_layer(fail_manifest_delta_reads_layer()); + let metadata = Arc::new(basic_region_metadata()); + let mut manager = env + .create_manifest_manager(CompressionType::Uncompressed, 0, Some(metadata.clone())) + .await + .unwrap() + .unwrap(); + manager.stop().await; + + let error = env + .create_manifest_manager(CompressionType::Uncompressed, 0, None) + .await + .expect_err("reopen must fail on the mocked delta read"); + assert_matches!( + &error, + ManifestDeltaNotFound { + version: 0, + path, + error, + .. + } if path.ends_with("00000000000000000000.json") + && error.kind() == object_store::ErrorKind::NotFound + ); + assert_eq!(StatusCode::StorageUnavailable, error.status_code()); + assert_eq!(RetryHint::Retryable, error.retry_hint()); +} + +#[tokio::test] +async fn install_preserves_listed_delta_not_found_retry_hint() { + let env = TestEnv::new() + .await + .with_mock_layer(fail_manifest_delta_reads_layer()); + let metadata = Arc::new(basic_region_metadata()); + let mut manager = env + .create_manifest_manager(CompressionType::Uncompressed, 0, Some(metadata)) + .await + .unwrap() + .unwrap(); + let mut store = manager.store(); + store + .save(1, &nop_action().encode().unwrap(), false) + .await + .unwrap(); + + let error = manager.install_manifest_to(1).await.unwrap_err(); + assert_matches!( + &error, + ManifestDeltaNotFound { + version: 1, + path, + .. + } if path.ends_with("00000000000000000001.json") + ); + assert_eq!(RetryHint::Retryable, error.retry_hint()); +} + +#[tokio::test] +async fn cancelled_waiter_keeps_pending_checkpoint_handle() { + let (blocker, mock_layer) = CheckpointTaskBlocker::block_cleanup(); + let env = TestEnv::new().await.with_mock_layer(mock_layer); + let metadata = Arc::new(basic_region_metadata()); + let mut manager = env + .create_manifest_manager(CompressionType::Uncompressed, 1, Some(metadata)) + .await + .unwrap() + .unwrap(); + + manager.update(nop_action(), false).await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked()) + .await + .expect("checkpoint cleanup did not start"); + + let wait_started = manager.checkpointer().pending_checkpoint_wait_started(); + let wait_started = wait_started.notified(); + let mut first_wait = Box::pin(manager.wait_for_pending_checkpoint()); + tokio::select! { + _ = wait_started => {} + _ = &mut first_wait => panic!("checkpoint waiter returned before cleanup finished"), + } + drop(first_wait); + assert!(manager.checkpointer().is_doing_checkpoint()); + + blocker.release(); + manager.wait_for_pending_checkpoint().await; + assert!(!manager.checkpointer().is_doing_checkpoint()); + + let (version, _) = manager + .store() + .load_last_checkpoint() + .await + .unwrap() + .expect("checkpoint must be published"); + assert_eq!(1, version); +} + +#[tokio::test] +async fn running_checkpoint_prevents_scheduling_another_checkpoint() { + let (blocker, mock_layer) = CheckpointTaskBlocker::block_cleanup(); + let env = TestEnv::new().await.with_mock_layer(mock_layer); + let metadata = Arc::new(basic_region_metadata()); + let mut manager = env + .create_manifest_manager(CompressionType::Uncompressed, 1, Some(metadata)) + .await + .unwrap() + .unwrap(); + + manager.update(nop_action(), false).await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked()) + .await + .expect("checkpoint cleanup did not start"); + + // The second update is eligible for checkpointing, but the first task still + // owns the pending handle and must prevent another task from being scheduled. + manager.update(nop_action(), false).await.unwrap(); + blocker.release(); + manager.wait_for_pending_checkpoint().await; + + assert_eq!(2, manager.manifest().manifest_version); + assert_eq!(1, manager.checkpointer().last_checkpoint_version()); +} + +#[tokio::test] +async fn completed_checkpoint_is_reaped_before_scheduling_the_next_one() { + let (blocker, mock_layer) = CheckpointTaskBlocker::block_cleanup(); + let env = TestEnv::new().await.with_mock_layer(mock_layer); + let metadata = Arc::new(basic_region_metadata()); + let mut manager = env + .create_manifest_manager(CompressionType::Uncompressed, 1, Some(metadata)) + .await + .unwrap() + .unwrap(); + + manager.update(nop_action(), false).await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked()) + .await + .expect("first checkpoint cleanup did not start"); + blocker.release(); + + tokio::time::timeout(Duration::from_secs(5), async { + while manager.checkpointer().is_doing_checkpoint() { + tokio::task::yield_now().await; + } + }) + .await + .expect("first checkpoint did not finish"); + assert_eq!(1, manager.checkpointer().last_checkpoint_version()); + + // Do not explicitly wait/reap the completed handle. The next eligible + // update must reap it before scheduling another checkpoint. + blocker.arm_next_close(); + manager.update(nop_action(), false).await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), blocker.wait_until_blocked()) + .await + .expect("second checkpoint cleanup did not start"); + blocker.release(); + manager.wait_for_pending_checkpoint().await; + + assert_eq!(2, manager.checkpointer().last_checkpoint_version()); +} diff --git a/src/mito2/src/region.rs b/src/mito2/src/region.rs index c5a9021bde5..76204da9cf6 100644 --- a/src/mito2/src/region.rs +++ b/src/mito2/src/region.rs @@ -483,6 +483,7 @@ impl MitoRegion { let mut manager: RwLockWriteGuard<'_, RegionManifestManager> = self.manifest_ctx.manifest_manager.write().await; let current_state = self.state(); + let mut wait_for_checkpoint = false; let hook_payload: Option = match state { SettableRegionRoleState::Leader => { @@ -545,10 +546,12 @@ impl MitoRegion { ); self.exit_staging()?; self.set_role(RegionRole::Follower); + wait_for_checkpoint = true; } RegionRoleState::Leader(_) => { info!("Demoting region {} from leader to follower", self.region_id); self.set_role(RegionRole::Follower); + wait_for_checkpoint = true; } RegionRoleState::Follower => { // Already in desired state - no-op @@ -568,14 +571,17 @@ impl MitoRegion { ); self.exit_staging()?; self.set_role(RegionRole::DowngradingLeader); + wait_for_checkpoint = true; } RegionRoleState::Leader(RegionLeaderState::Writable) => { info!("Starting downgrade for region {}", self.region_id); self.set_role(RegionRole::DowngradingLeader); + wait_for_checkpoint = true; } RegionRoleState::Leader(RegionLeaderState::Downgrading) => { // Already in desired state - no-op info!("Region {} already in downgrading mode", self.region_id); + wait_for_checkpoint = true; } _ => { warn!( @@ -588,6 +594,13 @@ impl MitoRegion { } }; + // The state is changed before waiting, so no new writable-leader work + // can race with the barrier. Keep the manager lock while joining to + // serialize the barrier with checkpoint scheduling. + if wait_for_checkpoint { + manager.wait_for_pending_checkpoint().await; + } + // Hack(zhongzc): If we have just become leader (writable), persist any backfilled metadata. let mut backfill_hook_payload: Option = None; if self.state() == RegionRoleState::Leader(RegionLeaderState::Writable) { @@ -1371,10 +1384,14 @@ impl ManifestContext { // Clone before `action_list` is moved into `update` so the hook still // sees what was written. let action_list_for_hook = self.hook.as_ref().map(|_| action_list.clone()); - let version = manager - .update(action_list, is_staging) - .await - .inspect_err(|e| error!(e; "Failed to update manifest, region_id: {}", region_id))?; + let version = if !is_staging + && self.state.load() == RegionRoleState::Leader(RegionLeaderState::Downgrading) + { + manager.update_normal_without_checkpoint(action_list).await + } else { + manager.update(action_list, is_staging).await + } + .inspect_err(|e| error!(e; "Failed to update manifest, region_id: {}", region_id))?; Ok(PendingManifestHook::new( region_id, diff --git a/src/mito2/src/test_util.rs b/src/mito2/src/test_util.rs index 8a63d749b5b..5bf55edb07c 100644 --- a/src/mito2/src/test_util.rs +++ b/src/mito2/src/test_util.rs @@ -52,7 +52,9 @@ use log_store::raft_engine::log_store::RaftEngineLogStore; use log_store::test_util::log_store_util; use moka::future::CacheBuilder; use object_store::ObjectStore; -use object_store::layers::mock::MockLayer; +use object_store::layers::mock::{ + Buffer, Deleter, Metadata, MockLayer, MockLayerBuilder, OpDelete, Result as MockResult, Writer, +}; use object_store::manager::{ObjectStoreManager, ObjectStoreManagerRef}; use object_store::services::Fs; use rskafka::client::partition::{Compression, UnknownTopicHandling}; @@ -67,6 +69,7 @@ use store_api::region_request::{ RegionOpenRequest, RegionPutRequest, RegionRequest, }; use store_api::storage::{ColumnId, RegionId}; +use tokio::sync::Notify; use crate::cache::write_cache::{WriteCache, WriteCacheRef}; use crate::config::MitoConfig; @@ -75,6 +78,7 @@ use crate::engine::{MITO_ENGINE_NAME, MitoEngine}; use crate::error::Result; use crate::flush::{WriteBufferManager, WriteBufferManagerRef}; use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions}; +use crate::manifest::storage::{is_checkpoint_file, is_delta_file}; use crate::read::{Batch, BatchBuilder, BatchReader}; use crate::region::opener::{PartitionExprFetcher, PartitionExprFetcherRef}; use crate::sst::FormatType; @@ -85,6 +89,133 @@ use crate::sst::index::puffin_manager::PuffinManagerFactory; use crate::time_provider::{StdTimeProvider, TimeProviderRef}; use crate::worker::WorkerGroup; +/// Controls a mock object-store layer that blocks a checkpoint task once. +#[derive(Clone)] +pub(crate) struct CheckpointTaskBlocker { + entered: Arc, + release: Arc, + armed: Arc, +} + +impl CheckpointTaskBlocker { + /// Blocks normal manifest checkpoint cleanup at batch-delete close. + pub(crate) fn block_cleanup() -> (Self, MockLayer) { + let blocker = Self { + entered: Arc::new(Notify::new()), + release: Arc::new(Notify::new()), + armed: Arc::new(AtomicBool::new(true)), + }; + let factory_blocker = blocker.clone(); + let layer = MockLayerBuilder::default() + .deleter_factory(Arc::new(move |inner| { + Box::new(BlockingCheckpointDeleter { + inner, + blocker: factory_blocker.clone(), + has_manifest_cleanup_target: false, + }) + })) + .build() + .unwrap(); + (blocker, layer) + } + + /// Blocks publication of `_last_checkpoint` at writer close. + pub(crate) fn block_last_checkpoint_write() -> (Self, MockLayer) { + let blocker = Self { + entered: Arc::new(Notify::new()), + release: Arc::new(Notify::new()), + armed: Arc::new(AtomicBool::new(true)), + }; + let factory_blocker = blocker.clone(); + let layer = MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _args, inner| { + Box::new(BlockingCheckpointWriter { + path: path.to_string(), + inner, + blocker: factory_blocker.clone(), + }) + })) + .build() + .unwrap(); + (blocker, layer) + } + + pub(crate) async fn wait_until_blocked(&self) { + self.entered.notified().await; + } + + pub(crate) fn release(&self) { + self.release.notify_one(); + } + + pub(crate) fn arm_next_close(&self) { + self.armed.store(true, Ordering::Release); + } + + async fn block_once(&self) { + if self.armed.swap(false, Ordering::AcqRel) { + self.entered.notify_one(); + self.release.notified().await; + } + } +} + +struct BlockingCheckpointDeleter { + inner: Deleter, + blocker: CheckpointTaskBlocker, + has_manifest_cleanup_target: bool, +} + +impl object_store::layers::mock::Delete for BlockingCheckpointDeleter { + async fn delete(&mut self, path: &str, args: OpDelete) -> MockResult<()> { + self.inner.delete(path, args).await?; + if is_manifest_checkpoint_file(path) { + self.has_manifest_cleanup_target = true; + } + Ok(()) + } + + async fn close(&mut self) -> MockResult<()> { + if self.has_manifest_cleanup_target { + self.blocker.block_once().await; + } + self.inner.close().await + } +} + +fn is_manifest_checkpoint_file(path: &str) -> bool { + // The mock deleter receives paths relative to the listed manifest + // directory, so the normal/staging directory segments are unavailable. + let path = Path::new(path); + let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else { + return false; + }; + is_delta_file(file_name) || is_checkpoint_file(file_name) +} + +struct BlockingCheckpointWriter { + path: String, + inner: Writer, + blocker: CheckpointTaskBlocker, +} + +impl object_store::layers::mock::Write for BlockingCheckpointWriter { + async fn write(&mut self, bs: Buffer) -> MockResult<()> { + self.inner.write(bs).await + } + + async fn close(&mut self) -> MockResult { + if self.path.ends_with("_last_checkpoint") { + self.blocker.block_once().await; + } + self.inner.close().await + } + + async fn abort(&mut self) -> MockResult<()> { + self.inner.abort().await + } +} + pub(crate) fn new_noop_file_purger() -> FilePurgerRef { Arc::new(NoopFilePurger) } diff --git a/src/mito2/src/worker/handle_enter_staging.rs b/src/mito2/src/worker/handle_enter_staging.rs index 000525021d3..f9e24541241 100644 --- a/src/mito2/src/worker/handle_enter_staging.rs +++ b/src/mito2/src/worker/handle_enter_staging.rs @@ -126,6 +126,10 @@ impl RegionWorkerLoop { // First step: clear all staging manifest files. { let mut manager = region.manifest_ctx.manifest_manager.write().await; + // `set_entering_staging` has already stopped new normal manifest + // publications. Wait for an older checkpoint to finish cleanup + // before exposing staging to repartition remap readers. + manager.wait_for_pending_checkpoint().await; manager .clear_staging_manifest_and_dir() .await From 174577164aaeea2c19a36d214d4eb1cbfd1e6806 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Mon, 24 Aug 2026 09:28:49 +0000 Subject: [PATCH 70/92] feat: otlp duration_nano and trace_flag signed integer coercion (#8816) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(servers): add signed→unsigned int coercion to OTLP ingest path Phase 0 of transitioning built-in data models from unsigned to signed integers (#8793): add lossless Int64→UInt64 and Int32→UInt32 coercion arms so existing UInt64/UInt32 columns (e.g. trace `duration_nano`, log `trace_flags`) keep accepting new signed ingest without an ALTER TABLE. The OTLP ingest path already reconciles every incoming column against the existing table schema and treats it as authoritative. With these arms, `choose_trace_reconcile_decision` returns `UseExisting(UInt64/Uint32)` for an existing unsigned column receiving signed data: the table keeps its type byte-for-byte and the request value is coerced. No persisted format is mutated; existing data stays readable as-is. This is the safety net that makes the actual schema flip (Phase 1) safe. Only the signed→unsigned direction is supported — the reverse would be lossy for values above the signed range and is intentionally rejected. Tests cover both new arms plus an end-to-end log test proving an existing UInt64 column coerces an incoming Int64 value while keeping its type. Signed-off-by: Ning Sun * feat: add compatibility layer for uint trace/log fields Signed-off-by: Ning Sun * fix: jaeger test Signed-off-by: Ning Sun * fix: keep trace v0 unsigned, reject negative span durations - Keep the frozen v0 data model on UInt64 duration_nano: the signed ingest compatibility layer only runs on the v1 path, so flipping v0 would break writes into every pre-existing v0 table at mito's schema check. Pin the schema with unit and integration tests. - Reject spans whose end precedes their start (or whose duration does not fit i64) on the v1 path instead of wrapping: new Int64 tables and existing UInt64 tables now fail identically, rather than storing negative durations that break the Jaeger query API. - Extract is_supported_signed_to_unsigned_coercion so the trace and log ingest paths share one supported-pair predicate and cannot drift. Signed-off-by: Ning Sun * fix: clamp negative span durations to zero, revert semantic_graph comment - Record duration 0 for spans whose end precedes their start instead of erroring: a malformed span no longer fails the request, and the value written is always a non-negative, in-range i64 so new Int64 tables and existing UInt64 tables (via the checked coercion) behave identically. Durations above i64::MAX saturate rather than wrap. - Revert the doc-comment tweak on the semantic_graph test fixture; the file is untouched by this PR again. Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- .../src/instance/otlp/trace_ingest.rs | 23 +- .../src/instance/otlp/trace_ingest/tests.rs | 28 +++ src/frontend/src/instance/otlp/trace_types.rs | 54 ++++ src/servers/src/otlp/coerce.rs | 108 +++++++- src/servers/src/otlp/logs.rs | 231 +++++++++++++++++- src/servers/src/otlp/trace/v0.rs | 27 +- src/servers/src/otlp/trace/v1.rs | 67 ++++- tests-integration/tests/http.rs | 23 +- 8 files changed, 532 insertions(+), 29 deletions(-) diff --git a/src/frontend/src/instance/otlp/trace_ingest.rs b/src/frontend/src/instance/otlp/trace_ingest.rs index 671f6c6f7cf..5951728e6f4 100644 --- a/src/frontend/src/instance/otlp/trace_ingest.rs +++ b/src/frontend/src/instance/otlp/trace_ingest.rs @@ -31,7 +31,7 @@ use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; use pipeline::{GreptimePipelineParams, PipelineWay}; use servers::error::{self, Result as ServerResult}; use servers::otlp; -use servers::otlp::coerce::{coerce_value_data, trace_value_datatype}; +use servers::otlp::coerce::{coerce_value_data, is_supported_trace_coercion, trace_value_datatype}; use servers::otlp::trace::span::{TraceSpan, TraceSpanGroup}; use servers::otlp::trace::v1::{TraceBatchSchema, TraceBinaryType, TraceRetryColumns}; use servers::otlp::trace::{SERVICE_NAME_COLUMN, TraceAuxData}; @@ -2025,10 +2025,23 @@ fn trace_logical_types_incompatible( right_datatype: ColumnDataType, right_concrete_type: &ConcreteDataType, ) -> bool { - left_concrete_type != right_concrete_type - && (left_datatype == right_datatype - || !is_trace_reconcile_candidate_type(left_datatype) - || !is_trace_reconcile_candidate_type(right_datatype)) + if left_concrete_type == right_concrete_type { + return false; + } + // A supported coercion in either direction means the two types can be + // reconciled, so they are not logically incompatible. This lets a signed + // request reach the coercion path instead of being excluded up front when + // the existing column is unsigned (e.g. trace `duration_nano` written as + // Int64 into an existing UInt64 column during the unsigned -> signed + // transition). + if is_supported_trace_coercion(left_datatype, right_datatype) + || is_supported_trace_coercion(right_datatype, left_datatype) + { + return false; + } + left_datatype == right_datatype + || !is_trace_reconcile_candidate_type(left_datatype) + || !is_trace_reconcile_candidate_type(right_datatype) } fn chunk_owned(items: Vec, chunk_size: usize) -> Vec> { diff --git a/src/frontend/src/instance/otlp/trace_ingest/tests.rs b/src/frontend/src/instance/otlp/trace_ingest/tests.rs index ec5bad89754..2445da95bb8 100644 --- a/src/frontend/src/instance/otlp/trace_ingest/tests.rs +++ b/src/frontend/src/instance/otlp/trace_ingest/tests.rs @@ -659,6 +659,34 @@ fn test_trace_request_schema_isolates_non_candidate_batch() { ); } +#[test] +fn test_trace_request_schema_keeps_signed_request_into_existing_unsigned_column() { + // Regression for the unsigned -> signed transition: an existing + // `duration_nano` UInt64 column must NOT exclude a new signed (Int64) + // batch as "logically incompatible". The batch must reach the + // reconciliation path, which coerces Int64 -> UInt64 in place (no ALTER). + // Excluding it up front made ingestion error after the flip. + let column_name = "duration_nano"; + let int64_schema = field_schema(column_name, ColumnDataType::Int64); + let existing_uint64_schema = + DatatypesSchemaBuilder::try_from_columns(vec![DatatypesColumnSchema::new( + column_name, + ConcreteDataType::uint64_datatype(), + true, + )]) + .unwrap() + .build() + .unwrap(); + + let mut request_schema = TraceRequestSchema::default(); + request_schema.observe_trace_column(0, &int64_schema, Some(ColumnDataType::Int64)); + + assert_eq!( + request_schema.incompatible_schema_observations(Some(&existing_uint64_schema)), + HashMap::new() + ); +} + #[test] fn test_trace_request_schema_isolates_non_candidate_new_column_batch() { let column_name = "span_attributes.payload"; diff --git a/src/frontend/src/instance/otlp/trace_types.rs b/src/frontend/src/instance/otlp/trace_types.rs index 603cc66c0d1..0295d12e685 100644 --- a/src/frontend/src/instance/otlp/trace_types.rs +++ b/src/frontend/src/instance/otlp/trace_types.rs @@ -351,6 +351,24 @@ mod tests { ); } + #[test] + fn test_choose_trace_reconcile_decision_existing_uint64_keeps_uint64() { + // Backward-compat for the unsigned -> signed transition: an existing + // table whose `duration_nano` is still UInt64 must keep that type (no + // ALTER) when new signed (Int64) ingest arrives, coercing the value in + // place. This is the no-ALTER guarantee for the trace path; it relies on + // the Int64 -> UInt64 coercion arm added in Phase 0. + assert_eq!( + choose_trace_reconcile_decision( + "duration_nano", + &[ColumnDataType::Int64], + Some(ColumnDataType::Uint64) + ) + .unwrap(), + Some(TraceReconcileDecision::UseExisting(ColumnDataType::Uint64)) + ); + } + #[test] fn test_choose_trace_reconcile_decision_existing_int64_widens_to_float64() { assert_eq!( @@ -612,6 +630,42 @@ mod tests { ); } + #[test] + fn test_prepare_trace_column_rewrites_coerces_int64_into_existing_uint64() { + // Existing-table backward-compat for the trace path: new signed ingest + // arrives as Int64, but the existing `duration_nano` column is UInt64, so + // the rewrite coerces the value into the existing type in place (no + // ALTER). Mirrors what happens for a table created before the + // unsigned -> signed flip. + let mut rows = Rows { + schema: vec![ColumnSchema { + datatype: ColumnDataType::Int64 as i32, + ..Default::default() + }], + rows: vec![Row { + values: vec![Value { + value_data: Some(ValueData::I64Value(42)), + }], + }], + }; + let pending_rewrites = vec![PendingTraceColumnRewrite { + col_idx: 0, + target_type: ColumnDataType::Uint64, + column_name: "duration_nano".to_string(), + }]; + + let prepared = + prepare_trace_column_rewrites(&rows.rows, pending_rewrites, "trace_type_atomicity") + .unwrap(); + + prepared.apply(&mut rows); + assert_eq!(rows.schema[0].datatype, ColumnDataType::Uint64 as i32); + assert_eq!( + rows.rows[0].values[0].value_data, + Some(ValueData::U64Value(42)) + ); + } + #[test] fn test_prepare_trace_column_rewrites_boolean_rejects_invalid_string_parse() { let rows = vec![Row { diff --git a/src/servers/src/otlp/coerce.rs b/src/servers/src/otlp/coerce.rs index 2d87744c5c2..2e7dd36ed7a 100644 --- a/src/servers/src/otlp/coerce.rs +++ b/src/servers/src/otlp/coerce.rs @@ -30,6 +30,37 @@ pub enum TraceCoerceError { // - String to Int64 // - String to Float64 // - String to Boolean +// +// Lossless signed-to-unsigned integer casts. These let the built-in data +// models move from unsigned to signed integers while existing unsigned tables +// keep accepting new signed ingest without an `ALTER TABLE`: an existing +// UInt64/UInt32 column coerces an incoming Int64/Int32 request into the +// existing type: +// - Int64 to UInt64 (e.g. trace `duration_nano` on the v1 path): checked, +// so negative values are rejected rather than silently wrapping. Counts and +// durations are non-negative by construction; a negative request is a +// malformed value (e.g. a span whose end precedes its start). +// - Int32 to UInt32 (e.g. log `trace_flags`): bit-preserving, because +// `trace_flags` is a bit field whose high bits may legitimately be set +// (e.g. the W3C sampled flag); bit patterns must round-trip exactly. + +/// The signed→unsigned integer coercions that let the built-in data models +/// move from unsigned to signed integers while existing unsigned tables keep +/// accepting new signed ingest without an `ALTER TABLE` (see the pair +/// descriptions in the module-level comment above). Kept as one predicate so +/// the trace and log ingest paths share the same supported pair set and +/// cannot drift. +pub fn is_supported_signed_to_unsigned_coercion( + request_type: ColumnDataType, + target_type: ColumnDataType, +) -> bool { + matches!( + (request_type, target_type), + (ColumnDataType::Int64, ColumnDataType::Uint64) + | (ColumnDataType::Int32, ColumnDataType::Uint32) + ) +} + pub fn is_supported_trace_coercion( request_type: ColumnDataType, target_type: ColumnDataType, @@ -43,7 +74,7 @@ pub fn is_supported_trace_coercion( | (ColumnDataType::String, ColumnDataType::Int64) | (ColumnDataType::String, ColumnDataType::Float64) | (ColumnDataType::String, ColumnDataType::Boolean) - ) + ) || is_supported_signed_to_unsigned_coercion(request_type, target_type) } pub fn coerce_value_data( @@ -88,6 +119,18 @@ pub fn coerce_non_null_value( (ColumnDataType::String, ColumnDataType::Boolean, ValueData::StringValue(s)) => { s.parse::().ok().map(ValueData::BoolValue) } + // Checked signed -> unsigned cast for built-in fields moving to signed + // types (durations, counts are always non-negative). Negative values + // are rejected instead of wrapping so the coercion stays lossless. + (ColumnDataType::Int64, ColumnDataType::Uint64, ValueData::I64Value(n)) => { + u64::try_from(*n).ok().map(ValueData::U64Value) + } + // Bit-preserving cast for the `trace_flags` bit field: high bits may + // legitimately be set, and the bit pattern must survive the round-trip + // into an existing UInt32 column unchanged. + (ColumnDataType::Int32, ColumnDataType::Uint32, ValueData::I32Value(n)) => { + Some(ValueData::U32Value(*n as u32)) + } _ => None, } } @@ -240,6 +283,55 @@ mod tests { assert_eq!(result, Err(TraceCoerceError::Unsupported)); } + #[test] + fn test_coerce_int64_to_uint64() { + // Non-negative durations coerce losslessly into an existing UInt64 + // column, so built-in fields moving to signed keep accepting writes. + let result = coerce_value_data( + &Some(ValueData::I64Value(123)), + ColumnDataType::Uint64, + ColumnDataType::Int64, + ); + assert_eq!(result, Ok(Some(ValueData::U64Value(123)))); + } + + #[test] + fn test_coerce_negative_int64_to_uint64_rejected() { + // Negative values must not wrap into the unsigned column: the + // signed -> unsigned coercion is only lossless for non-negative + // inputs, so anything else is rejected instead of silently wrapping + // to a huge number. + let result = coerce_value_data( + &Some(ValueData::I64Value(-1)), + ColumnDataType::Uint64, + ColumnDataType::Int64, + ); + assert_eq!(result, Err(TraceCoerceError::Unsupported)); + } + + #[test] + fn test_coerce_int32_to_uint32() { + let result = coerce_value_data( + &Some(ValueData::I32Value(7)), + ColumnDataType::Uint32, + ColumnDataType::Int32, + ); + assert_eq!(result, Ok(Some(ValueData::U32Value(7)))); + } + + #[test] + fn test_coerce_uint_to_int_not_supported() { + // Only the signed -> unsigned direction is supported (the direction + // the no-ALTER transition needs); the reverse would be lossy for + // values above the signed range and is intentionally rejected. + let result = coerce_value_data( + &Some(ValueData::U64Value(9)), + ColumnDataType::Int64, + ColumnDataType::Uint64, + ); + assert_eq!(result, Err(TraceCoerceError::Unsupported)); + } + #[test] fn test_coerce_none_value() { let result = coerce_value_data(&None, ColumnDataType::Float64, ColumnDataType::Int64); @@ -280,6 +372,20 @@ mod tests { ColumnDataType::Binary, ColumnDataType::Json )); + // Signed -> unsigned casts are supported (built-in no-ALTER transition). + assert!(is_supported_trace_coercion( + ColumnDataType::Int64, + ColumnDataType::Uint64 + )); + assert!(is_supported_trace_coercion( + ColumnDataType::Int32, + ColumnDataType::Uint32 + )); + // The reverse direction is intentionally not supported (lossy). + assert!(!is_supported_trace_coercion( + ColumnDataType::Uint64, + ColumnDataType::Int64 + )); } #[test] diff --git a/src/servers/src/otlp/logs.rs b/src/servers/src/otlp/logs.rs index 33d4f4eb5e5..74f1a3d03a6 100644 --- a/src/servers/src/otlp/logs.rs +++ b/src/servers/src/otlp/logs.rs @@ -44,7 +44,7 @@ use crate::error::{ UnsupportedJsonDataTypeForTagSnafu, }; use crate::http::event::PipelineIngestRequest; -use crate::otlp::coerce::coerce_value_data; +use crate::otlp::coerce::{coerce_value_data, is_supported_signed_to_unsigned_coercion}; use crate::otlp::trace::attributes::OtlpAnyValue; use crate::otlp::utils::{bytes_to_hex_string, key_value_to_jsonb}; use crate::pipeline::run_pipeline; @@ -257,7 +257,7 @@ fn build_otlp_logs_identity_schema() -> Vec { ), ( "trace_flags", - ColumnDataType::Uint32, + ColumnDataType::Int32, SemanticType::Field, None, None, @@ -440,7 +440,7 @@ fn build_otlp_build_in_row( value_data: Some(ValueData::BinaryValue(log_attr.to_vec())), }, GreptimeValue { - value_data: Some(ValueData::U32Value(log.flags)), + value_data: Some(ValueData::I32Value(log.flags as i32)), }, GreptimeValue { value_data: parse_ctx.scope_name.clone().map(ValueData::StringValue), @@ -549,14 +549,17 @@ fn decide_column_schema_and_convert_value( key: column_name, } .fail(), - JsonbNumber::UInt64(u) => Ok(Some(( - GreptimeValue { - value_data: Some(ValueData::U64Value(u)), - }, - ColumnDataType::Uint64, - SemanticType::Tag, - None, - ))), + JsonbNumber::UInt64(u) => { + let value = jsonb_uint64_to_log_value(u, column_name)?; + Ok(Some(( + GreptimeValue { + value_data: Some(ValueData::I64Value(value)), + }, + ColumnDataType::Int64, + SemanticType::Tag, + None, + ))) + } }, JsonbValue::Bool(b) => Ok(Some(( GreptimeValue { @@ -615,6 +618,21 @@ fn decide_existing_column_schema_and_convert_value( ))) } +/// Converts a JSON `UInt64` number into the signed value used by built-in +/// log columns. `JsonbNumber::UInt64` only arises for values that do not fit +/// in `i64` (see `decide_column_schema_and_convert_value`), so the conversion +/// is fallible: wrapping to a negative number would silently corrupt the +/// value (a counter reading back as negative), so out-of-range values are +/// rejected instead. +fn jsonb_uint64_to_log_value(u: u64, column_name: &str) -> Result { + i64::try_from(u).map_err(|_| InvalidParameterSnafu { + reason: format!( + "uint64 value {u} in column '{column_name}' exceeds the i64 range supported by built-in log columns" + ), + } + .build()) +} + fn jsonb_value_to_log_value_data( column_name: &str, value: JsonbValue, @@ -635,7 +653,10 @@ fn jsonb_value_to_log_value_data( key: column_name, } .fail(), - JsonbNumber::UInt64(u) => Ok(Some((ValueData::U64Value(u), ColumnDataType::Uint64))), + JsonbNumber::UInt64(u) => Ok(Some(( + ValueData::I64Value(jsonb_uint64_to_log_value(u, column_name)?), + ColumnDataType::Int64, + ))), }, JsonbValue::Bool(b) => Ok(Some((ValueData::BoolValue(b), ColumnDataType::Boolean))), JsonbValue::Array(_) | JsonbValue::Object(_) => UnsupportedJsonDataTypeForTagSnafu { @@ -718,6 +739,14 @@ fn coerce_log_value_data( return align_timestamp_value(value_data, target_unit, column_name, table_name).map(Some); } + // Lossless signed -> unsigned integer cast for built-in fields moving from + // unsigned to signed types (e.g. an existing UInt64/UInt32 log column + // receiving new Int64/Int32 ingest). Lets existing tables keep their + // unsigned columns without an `ALTER`. In the mutually-exclusive `else` of + // the String branch so the move stays local to non-String targets; the + // shared predicate keeps the supported pairs identical to the trace path + // (the generic String -> numeric coercions keep their pre-existing + // rejection on this path). if target_type == ColumnDataType::String { if let Ok(value_data) = coerce_value_data(&Some(value_data.clone()), target_type, request_type) @@ -727,6 +756,11 @@ fn coerce_log_value_data( if let Some(value_data) = stringify_scalar_value(value_data) { return Ok(Some(value_data)); } + } else if is_supported_signed_to_unsigned_coercion(request_type, target_type) + && let Ok(Some(value_data)) = + coerce_value_data(&Some(value_data), target_type, request_type) + { + return Ok(Some(value_data)); } InvalidParameterSnafu { @@ -1089,6 +1123,16 @@ mod tests { ExistingLogSchema::try_from_schema_parts(&columns, primary_key_indices).unwrap() } + fn existing_uint32_trace_flags_schema() -> ExistingLogSchema { + existing_schema( + vec![ + time_column(ConcreteDataType::timestamp_nanosecond_datatype()), + column("trace_flags", ConcreteDataType::uint32_datatype()), + ], + &[], + ) + } + fn kv(key: &str, value: OtlpValue) -> KeyValue { KeyValue { key: key.to_string(), @@ -1098,12 +1142,20 @@ mod tests { } fn request_with_log_attrs(attrs: Vec) -> ExportLogsServiceRequest { + request_with_log_attrs_and_flags(attrs, 0) + } + + fn request_with_log_attrs_and_flags( + attrs: Vec, + flags: u32, + ) -> ExportLogsServiceRequest { ExportLogsServiceRequest { resource_logs: vec![ResourceLogs { scope_logs: vec![ScopeLogs { log_records: vec![LogRecord { time_unix_nano: 1_234_000_000, trace_id: vec![1; 16], + flags, attributes: attrs, ..Default::default() }], @@ -1151,6 +1203,83 @@ mod tests { ); } + #[test] + fn test_fresh_table_trace_flags_is_int32() { + // Phase 1 of the unsigned -> signed transition: a fresh log table now + // creates `trace_flags` as Int32. Existing tables keep UInt32 (covered + // by the coercion test above); new tables start signed. + let rows = parse_with_select(request_with_log_attrs(vec![]), "", None).unwrap(); + let idx = column_index(&rows, "trace_flags"); + + assert_eq!(rows.schema[idx].datatype, ColumnDataType::Int32 as i32); + assert_eq!( + rows.rows[0].values[idx].value_data, + Some(ValueData::I32Value(0)) + ); + } + + #[test] + fn test_existing_uint32_trace_flags_keeps_type_and_coerces_int32_request() { + // An existing UInt32 `trace_flags` column keeps its type when new + // signed (Int32) ingest arrives, without an `ALTER`. + let existing = existing_uint32_trace_flags_schema(); + + let rows = parse_with_select(request_with_log_attrs(vec![]), "", Some(&existing)).unwrap(); + let idx = column_index(&rows, "trace_flags"); + + assert_eq!(rows.schema[idx].datatype, ColumnDataType::Uint32 as i32); + assert_eq!( + rows.rows[0].values[idx].value_data, + Some(ValueData::U32Value(0)) + ); + } + + #[test] + fn test_existing_uint32_trace_flags_preserves_nonzero_flags_value() { + // A non-zero trace_flags bit pattern (e.g. the W3C sampled flag as sent + // by common encoders) round-trips exactly through the Int32 -> UInt32 + // coercion into an existing unsigned column. Guards that the cast is + // bit-preserving for realistic values, not just the default 0. + let existing = existing_uint32_trace_flags_schema(); + + let rows = parse_with_select( + request_with_log_attrs_and_flags(vec![], 256), + "", + Some(&existing), + ) + .unwrap(); + let idx = column_index(&rows, "trace_flags"); + + assert_eq!(rows.schema[idx].datatype, ColumnDataType::Uint32 as i32); + assert_eq!( + rows.rows[0].values[idx].value_data, + Some(ValueData::U32Value(256)) + ); + } + + #[test] + fn test_existing_uint32_trace_flags_preserves_high_bit_flags_value() { + // The full u32 range round-trips: a value with the sign bit set proves + // the Int32 -> UInt32 coercion is bit-exact (`as i32` then `as u32` + // preserve the pattern), so no flags value is corrupted on the + // unsigned -> signed transition. + let existing = existing_uint32_trace_flags_schema(); + + let rows = parse_with_select( + request_with_log_attrs_and_flags(vec![], 0x8000_0000), + "", + Some(&existing), + ) + .unwrap(); + let idx = column_index(&rows, "trace_flags"); + + assert_eq!(rows.schema[idx].datatype, ColumnDataType::Uint32 as i32); + assert_eq!( + rows.rows[0].values[idx].value_data, + Some(ValueData::U32Value(0x8000_0000)) + ); + } + #[test] fn test_existing_primary_key_updates_builtin_column_semantic_type() { let existing = existing_schema( @@ -1326,4 +1455,82 @@ mod tests { SemanticType::Field as i32 ); } + + #[test] + fn test_existing_uint64_column_keeps_type_and_coerces_int64_request() { + // An existing UInt64 column keeps its type when new signed (Int64) + // ingest arrives, without an `ALTER`. The built-in models move to + // signed while existing unsigned tables stay byte-for-byte unchanged. + let existing = existing_schema( + vec![ + time_column(ConcreteDataType::timestamp_nanosecond_datatype()), + column("counter", ConcreteDataType::uint64_datatype()), + ], + &[], + ); + + let rows = parse_with_select( + request_with_log_attrs(vec![kv("counter", OtlpValue::IntValue(42))]), + "counter", + Some(&existing), + ) + .unwrap(); + let idx = column_index(&rows, "counter"); + + assert_eq!(rows.schema[idx].datatype, ColumnDataType::Uint64 as i32); + assert_eq!( + rows.rows[0].values[idx].value_data, + Some(ValueData::U64Value(42)) + ); + } + + #[test] + fn test_jsonb_uint64_exceeding_i64_range_rejected() { + // A JSON number that only fits in u64 must be rejected, not wrapped: + // storing it as a negative i64 would silently corrupt the value. The + // unsigned -> signed transition is only lossless for values that fit + // the signed range. + let err = jsonb_value_to_log_value_data( + "counter", + JsonbValue::Number(JsonbNumber::UInt64(u64::MAX)), + false, + ) + .unwrap_err(); + + assert!( + err.to_string() + .contains("exceeds the i64 range supported by built-in log columns") + ); + } + + #[test] + fn test_existing_int64_column_rejects_numeric_string_value() { + // Regression guard for the signed -> unsigned transition scope: even a + // parseable numeric string must keep its pre-existing rejection when + // targeting a non-String column. Only the signed -> unsigned integer + // pairs gained coercion on this path, not the generic String -> + // numeric ones. + let existing = existing_schema( + vec![ + time_column(ConcreteDataType::timestamp_nanosecond_datatype()), + column("counter", ConcreteDataType::int64_datatype()), + ], + &[], + ); + + let err = parse_with_select( + request_with_log_attrs(vec![kv( + "counter", + OtlpValue::StringValue("42".to_string()), + )]), + "counter", + Some(&existing), + ) + .unwrap_err(); + + assert!( + err.to_string() + .contains("failed to align log column 'counter'") + ); + } } diff --git a/src/servers/src/otlp/trace/v0.rs b/src/servers/src/otlp/trace/v0.rs index fa10dcc00f8..1bc6752e900 100644 --- a/src/servers/src/otlp/trace/v0.rs +++ b/src/servers/src/otlp/trace/v0.rs @@ -108,6 +108,11 @@ pub fn write_span_to_row(writer: &mut TableData, span: TraceSpan) -> Result<()> span.end_in_nanosecond as i64, )), ), + // The v0 data model is frozen: `duration_nano` stays UInt64. The + // signed→unsigned ingest compatibility layer only runs on the v1 path + // (see `trace_ingest.rs`), so flipping v0 here would break writes into + // every pre-existing v0 table at mito's schema check. New v1 tables + // are signed; v0 tables keep the legacy unsigned column indefinitely. make_column_data( DURATION_NANO_COLUMN, ColumnDataType::Uint64, @@ -231,10 +236,13 @@ fn write_trace_operations_to_row( #[cfg(test)] mod tests { + use api::v1::ColumnDataType; + use api::v1::value::ValueData; + use super::{build_aux_table_requests, build_trace_table_data}; - use crate::otlp::trace::TraceAuxData; use crate::otlp::trace::attributes::Attributes; use crate::otlp::trace::span::{SpanEvents, SpanLinks, TraceSpan}; + use crate::otlp::trace::{DURATION_NANO_COLUMN, TraceAuxData}; fn make_span(service_name: &str, trace_id: &str, span_id: &str) -> TraceSpan { TraceSpan { @@ -271,6 +279,23 @@ mod tests { assert_eq!(rows.len(), 1); } + #[test] + fn test_v0_duration_nano_stays_uint64() { + // The v0 data model is frozen on the unsigned `duration_nano`. This + // pins the schema so the unsigned→signed transition of the built-in + // models (which flips only v1) cannot silently leak into v0 and break + // writes into pre-existing v0 tables. + let writer = build_trace_table_data(&[make_span("svc-a", "trace-a", "span-a")]).unwrap(); + let (schema, rows) = writer.into_schema_and_rows(); + + let idx = schema + .iter() + .position(|c| c.column_name == DURATION_NANO_COLUMN) + .unwrap(); + assert_eq!(schema[idx].datatype, ColumnDataType::Uint64 as i32); + assert_eq!(rows[0].values[idx].value_data, Some(ValueData::U64Value(1))); + } + #[test] fn test_build_aux_table_requests_deduplicates_services_and_operations() { let spans = vec![ diff --git a/src/servers/src/otlp/trace/v1.rs b/src/servers/src/otlp/trace/v1.rs index 8ddeac24f33..f55ab48251e 100644 --- a/src/servers/src/otlp/trace/v1.rs +++ b/src/servers/src/otlp/trace/v1.rs @@ -79,7 +79,7 @@ impl FixedTraceColumnIndexes { Ok(Self { timestamp, timestamp_end: field("timestamp_end", ColumnDataType::TimestampNanosecond)?, - duration_nano: field(DURATION_NANO_COLUMN, ColumnDataType::Uint64)?, + duration_nano: field(DURATION_NANO_COLUMN, ColumnDataType::Int64)?, parent_span_id: field(PARENT_SPAN_ID_COLUMN, ColumnDataType::String)?, trace_id: field(TRACE_ID_COLUMN, ColumnDataType::String)?, span_id: field(SPAN_ID_COLUMN, ColumnDataType::String)?, @@ -342,6 +342,22 @@ pub fn write_span_to_row(writer: &mut TableData, span: TraceSpan) -> Result<()> write_span_to_row_inner(writer, span, row_index, &fixed_columns, None) } +/// Computes the span duration as the signed `duration_nano` value written by +/// the v1 data model. +/// +/// A span whose end precedes its start carries no meaningful duration; it +/// clamps to 0 instead of wrapping — a negative `duration_nano` would fail +/// the checked Int64→UInt64 coercion into pre-existing unsigned tables and +/// would break unsigned readers such as the Jaeger query API. A duration +/// that does not fit `i64` saturates at `i64::MAX`. Clamping at the source +/// keeps new Int64 tables and existing UInt64 tables behaving identically: +/// the written value is always a non-negative, in-range `i64`. +fn span_duration_nano(span: &TraceSpan) -> i64 { + span.end_in_nanosecond + .saturating_sub(span.start_in_nanosecond) + .min(i64::MAX as u64) as i64 +} + /// Writes one span and optionally records its dynamic columns for reconciliation. fn write_span_to_row_inner( writer: &mut TableData, @@ -367,9 +383,7 @@ fn write_span_to_row_inner( ), ( fixed_columns.duration_nano, - Some(ValueData::U64Value( - span.end_in_nanosecond - span.start_in_nanosecond, - )), + Some(ValueData::I64Value(span_duration_nano(&span))), ), ( fixed_columns.parent_span_id, @@ -689,6 +703,49 @@ mod tests { } } + #[test] + fn test_span_end_before_start_records_zero_duration() { + // A span whose end precedes its start carries no meaningful duration; + // it clamps to 0 instead of wrapping. A negative value would fail the + // checked Int64→UInt64 coercion into pre-existing unsigned tables and + // break unsigned readers such as the Jaeger query API. + let mut span = make_span("svc", "trace", "span"); + span.start_in_nanosecond = 200; + span.end_in_nanosecond = 100; + + let (schema, rows) = build_trace_table_data(&[span]) + .unwrap() + .into_schema_and_rows(); + + let idx = schema + .iter() + .position(|c| c.column_name == DURATION_NANO_COLUMN) + .unwrap(); + assert_eq!(rows[0].values[idx].value_data, Some(ValueData::I64Value(0))); + } + + #[test] + fn test_span_duration_above_i64_max_saturates() { + // The companion clamp: a duration that fits u64 but not i64 saturates + // at i64::MAX rather than wrapping to a negative number. + let mut span = make_span("svc", "trace", "span"); + span.start_in_nanosecond = 0; + span.end_in_nanosecond = u64::MAX; + + let (schema, rows) = build_trace_table_data(&[span]) + .unwrap() + .into_schema_and_rows(); + + let idx = schema + .iter() + .position(|c| c.column_name == DURATION_NANO_COLUMN) + .unwrap(); + assert_eq!( + rows[0].values[idx].value_data, + Some(ValueData::I64Value(i64::MAX)) + ); + } + #[test] fn test_fixed_trace_columns_keep_schema_and_values_aligned() { let table_data = build_trace_table_data(&[make_span("svc", "trace", "span")]).unwrap(); @@ -702,7 +759,7 @@ mod tests { "timestamp_end", Some(ValueData::TimestampNanosecondValue(2)), ), - (DURATION_NANO_COLUMN, Some(ValueData::U64Value(1))), + (DURATION_NANO_COLUMN, Some(ValueData::I64Value(1))), (PARENT_SPAN_ID_COLUMN, None), ( TRACE_ID_COLUMN, diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 90cbb8c93fe..f89d75e9ef0 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -6860,6 +6860,19 @@ pub async fn test_otlp_traces_v0(store_type: StorageType) { ) .await; + // The v0 data model is frozen on the unsigned `duration_nano`: pin the + // schema so the unsigned→signed flip of the built-in models (which flips + // only v1) cannot leak into v0 and break writes into pre-existing v0 + // tables. + validate_data( + "otlp_traces_v0_schema", + &client, + "select column_name, lower(data_type) from information_schema.columns \ + where table_name = 'opentelemetry_traces' and column_name = 'duration_nano';", + r#"[["duration_nano","bigint unsigned"]]"#, + ) + .await; + // drop table let res = client .get("/v1/sql?sql=drop table opentelemetry_traces;") @@ -7075,7 +7088,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; - let expected_ddl = r#"[["mytable","CREATE TABLE IF NOT EXISTS \"mytable\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT UNSIGNED NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '1',\n trace_id >= '1' AND trace_id < '2',\n trace_id >= '2' AND trace_id < '3',\n trace_id >= '3' AND trace_id < '4',\n trace_id >= '4' AND trace_id < '5',\n trace_id >= '5' AND trace_id < '6',\n trace_id >= '6' AND trace_id < '7',\n trace_id >= '7' AND trace_id < '8',\n trace_id >= '8' AND trace_id < '9',\n trace_id >= '9' AND trace_id < 'a',\n trace_id >= 'a' AND trace_id < 'b',\n trace_id >= 'b' AND trace_id < 'c',\n trace_id >= 'c' AND trace_id < 'd',\n trace_id >= 'd' AND trace_id < 'e',\n trace_id >= 'e' AND trace_id < 'f',\n trace_id >= 'f'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["mytable","CREATE TABLE IF NOT EXISTS \"mytable\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '1',\n trace_id >= '1' AND trace_id < '2',\n trace_id >= '2' AND trace_id < '3',\n trace_id >= '3' AND trace_id < '4',\n trace_id >= '4' AND trace_id < '5',\n trace_id >= '5' AND trace_id < '6',\n trace_id >= '6' AND trace_id < '7',\n trace_id >= '7' AND trace_id < '8',\n trace_id >= '8' AND trace_id < '9',\n trace_id >= '9' AND trace_id < 'a',\n trace_id >= 'a' AND trace_id < 'b',\n trace_id >= 'b' AND trace_id < 'c',\n trace_id >= 'c' AND trace_id < 'd',\n trace_id >= 'd' AND trace_id < 'e',\n trace_id >= 'e' AND trace_id < 'f',\n trace_id >= 'f'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -7128,7 +7141,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; assert_eq!(StatusCode::OK, res.status()); - let expected_ddl = r#"[["trace_table_part1","CREATE TABLE IF NOT EXISTS \"trace_table_part1\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT UNSIGNED NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\n\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["trace_table_part1","CREATE TABLE IF NOT EXISTS \"trace_table_part1\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\n\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -7165,7 +7178,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; assert_eq!(StatusCode::OK, res.status()); - let expected_ddl = r#"[["trace_table_part4","CREATE TABLE IF NOT EXISTS \"trace_table_part4\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT UNSIGNED NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '4',\n trace_id >= '4' AND trace_id < '8',\n trace_id >= '8' AND trace_id < 'c',\n trace_id >= 'c'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["trace_table_part4","CREATE TABLE IF NOT EXISTS \"trace_table_part4\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '4',\n trace_id >= '4' AND trace_id < '8',\n trace_id >= '8' AND trace_id < 'c',\n trace_id >= 'c'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -7202,7 +7215,7 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { ) .await; assert_eq!(StatusCode::OK, res.status()); - let expected_ddl = r#"[["trace_table_part32","CREATE TABLE IF NOT EXISTS \"trace_table_part32\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT UNSIGNED NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '08',\n trace_id >= '08' AND trace_id < '10',\n trace_id >= '10' AND trace_id < '18',\n trace_id >= '18' AND trace_id < '20',\n trace_id >= '20' AND trace_id < '28',\n trace_id >= '28' AND trace_id < '30',\n trace_id >= '30' AND trace_id < '38',\n trace_id >= '38' AND trace_id < '40',\n trace_id >= '40' AND trace_id < '48',\n trace_id >= '48' AND trace_id < '50',\n trace_id >= '50' AND trace_id < '58',\n trace_id >= '58' AND trace_id < '60',\n trace_id >= '60' AND trace_id < '68',\n trace_id >= '68' AND trace_id < '70',\n trace_id >= '70' AND trace_id < '78',\n trace_id >= '78' AND trace_id < '80',\n trace_id >= '80' AND trace_id < '88',\n trace_id >= '88' AND trace_id < '90',\n trace_id >= '90' AND trace_id < '98',\n trace_id >= '98' AND trace_id < 'a0',\n trace_id >= 'a0' AND trace_id < 'a8',\n trace_id >= 'a8' AND trace_id < 'b0',\n trace_id >= 'b0' AND trace_id < 'b8',\n trace_id >= 'b8' AND trace_id < 'c0',\n trace_id >= 'c0' AND trace_id < 'c8',\n trace_id >= 'c8' AND trace_id < 'd0',\n trace_id >= 'd0' AND trace_id < 'd8',\n trace_id >= 'd8' AND trace_id < 'e0',\n trace_id >= 'e0' AND trace_id < 'e8',\n trace_id >= 'e8' AND trace_id < 'f0',\n trace_id >= 'f0' AND trace_id < 'f8',\n trace_id >= 'f8'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; + let expected_ddl = r#"[["trace_table_part32","CREATE TABLE IF NOT EXISTS \"trace_table_part32\" (\n \"timestamp\" TIMESTAMP(9) NOT NULL,\n \"timestamp_end\" TIMESTAMP(9) NULL,\n \"duration_nano\" BIGINT NULL,\n \"parent_span_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"trace_id\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_id\" STRING NULL,\n \"span_kind\" STRING NULL,\n \"span_name\" STRING NULL,\n \"span_status_code\" STRING NULL,\n \"span_status_message\" STRING NULL,\n \"trace_state\" STRING NULL,\n \"scope_name\" STRING NULL,\n \"scope_version\" STRING NULL,\n \"service_name\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\n \"span_attributes.net.peer.ip\" STRING NULL,\n \"span_attributes.peer.service\" STRING NULL,\n \"span_events\" JSON NULL,\n \"span_links\" JSON NULL,\n TIME INDEX (\"timestamp\"),\n PRIMARY KEY (\"service_name\")\n)\nPARTITION ON COLUMNS (\"trace_id\") (\n trace_id < '08',\n trace_id >= '08' AND trace_id < '10',\n trace_id >= '10' AND trace_id < '18',\n trace_id >= '18' AND trace_id < '20',\n trace_id >= '20' AND trace_id < '28',\n trace_id >= '28' AND trace_id < '30',\n trace_id >= '30' AND trace_id < '38',\n trace_id >= '38' AND trace_id < '40',\n trace_id >= '40' AND trace_id < '48',\n trace_id >= '48' AND trace_id < '50',\n trace_id >= '50' AND trace_id < '58',\n trace_id >= '58' AND trace_id < '60',\n trace_id >= '60' AND trace_id < '68',\n trace_id >= '68' AND trace_id < '70',\n trace_id >= '70' AND trace_id < '78',\n trace_id >= '78' AND trace_id < '80',\n trace_id >= '80' AND trace_id < '88',\n trace_id >= '88' AND trace_id < '90',\n trace_id >= '90' AND trace_id < '98',\n trace_id >= '98' AND trace_id < 'a0',\n trace_id >= 'a0' AND trace_id < 'a8',\n trace_id >= 'a8' AND trace_id < 'b0',\n trace_id >= 'b0' AND trace_id < 'b8',\n trace_id >= 'b8' AND trace_id < 'c0',\n trace_id >= 'c0' AND trace_id < 'c8',\n trace_id >= 'c8' AND trace_id < 'd0',\n trace_id >= 'd0' AND trace_id < 'd8',\n trace_id >= 'd8' AND trace_id < 'e0',\n trace_id >= 'e0' AND trace_id < 'e8',\n trace_id >= 'e8' AND trace_id < 'f0',\n trace_id >= 'f0' AND trace_id < 'f8',\n trace_id >= 'f8'\n)\nENGINE=mito\nWITH(\n 'comment' = 'Created on insertion',\n append_mode = 'true',\n 'greptime.semantic.entity.service.id' = 'service_name',\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\n 'greptime.semantic.signal_type' = 'trace',\n 'greptime.semantic.source' = 'opentelemetry',\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\n table_data_model = 'greptime_trace_v1'\n)"]]"#; validate_data( "otlp_traces", &client, @@ -9601,7 +9614,7 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { .await; assert_eq!(StatusCode::OK, res.status()); - let trace_table_sql = "[[\"mytable\",\"CREATE TABLE IF NOT EXISTS \\\"mytable\\\" (\\n \\\"timestamp\\\" TIMESTAMP(9) NOT NULL,\\n \\\"timestamp_end\\\" TIMESTAMP(9) NULL,\\n \\\"duration_nano\\\" BIGINT UNSIGNED NULL,\\n \\\"parent_span_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"trace_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_id\\\" STRING NULL,\\n \\\"span_kind\\\" STRING NULL,\\n \\\"span_name\\\" STRING NULL,\\n \\\"span_status_code\\\" STRING NULL,\\n \\\"span_status_message\\\" STRING NULL,\\n \\\"trace_state\\\" STRING NULL,\\n \\\"scope_name\\\" STRING NULL,\\n \\\"scope_version\\\" STRING NULL,\\n \\\"service_name\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_attributes.operation.type\\\" STRING NULL,\\n \\\"span_attributes.net.peer.ip\\\" STRING NULL,\\n \\\"span_attributes.peer.service\\\" STRING NULL,\\n \\\"span_events\\\" JSON NULL,\\n \\\"span_links\\\" JSON NULL,\\n TIME INDEX (\\\"timestamp\\\"),\\n PRIMARY KEY (\\\"service_name\\\")\\n)\\nPARTITION ON COLUMNS (\\\"trace_id\\\") (\\n trace_id < '1',\\n trace_id >= '1' AND trace_id < '2',\\n trace_id >= '2' AND trace_id < '3',\\n trace_id >= '3' AND trace_id < '4',\\n trace_id >= '4' AND trace_id < '5',\\n trace_id >= '5' AND trace_id < '6',\\n trace_id >= '6' AND trace_id < '7',\\n trace_id >= '7' AND trace_id < '8',\\n trace_id >= '8' AND trace_id < '9',\\n trace_id >= '9' AND trace_id < 'a',\\n trace_id >= 'a' AND trace_id < 'b',\\n trace_id >= 'b' AND trace_id < 'c',\\n trace_id >= 'c' AND trace_id < 'd',\\n trace_id >= 'd' AND trace_id < 'e',\\n trace_id >= 'e' AND trace_id < 'f',\\n trace_id >= 'f'\\n)\\nENGINE=mito\\nWITH(\\n 'comment' = 'Created on insertion',\\n append_mode = 'true',\\n 'greptime.semantic.entity.service.id' = 'service_name',\\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\\n 'greptime.semantic.signal_type' = 'trace',\\n 'greptime.semantic.source' = 'opentelemetry',\\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\\n table_data_model = 'greptime_trace_v1',\\n ttl = '7days'\\n)\"]]"; + let trace_table_sql = "[[\"mytable\",\"CREATE TABLE IF NOT EXISTS \\\"mytable\\\" (\\n \\\"timestamp\\\" TIMESTAMP(9) NOT NULL,\\n \\\"timestamp_end\\\" TIMESTAMP(9) NULL,\\n \\\"duration_nano\\\" BIGINT NULL,\\n \\\"parent_span_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"trace_id\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_id\\\" STRING NULL,\\n \\\"span_kind\\\" STRING NULL,\\n \\\"span_name\\\" STRING NULL,\\n \\\"span_status_code\\\" STRING NULL,\\n \\\"span_status_message\\\" STRING NULL,\\n \\\"trace_state\\\" STRING NULL,\\n \\\"scope_name\\\" STRING NULL,\\n \\\"scope_version\\\" STRING NULL,\\n \\\"service_name\\\" STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),\\n \\\"span_attributes.operation.type\\\" STRING NULL,\\n \\\"span_attributes.net.peer.ip\\\" STRING NULL,\\n \\\"span_attributes.peer.service\\\" STRING NULL,\\n \\\"span_events\\\" JSON NULL,\\n \\\"span_links\\\" JSON NULL,\\n TIME INDEX (\\\"timestamp\\\"),\\n PRIMARY KEY (\\\"service_name\\\")\\n)\\nPARTITION ON COLUMNS (\\\"trace_id\\\") (\\n trace_id < '1',\\n trace_id >= '1' AND trace_id < '2',\\n trace_id >= '2' AND trace_id < '3',\\n trace_id >= '3' AND trace_id < '4',\\n trace_id >= '4' AND trace_id < '5',\\n trace_id >= '5' AND trace_id < '6',\\n trace_id >= '6' AND trace_id < '7',\\n trace_id >= '7' AND trace_id < '8',\\n trace_id >= '8' AND trace_id < '9',\\n trace_id >= '9' AND trace_id < 'a',\\n trace_id >= 'a' AND trace_id < 'b',\\n trace_id >= 'b' AND trace_id < 'c',\\n trace_id >= 'c' AND trace_id < 'd',\\n trace_id >= 'd' AND trace_id < 'e',\\n trace_id >= 'e' AND trace_id < 'f',\\n trace_id >= 'f'\\n)\\nENGINE=mito\\nWITH(\\n 'comment' = 'Created on insertion',\\n append_mode = 'true',\\n 'greptime.semantic.entity.service.id' = 'service_name',\\n 'greptime.semantic.pipeline' = 'greptime_trace_v1',\\n 'greptime.semantic.signal_type' = 'trace',\\n 'greptime.semantic.source' = 'opentelemetry',\\n 'greptime.semantic.trace.conventions' = 'https://opentelemetry.io/schemas/1.4.0',\\n table_data_model = 'greptime_trace_v1',\\n ttl = '7days'\\n)\"]]"; validate_data( "trace_v1_create_table", &client, From 4ac3423261488fe9d84ed1e2068c79180821072e Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Mon, 24 Aug 2026 10:02:32 +0000 Subject: [PATCH 71/92] fix(mito2): split SSTs at primary key series boundaries (#8888) * fix(mito2): split SSTs at series boundaries Signed-off-by: Lei, HUANG * test(mito2): cover SST splitting without primary key Also document the sortedness precondition and the series boundary split semantics on write_all_flat/write_all_flat_as_primary_key and the new split helpers. Signed-off-by: Lei, HUANG * test(mito2): avoid per-row Vec allocation for empty primary key Signed-off-by: Lei, HUANG * Update src/mito2/src/sst/parquet/writer.rs Co-authored-by: Yingwen --------- Signed-off-by: Lei, HUANG Co-authored-by: Yingwen --- src/mito2/src/sst/parquet.rs | 223 ++++++++++++++++++++++++++- src/mito2/src/sst/parquet/writer.rs | 229 +++++++++++++++++++++------- 2 files changed, 395 insertions(+), 57 deletions(-) diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 5c40bc36bed..e67dcf2bee5 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -152,6 +152,7 @@ mod tests { use crate::sst::index::inverted_index::applier::builder::InvertedIndexApplierBuilder; use crate::sst::index::{IndexBuildType, Indexer, IndexerBuilder, IndexerBuilderImpl}; use crate::sst::parquet::flat_format::FlatWriteFormat; + use crate::sst::parquet::metadata::extract_primary_key_range; use crate::sst::parquet::reader::{ParquetReader, ParquetReaderBuilder, ReaderMetrics}; use crate::sst::parquet::row_selection::RowGroupSelection; use crate::sst::parquet::writer::ParquetWriter; @@ -687,12 +688,13 @@ mod tests { let object_store = env.init_object_store_manager(); let metadata = Arc::new(sst_region_metadata()); let batches = vec![ - new_record_batch_by_range(&["a", "d"], 0, 1000), - new_record_batch_by_range(&["b", "f"], 0, 1000), - new_record_batch_by_range(&["c", "g"], 0, 1000), - new_record_batch_by_range(&["b", "h"], 100, 200), - new_record_batch_by_range(&["b", "h"], 200, 300), - new_record_batch_by_range(&["b", "h"], 300, 1000), + new_record_batch_by_range(&["a", "a"], 0, 1000), + new_record_batch_by_range(&["b", "b"], 0, 1000), + new_record_batch_by_range(&["c", "c"], 0, 1000), + new_record_batch_by_range(&["d", "d"], 100, 200), + new_record_batch_by_range(&["d", "d"], 200, 300), + new_record_batch_by_range(&["d", "d"], 300, 1000), + new_record_batch_by_range(&["e", "e"], 0, 100), ]; let total_rows: usize = batches.iter().map(|batch| batch.num_rows()).sum(); @@ -745,6 +747,159 @@ mod tests { assert_eq!(total_rows, rows_read); } + #[tokio::test] + async fn test_split_file_at_series_boundary_inside_batch() { + let mut env = TestEnv::new().await; + let object_store = env.init_object_store_manager(); + let metadata = Arc::new(sst_region_metadata()); + let first_batch_rows = (0..1000).map(|ts| ("a", "a", ts)).collect::>(); + let second_batch_rows = (0..1000).map(|ts| ("b", "b", ts)).collect::>(); + let mut third_batch_rows = (1000..2000).map(|ts| ("b", "b", ts)).collect::>(); + third_batch_rows.extend((0..1000).map(|ts| ("c", "c", ts))); + let batches = vec![ + new_record_batch_from_rows(&first_batch_rows), + new_record_batch_from_rows(&second_batch_rows), + new_record_batch_from_rows(&third_batch_rows), + ]; + let total_rows = batches.iter().map(RecordBatch::num_rows).sum::(); + let source = new_flat_source_from_record_batches(batches); + let write_opts = WriteOptions { + row_group_size: 50, + max_file_size: Some(1), + ..Default::default() + }; + let path_provider = RegionFilePathFactory { + table_dir: "test_series_boundary".to_string(), + path_type: PathType::Bare, + }; + let mut metrics = Metrics::new(WriteType::Compaction); + let mut writer = ParquetWriter::new_with_object_store( + object_store, + metadata.clone(), + IndexConfig::default(), + NoopIndexBuilder, + path_provider, + &mut metrics, + ) + .await; + + let files = writer + .write_all_flat(source, None, &write_opts) + .await + .unwrap(); + + assert!(files.len() > 1); + assert_eq!( + total_rows, + files.iter().map(|file| file.num_rows).sum::() + ); + let primary_key_ranges = files + .iter() + .map(|file| { + extract_primary_key_range(file.file_metadata.as_ref().unwrap(), metadata.as_ref()) + .unwrap() + }) + .collect::>(); + assert!( + primary_key_ranges + .windows(2) + .all(|ranges| ranges[0].1 < ranges[1].0) + ); + } + + #[tokio::test] + async fn test_oversized_single_series_stays_in_one_file() { + let mut env = TestEnv::new().await; + let object_store = env.init_object_store_manager(); + let metadata = Arc::new(sst_region_metadata()); + let first_batch_rows = (0..1000).map(|ts| ("a", "a", ts)).collect::>(); + let second_batch_rows = (1000..2000).map(|ts| ("a", "a", ts)).collect::>(); + let third_batch_rows = (2000..3000).map(|ts| ("a", "a", ts)).collect::>(); + let batches = vec![ + new_record_batch_from_rows(&first_batch_rows), + new_record_batch_from_rows(&second_batch_rows), + new_record_batch_from_rows(&third_batch_rows), + ]; + let total_rows = batches.iter().map(RecordBatch::num_rows).sum::(); + let source = new_flat_source_from_record_batches(batches); + let write_opts = WriteOptions { + row_group_size: 50, + max_file_size: Some(1), + ..Default::default() + }; + let path_provider = RegionFilePathFactory { + table_dir: "test_oversized_series".to_string(), + path_type: PathType::Bare, + }; + let mut metrics = Metrics::new(WriteType::Compaction); + let mut writer = ParquetWriter::new_with_object_store( + object_store, + metadata, + IndexConfig::default(), + NoopIndexBuilder, + path_provider, + &mut metrics, + ) + .await; + + let files = writer + .write_all_flat_as_primary_key(source, None, &write_opts) + .await + .unwrap(); + + assert_eq!(1, files.len()); + assert_eq!(total_rows, files[0].num_rows); + assert!(files[0].file_size > write_opts.max_file_size.unwrap() as u64); + } + + #[tokio::test] + async fn test_write_multiple_files_without_primary_key() { + let mut env = TestEnv::new().await; + let object_store = env.init_object_store_manager(); + let metadata = Arc::new(sst_region_metadata_without_primary_key()); + let batch_rows = 1000; + let batches = vec![ + new_record_batch_without_primary_key(0, batch_rows), + new_record_batch_without_primary_key(batch_rows, 2 * batch_rows), + new_record_batch_without_primary_key(2 * batch_rows, 3 * batch_rows), + ]; + let total_rows = batches.iter().map(RecordBatch::num_rows).sum::(); + let source = new_flat_source_from_record_batches(batches); + let write_opts = WriteOptions { + row_group_size: 50, + max_file_size: Some(1), + ..Default::default() + }; + let path_provider = RegionFilePathFactory { + table_dir: "test_no_primary_key".to_string(), + path_type: PathType::Bare, + }; + let mut metrics = Metrics::new(WriteType::Compaction); + let mut writer = ParquetWriter::new_with_object_store( + object_store, + metadata, + IndexConfig::default(), + NoopIndexBuilder, + path_provider, + &mut metrics, + ) + .await; + + let files = writer + .write_all_flat(source, None, &write_opts) + .await + .unwrap(); + + // Regions without a primary key keep splitting at batch boundaries: + // the limit produces multiple files and no batch is sliced. + assert!(files.len() > 1); + assert!(files.iter().all(|file| file.num_rows % batch_rows == 0)); + assert_eq!( + total_rows, + files.iter().map(|file| file.num_rows).sum::() + ); + } + #[tokio::test] async fn test_write_read_with_index() { let mut env = TestEnv::new().await; @@ -1145,6 +1300,62 @@ mod tests { assert!(reader.next_record_batch().await.unwrap().is_none()); } + /// Creates a new region metadata without primary key for testing SSTs. + /// + /// Schema: field_0, ts + fn sst_region_metadata_without_primary_key() -> RegionMetadata { + let mut builder = RegionMetadataBuilder::new(REGION_ID); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "field_0".to_string(), + ConcreteDataType::uint64_datatype(), + true, + ), + semantic_type: SemanticType::Field, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 1, + }) + .primary_key(vec![]); + builder.build().unwrap() + } + + /// Creates a flat format RecordBatch for regions without a primary key. + fn new_record_batch_without_primary_key(start: usize, end: usize) -> RecordBatch { + assert!(end >= start); + let metadata = Arc::new(sst_region_metadata_without_primary_key()); + let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default()); + + let num_rows = end - start; + let mut pk_builder = BinaryDictionaryBuilder::::new(); + // Regions without a primary key encode it as empty bytes. + for _ in 0..num_rows { + pk_builder.append([]).unwrap(); + } + + RecordBatch::try_new( + flat_schema, + vec![ + Arc::new(UInt64Array::from_iter_values(start as u64..end as u64)) as ArrayRef, + Arc::new(TimestampMillisecondArray::from_iter_values( + start as i64..end as i64, + )) as ArrayRef, + Arc::new(pk_builder.finish()) as ArrayRef, + Arc::new(UInt64Array::from_value(1000, num_rows)) as ArrayRef, + Arc::new(UInt8Array::from_value(OpType::Put as u8, num_rows)) as ArrayRef, + ], + ) + .unwrap() + } + fn new_record_batch_from_rows(rows: &[(&str, &str, i64)]) -> RecordBatch { let metadata = Arc::new(sst_region_metadata()); let flat_schema = to_flat_sst_arrow_schema(&metadata, &FlatSchemaOptions::default()); diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index a6a4cd5cd06..b9d5016173e 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/src/mito2/src/sst/parquet/writer.rs @@ -23,11 +23,12 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::task::{Context, Poll}; use std::time::Instant; +use bytes::Bytes; use common_telemetry::debug; use common_time::Timestamp; use datatypes::arrow::array::{ - ArrayRef, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, - TimestampSecondArray, + ArrayRef, BinaryArray, TimestampMicrosecondArray, TimestampMillisecondArray, + TimestampNanosecondArray, TimestampSecondArray, UInt32Array, }; use datatypes::arrow::compute::{max, min}; use datatypes::arrow::datatypes::{DataType, SchemaRef, TimeUnit}; @@ -40,7 +41,7 @@ use parquet::file::metadata::KeyValue; use parquet::file::properties::{WriterProperties, WriterPropertiesBuilder}; use parquet::schema::types::ColumnPath; use smallvec::smallvec; -use snafu::ResultExt; +use snafu::{OptionExt, ResultExt}; use store_api::metadata::RegionMetadataRef; use store_api::storage::consts::{OP_TYPE_COLUMN_NAME, SEQUENCE_COLUMN_NAME}; use store_api::storage::{FileId, SequenceNumber}; @@ -50,13 +51,16 @@ use tokio_util::compat::{Compat, FuturesAsyncWriteCompatExt}; use crate::access_layer::{FilePathProvider, Metrics, SstInfoArray, TempFileCleaner}; use crate::config::{IndexBuildMode, IndexConfig}; use crate::error::{ - InvalidMetadataSnafu, OpenDalSnafu, Result, UnexpectedSnafu, WriteParquetSnafu, + InvalidMetadataSnafu, InvalidRecordBatchSnafu, OpenDalSnafu, Result, UnexpectedSnafu, + WriteParquetSnafu, }; use crate::read::FlatSource; use crate::sst::file::RegionFileId; use crate::sst::index::{IndexOutput, Indexer, IndexerBuilder}; -use crate::sst::parquet::flat_format::{FlatWriteFormat, time_index_column_index}; -use crate::sst::parquet::format::PrimaryKeyWriteFormat; +use crate::sst::parquet::flat_format::{ + FlatWriteFormat, primary_key_column_index, time_index_column_index, +}; +use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyWriteFormat}; use crate::sst::parquet::{PARQUET_METADATA_KEY, SstInfo, WriteOptions}; use crate::sst::{ DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator, @@ -85,6 +89,71 @@ impl FlatBatchConverter { } } +/// Result of splitting a batch at the next series boundary. +enum SeriesBoundarySplit { + /// The whole batch belongs to the current series and stays in the current file. + Continue(RecordBatch), + /// The batch crosses a series boundary. + Split { + /// Remaining rows of the current series, appended to the current file. + current_file_tail: Option, + /// Rows of the following series that start the next file. + next_file_head: RecordBatch, + }, +} + +/// Returns the dictionary keys and values of the encoded `__primary_key` column. +fn encoded_primary_keys(batch: &RecordBatch) -> Result<(&UInt32Array, &BinaryArray)> { + let column = batch.column(primary_key_column_index(batch.num_columns())); + let primary_keys = column + .as_any() + .downcast_ref::() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!( + "expected dictionary primary key column, got {:?}", + column.data_type() + ), + })?; + let values = primary_keys + .values() + .as_any() + .downcast_ref::() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!( + "expected binary primary key values, got {:?}", + primary_keys.values().data_type() + ), + })?; + Ok((primary_keys.keys(), values)) +} + +/// Splits `batch` at the first row whose encoded primary key differs from +/// `current_primary_key`, the last primary key written to the current file. +/// +/// The writer finishes an oversized file at such a series boundary so that the +/// primary key ranges of output files never overlap, which allows pickers like +/// TWCS to detect overlapping files by their time and primary key ranges. A +/// series is never split in the middle: if no row starts a new series, the whole +/// batch stays in the current file even if it already exceeds the size limit. +fn split_at_next_series( + batch: RecordBatch, + current_primary_key: &[u8], +) -> Result { + let (keys, values) = encoded_primary_keys(&batch)?; + let Some(offset) = (0..batch.num_rows()) + .find(|&row| values.value(keys.value(row) as usize) != current_primary_key) + else { + return Ok(SeriesBoundarySplit::Continue(batch)); + }; + + let next_file_head = batch.slice(offset, batch.num_rows() - offset); + let current_file_tail = (offset > 0).then(|| batch.slice(0, offset)); + Ok(SeriesBoundarySplit::Split { + current_file_tail, + next_file_head, + }) +} + /// Parquet SST writer. pub struct ParquetWriter<'a, F: WriterFactory, I: IndexerBuilder, P: FilePathProvider> { /// Path provider that creates SST and index file paths according to file id. @@ -269,6 +338,13 @@ where /// Iterates FlatSource and writes all RecordBatch in flat format to Parquet file. /// + /// The source must yield batches globally sorted by the encoded primary key. + /// `opts.max_file_size` is a soft limit for regions with a primary key: the + /// current file is finished at the next series boundary after exceeding the + /// limit (see [split_at_next_series]), so a series larger than the limit + /// stays in one file. Regions without a primary key are split at batch + /// boundaries. + /// /// Returns the [SstInfo] if the SST is written. pub async fn write_all_flat( &mut self, @@ -304,6 +380,9 @@ where /// Iterates FlatSource and writes all RecordBatch in primary-key format to Parquet file. /// + /// The source must yield batches globally sorted by the encoded primary key. + /// See [Self::write_all_flat] for the `opts.max_file_size` semantics. + /// /// Returns the [SstInfo] if the SST is written. pub async fn write_all_flat_as_primary_key( &mut self, @@ -336,36 +415,55 @@ where let mut results = smallvec![]; let mut stats = SourceStats::default(); - while let Some(record_batch) = self - .write_next_flat_batch(&mut source, converter, opts) - .await - .transpose() - { - match record_batch { - Ok(batch) => { - stats.update_flat(&batch)?; - if matches!(self.index_config.build_mode, IndexBuildMode::Sync) { - let start = Instant::now(); - // safety: self.current_indexer must be set when first batch has been written. - self.current_indexer - .as_mut() - .unwrap() - .update_flat(&batch) - .await; - self.metrics.update_index += start.elapsed(); - } - if let Some(max_file_size) = opts.max_file_size - && self.bytes_written.load(Ordering::Relaxed) > max_file_size - { - self.finish_current_file(&mut results, &mut stats).await?; - } - } + loop { + let start = Instant::now(); + let batch = match source.next_batch().await { + Ok(Some(batch)) => batch, + Ok(None) => break, Err(e) => { - if let Some(indexer) = &mut self.current_indexer { - indexer.abort().await; - } + self.abort_current_indexer().await; return Err(e); } + }; + self.metrics.iter_source += start.elapsed(); + + if self.metadata.primary_key.is_empty() { + self.append_flat_batch(&batch, converter, opts, &mut stats) + .await?; + if self.exceeds_max_file_size(opts) { + self.finish_current_file(&mut results, &mut stats).await?; + } + } else if self.exceeds_max_file_size(opts) + && let Some(current_primary_key) = stats.last_primary_key.as_deref() + { + let series_split = match split_at_next_series(batch, current_primary_key) { + Ok(series_split) => series_split, + Err(e) => { + self.abort_current_indexer().await; + return Err(e); + } + }; + match series_split { + SeriesBoundarySplit::Continue(batch) => { + self.append_flat_batch(&batch, converter, opts, &mut stats) + .await?; + } + SeriesBoundarySplit::Split { + current_file_tail, + next_file_head, + } => { + if let Some(tail) = current_file_tail { + self.append_flat_batch(&tail, converter, opts, &mut stats) + .await?; + } + self.finish_current_file(&mut results, &mut stats).await?; + self.append_flat_batch(&next_file_head, converter, opts, &mut stats) + .await?; + } + } + } else { + self.append_flat_batch(&batch, converter, opts, &mut stats) + .await?; } } @@ -398,29 +496,53 @@ where .set_column_compression(op_type_col, Compression::UNCOMPRESSED) } - async fn write_next_flat_batch( + async fn append_flat_batch( &mut self, - source: &mut FlatSource, + batch: &RecordBatch, converter: &FlatBatchConverter, opts: &WriteOptions, - ) -> Result> { - let start = Instant::now(); - let Some(record_batch) = source.next_batch().await? else { - return Ok(None); - }; - self.metrics.iter_source += start.elapsed(); + stats: &mut SourceStats, + ) -> Result<()> { + let result = async { + let arrow_batch = converter.convert_batch(batch)?; + let start = Instant::now(); + self.maybe_init_writer(arrow_batch.schema_ref(), opts) + .await? + .write(&arrow_batch) + .await + .context(WriteParquetSnafu)?; + self.metrics.write_batch += start.elapsed(); - let arrow_batch = converter.convert_batch(&record_batch)?; + stats.update_flat(batch)?; + if matches!(self.index_config.build_mode, IndexBuildMode::Sync) { + let start = Instant::now(); + // safety: self.current_indexer must be set when first batch has been written. + self.current_indexer + .as_mut() + .unwrap() + .update_flat(batch) + .await; + self.metrics.update_index += start.elapsed(); + } + Ok(()) + } + .await; - let start = Instant::now(); - self.maybe_init_writer(arrow_batch.schema_ref(), opts) - .await? - .write(&arrow_batch) - .await - .context(WriteParquetSnafu)?; - self.metrics.write_batch += start.elapsed(); - // Return original flat batch for stats/indexer which use flat layout. - Ok(Some(record_batch)) + if result.is_err() { + self.abort_current_indexer().await; + } + result + } + + fn exceeds_max_file_size(&self, opts: &WriteOptions) -> bool { + opts.max_file_size + .is_some_and(|max_size| self.bytes_written.load(Ordering::Relaxed) >= max_size) + } + + async fn abort_current_indexer(&mut self) { + if let Some(indexer) = &mut self.current_indexer { + indexer.abort().await; + } } async fn maybe_init_writer( @@ -481,6 +603,8 @@ struct SourceStats { num_rows: usize, /// Time range of fetched batches. time_range: Option<(Timestamp, Timestamp)>, + /// Last primary key written to the current file. + last_primary_key: Option, /// Series estimator for computing num_series. series_estimator: SeriesEstimator, } @@ -493,6 +617,9 @@ impl SourceStats { self.num_rows += record_batch.num_rows(); self.series_estimator.update_flat(record_batch); + let (keys, values) = encoded_primary_keys(record_batch)?; + let key = keys.value(record_batch.num_rows() - 1); + self.last_primary_key = Some(Bytes::copy_from_slice(values.value(key as usize))); // Get the timestamp column by index let time_index_col_idx = time_index_column_index(record_batch.num_columns()); From 82444635f501113cccb56a2bf49372073120e456 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Mon, 24 Aug 2026 12:45:35 +0000 Subject: [PATCH 72/92] feat: support quantile and fraction queries on mixed histograms (#8874) * fix: test Signed-off-by: shuiyisong * fix: remove try_build_float_literal Signed-off-by: shuiyisong * fix: remove NonCommutative Signed-off-by: shuiyisong * fix: NaN Signed-off-by: shuiyisong * chore: add comments Signed-off-by: shuiyisong * test(query): cover frontend-only histogram fold planning Signed-off-by: shuiyisong * fix(promql): ignore unparseable histogram bucket labels Signed-off-by: shuiyisong * fix: typo Signed-off-by: shuiyisong * fix: sqlness Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/common/query/src/native_histogram.rs | 123 ++- .../substrait/src/extension_serializer.rs | 2 +- src/promql/src/extension_plan.rs | 4 +- .../src/extension_plan/histogram_fold.rs | 769 ++++++++++++++---- src/promql/src/functions/native_histogram.rs | 102 ++- src/query/src/dist_plan/analyzer/test.rs | 71 +- src/query/src/dist_plan/commutativity.rs | 26 +- src/query/src/promql/planner.rs | 709 +++++++++++++--- .../distributed/flow-tql/tsid_on_phy.result | 2 +- .../promql/histogram_multi_partition.result | 2 +- .../tsid_histogram_quantile_regression.result | 2 +- 11 files changed, 1532 insertions(+), 280 deletions(-) diff --git a/src/common/query/src/native_histogram.rs b/src/common/query/src/native_histogram.rs index b090ac33fcd..1c7bb09868e 100644 --- a/src/common/query/src/native_histogram.rs +++ b/src/common/query/src/native_histogram.rs @@ -215,6 +215,15 @@ struct Bucket { boundary_rule: BoundaryRule, } +/// Additional information produced while estimating a native histogram quantile. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum NativeHistogramQuantileInfo { + /// NaN observations made the estimated quantile skew higher. + NaNSkew, + /// The requested quantile fell beyond all populated buckets because of NaN observations. + NaNResult, +} + /// Query-time representation of a Prometheus native histogram. /// /// Bucket counts are absolute `f64` values even when the persisted payload used @@ -635,19 +644,27 @@ impl NativeHistogram { /// Returns negative or positive infinity outside `[0, 1]`, and `NaN` when /// the quantile cannot be estimated. pub fn quantile(&self, q: f64) -> f64 { + self.quantile_with_info(q).0 + } + + /// Estimates the value at quantile `q` and reports the effect of NaN observations. + pub fn quantile_with_info(&self, q: f64) -> (f64, Option) { if q < 0.0 { - return f64::NEG_INFINITY; + return (f64::NEG_INFINITY, None); } if q > 1.0 { - return f64::INFINITY; + return (f64::INFINITY, None); } if self.count == 0.0 || q.is_nan() { - return f64::NAN; + return (f64::NAN, None); } let Some(mut buckets) = self.all_buckets() else { - return f64::NAN; + return (f64::NAN, None); }; + let bucket_total = buckets.iter().map(|bucket| bucket.count).sum::(); + let has_nan_observations = self.sum.is_nan() && bucket_total < self.count; + let info = has_nan_observations.then_some(NativeHistogramQuantileInfo::NaNSkew); let rank = q * self.count; let mut count = 0.0; for bucket in &mut buckets { @@ -668,42 +685,65 @@ impl NativeHistogram { } else if self.uses_custom_buckets() { if bucket.lower == f64::NEG_INFINITY { if bucket.upper <= 0.0 { - return bucket.upper; + return (bucket.upper, info); } bucket.lower = 0.0; } else if bucket.upper == f64::INFINITY { - return bucket.lower; + return (bucket.lower, info); } } let rank_in_bucket = rank - (count - bucket.count); let fraction = rank_in_bucket / bucket.count; if self.uses_custom_buckets() || (bucket.lower <= 0.0 && bucket.upper >= 0.0) { - return bucket.lower + (bucket.upper - bucket.lower) * fraction; + return ( + bucket.lower + (bucket.upper - bucket.lower) * fraction, + info, + ); } let log_lower = bucket.lower.abs().log2(); let log_upper = bucket.upper.abs().log2(); if bucket.lower > 0.0 { - return 2.0_f64.powf(log_lower + (log_upper - log_lower) * fraction); + return ( + 2.0_f64.powf(log_lower + (log_upper - log_lower) * fraction), + info, + ); } - return -2.0_f64.powf(log_upper + (log_lower - log_upper) * (1.0 - fraction)); + return ( + -2.0_f64.powf(log_upper + (log_lower - log_upper) * (1.0 - fraction)), + info, + ); } - f64::NAN + let info = if self.sum.is_nan() && count < self.count { + Some(if count < rank { + NativeHistogramQuantileInfo::NaNResult + } else { + NativeHistogramQuantileInfo::NaNSkew + }) + } else { + None + }; + (f64::NAN, info) } /// Estimates the fraction of observations between `lower` and `upper`. pub fn fraction(&self, lower: f64, upper: f64) -> f64 { + self.fraction_with_info(lower, upper).0 + } + + /// Estimates a fraction and reports whether NaN observations were excluded. + pub fn fraction_with_info(&self, lower: f64, upper: f64) -> (f64, bool) { if self.count == 0.0 || lower.is_nan() || upper.is_nan() { - return f64::NAN; + return (f64::NAN, false); } if lower >= upper { - return 0.0; + return (0.0, false); } let Some(mut buckets) = self.all_buckets() else { - return f64::NAN; + return (f64::NAN, false); }; let count = if self.sum.is_nan() { buckets.iter().map(|bucket| bucket.count).sum() @@ -759,7 +799,10 @@ impl NativeHistogram { upper_rank = count; } - (upper_rank - lower_rank) / self.count + ( + (upper_rank - lower_rank) / self.count, + self.sum.is_nan() && count < self.count, + ) } /// Converts populated buckets to `(boundary rule, lower, upper, count)` strings. @@ -2141,7 +2184,57 @@ mod tests { histogram.count = 10.0; histogram.sum = f64::NAN; - assert_eq!(histogram.fraction(f64::NEG_INFINITY, f64::INFINITY), 0.8); + assert_eq!( + histogram.fraction_with_info(f64::NEG_INFINITY, f64::INFINITY), + (0.8, true) + ); + } + + #[test] + fn quantile_reports_nan_observation_effect() { + let mut histogram = histogram( + vec![Span { + offset: 0, + length: 1, + }], + vec![8.0], + ); + histogram.count = 10.0; + histogram.sum = f64::NAN; + + let (skewed, info) = histogram.quantile_with_info(0.5); + assert!(skewed.is_finite()); + assert_eq!(info, Some(NativeHistogramQuantileInfo::NaNSkew)); + + let (nan, info) = histogram.quantile_with_info(0.9); + assert!(nan.is_nan()); + assert_eq!(info, Some(NativeHistogramQuantileInfo::NaNResult)); + + histogram.count = 8.0; + assert_eq!(histogram.quantile_with_info(0.5).1, None); + + histogram.count = 3.0; + histogram.positive_spans.clear(); + histogram.positive_buckets.clear(); + let (nan, info) = histogram.quantile_with_info(0.0); + assert!(nan.is_nan()); + assert_eq!(info, Some(NativeHistogramQuantileInfo::NaNSkew)); + + histogram.schema = CUSTOM_BUCKETS_SCHEMA; + histogram.count = 10.0; + histogram.positive_spans = vec![Span { + offset: 0, + length: 1, + }]; + histogram.positive_buckets = vec![8.0]; + for (bound, offset, expected) in [(-1.0, 0, -1.0), (1.0, 1, 1.0)] { + histogram.custom_values = vec![bound]; + histogram.positive_spans[0].offset = offset; + assert_eq!( + histogram.quantile_with_info(0.5), + (expected, Some(NativeHistogramQuantileInfo::NaNSkew)) + ); + } } #[test] diff --git a/src/common/substrait/src/extension_serializer.rs b/src/common/substrait/src/extension_serializer.rs index e2e7da524c1..5e94a1a57ae 100644 --- a/src/common/substrait/src/extension_serializer.rs +++ b/src/common/substrait/src/extension_serializer.rs @@ -78,7 +78,7 @@ impl SerializerRegistry for ExtensionSerializer { .as_any() .downcast_ref::() .expect("Failed to downcast to HistogramFold"); - Ok(histogram_fold.serialize()) + histogram_fold.serialize() } name if name == UnionDistinctOn::name() => { let union_distinct_on = node diff --git a/src/promql/src/extension_plan.rs b/src/promql/src/extension_plan.rs index 973758032c2..29b31b7ca03 100644 --- a/src/promql/src/extension_plan.rs +++ b/src/promql/src/extension_plan.rs @@ -34,7 +34,9 @@ use datafusion::common::DFSchemaRef; use datafusion::error::{DataFusionError, Result as DataFusionResult}; use datatypes::data_type::DataType as _; pub use empty_metric::{EmptyMetric, EmptyMetricExec, EmptyMetricStream, build_special_time_expr}; -pub use histogram_fold::{HistogramFold, HistogramFoldExec, HistogramFoldStream}; +pub use histogram_fold::{ + HistogramFold, HistogramFoldExec, HistogramFoldOperation, HistogramFoldStream, +}; pub use instant_manipulate::{InstantManipulate, InstantManipulateExec, InstantManipulateStream}; pub use normalize::{SeriesNormalize, SeriesNormalizeExec, SeriesNormalizeStream}; pub use planner::PromExtensionPlanner; diff --git a/src/promql/src/extension_plan/histogram_fold.rs b/src/promql/src/extension_plan/histogram_fold.rs index 462f577fc67..5a5f2f8d466 100644 --- a/src/promql/src/extension_plan/histogram_fold.rs +++ b/src/promql/src/extension_plan/histogram_fold.rs @@ -33,7 +33,7 @@ use datafusion::physical_expr::{ EquivalenceProperties, LexRequirement, OrderingRequirements, PhysicalSortRequirement, }; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; -use datafusion::physical_plan::expressions::{CastExpr as PhyCast, Column as PhyColumn}; +use datafusion::physical_plan::expressions::{Column as PhyColumn, TryCastExpr as PhyTryCast}; use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, @@ -78,11 +78,42 @@ pub struct HistogramFold { ts_column: String, input: LogicalPlan, field_column: String, - quantile: OrderedF64, + /// Native histogram companion column for a mixed classic/native input. + histogram_column: Option, + operation: HistogramFoldOperation, output_schema: DFSchemaRef, unfix: Option, } +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd)] +pub enum HistogramFoldOperation { + Quantile(OrderedF64), + Fraction { + lower: OrderedF64, + upper: OrderedF64, + }, +} + +impl HistogramFoldOperation { + pub const fn function_name(self) -> &'static str { + match self { + Self::Quantile(_) => "histogram_quantile", + Self::Fraction { .. } => "histogram_fraction", + } + } +} + +impl std::fmt::Display for HistogramFoldOperation { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Quantile(quantile) => write!(f, "quantile={quantile}"), + Self::Fraction { lower, upper } => { + write!(f, "fraction=[{lower}, {upper}]") + } + } + } +} + #[derive(Debug, PartialEq, Eq, Hash, PartialOrd)] struct UnfixIndices { pub le_column_idx: u64, @@ -137,6 +168,21 @@ impl UserDefinedLogicalNodeCore for HistogramFold { return Some(vec![indices]); } + if let Some(histogram_column) = &self.histogram_column { + let mut necessary_indices = output_columns.to_vec(); + for column in [ + &self.le_column, + &self.ts_column, + &self.field_column, + histogram_column, + ] { + necessary_indices.push(input_schema.index_of_column_by_name(None, column)?); + } + necessary_indices.sort_unstable(); + necessary_indices.dedup(); + return Some(vec![necessary_indices]); + } + let mut necessary_indices = output_columns .iter() .map(|&output_column| { @@ -156,9 +202,13 @@ impl UserDefinedLogicalNodeCore for HistogramFold { fn fmt_for_explain(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!( f, - "HistogramFold: le={}, field={}, quantile={}", - self.le_column, self.field_column, self.quantile - ) + "HistogramFold: le={}, field={}", + self.le_column, self.field_column + )?; + if let Some(histogram) = &self.histogram_column { + write!(f, ", histogram={histogram}")?; + } + write!(f, ", {}", self.operation) } fn with_exprs_and_inputs( @@ -194,7 +244,8 @@ impl UserDefinedLogicalNodeCore for HistogramFold { ts_column, input, field_column, - quantile: self.quantile, + histogram_column: None, + operation: self.operation, output_schema, unfix: None, }) @@ -204,7 +255,8 @@ impl UserDefinedLogicalNodeCore for HistogramFold { ts_column: self.ts_column.clone(), input, field_column: self.field_column.clone(), - quantile: self.quantile, + histogram_column: self.histogram_column.clone(), + operation: self.operation, output_schema: self.output_schema.clone(), unfix: None, }) @@ -219,16 +271,42 @@ impl HistogramFold { ts_column: String, quantile: f64, input: LogicalPlan, + ) -> DataFusionResult { + Self::new_with_operation( + le_column, + field_column, + ts_column, + HistogramFoldOperation::Quantile(quantile.into()), + None, + input, + ) + } + + pub fn new_with_operation( + le_column: String, + field_column: String, + ts_column: String, + operation: HistogramFoldOperation, + histogram_column: Option, + input: LogicalPlan, ) -> DataFusionResult { let input_schema = input.schema(); Self::check_schema(input_schema, &le_column, &field_column, &ts_column)?; - let output_schema = Self::convert_schema(input_schema, &le_column)?; + if let Some(histogram_column) = &histogram_column { + Self::check_column(input_schema, histogram_column)?; + } + let output_schema = if histogram_column.is_some() { + input_schema.clone() + } else { + Self::convert_schema(input_schema, &le_column)? + }; Ok(Self { le_column, ts_column, input, field_column, - quantile: quantile.into(), + histogram_column, + operation, output_schema, unfix: None, }) @@ -244,23 +322,22 @@ impl HistogramFold { field_column: &str, ts_column: &str, ) -> DataFusionResult<()> { - let check_column = |col| { - if !input_schema.has_column_with_unqualified_name(col) { - Err(DataFusionError::SchemaError( - Box::new(datafusion::common::SchemaError::FieldNotFound { - field: Box::new(Column::new(None::, col)), - valid_fields: input_schema.columns(), - }), - Box::new(None), - )) - } else { - Ok(()) - } - }; + Self::check_column(input_schema, le_column)?; + Self::check_column(input_schema, ts_column)?; + Self::check_column(input_schema, field_column) + } - check_column(le_column)?; - check_column(ts_column)?; - check_column(field_column) + fn check_column(input_schema: &DFSchemaRef, column: &str) -> DataFusionResult<()> { + if !input_schema.has_column_with_unqualified_name(column) { + return Err(DataFusionError::SchemaError( + Box::new(datafusion::common::SchemaError::FieldNotFound { + field: Box::new(Column::new(None::, column)), + valid_fields: input_schema.columns(), + }), + Box::new(None), + )); + } + Ok(()) } pub fn to_execution_plan(&self, exec_input: Arc) -> Arc { @@ -272,6 +349,10 @@ impl HistogramFold { let field_column_index = input_schema .index_of_column_by_name(None, &self.field_column) .unwrap(); + let histogram_column_index = self + .histogram_column + .as_ref() + .map(|column| input_schema.index_of_column_by_name(None, column).unwrap()); let ts_column_index = input_schema .index_of_column_by_name(None, &self.ts_column) .unwrap(); @@ -282,7 +363,11 @@ impl HistogramFold { .iter() .enumerate() .filter_map(|(idx, field)| { - if idx == le_column_index || idx == field_column_index || idx == ts_column_index { + if idx == le_column_index + || idx == field_column_index + || Some(idx) == histogram_column_index + || idx == ts_column_index + { None } else { Some(Arc::new(PhyColumn::new(field.name(), idx)) as _) @@ -309,11 +394,12 @@ impl HistogramFold { Arc::new(HistogramFoldExec { le_column_index, field_column_index, + histogram_column_index, ts_column_index, input: exec_input, tag_columns, partition_exprs, - quantile: self.quantile.into(), + operation: self.operation, output_schema, metric: ExecutionPlanMetricsSet::new(), properties, @@ -343,18 +429,28 @@ impl HistogramFold { )?)) } - pub fn serialize(&self) -> Vec { + pub fn serialize(&self) -> DataFusionResult> { + if self.histogram_column.is_some() { + return Err(DataFusionError::NotImplemented( + "mixed HistogramFold is frontend-only".to_string(), + )); + } + let HistogramFoldOperation::Quantile(quantile) = self.operation else { + return Err(DataFusionError::NotImplemented( + "HistogramFold fraction is frontend-only".to_string(), + )); + }; let le_column_idx = serialize_column_index(self.input.schema(), &self.le_column); let ts_column_idx = serialize_column_index(self.input.schema(), &self.ts_column); let field_column_idx = serialize_column_index(self.input.schema(), &self.field_column); - pb::HistogramFold { + Ok(pb::HistogramFold { le_column_idx, ts_column_idx, field_column_idx, - quantile: self.quantile.into(), + quantile: quantile.into(), } - .encode_to_vec() + .encode_to_vec()) } pub fn deserialize(bytes: &[u8]) -> Result { @@ -375,7 +471,8 @@ impl HistogramFold { ts_column: String::new(), input: placeholder_plan, field_column: String::new(), - quantile: pb_histogram_fold.quantile.into(), + histogram_column: None, + operation: HistogramFoldOperation::Quantile(pb_histogram_fold.quantile.into()), output_schema: Arc::new(DFSchema::empty()), unfix: Some(unfix), }) @@ -401,7 +498,11 @@ impl PartialOrd for HistogramFold { Some(core::cmp::Ordering::Equal) => {} ord => return ord, } - self.quantile.partial_cmp(&other.quantile) + match self.histogram_column.partial_cmp(&other.histogram_column) { + Some(core::cmp::Ordering::Equal) => {} + ord => return ord, + } + self.operation.partial_cmp(&other.operation) } } @@ -413,11 +514,13 @@ pub struct HistogramFoldExec { output_schema: SchemaRef, /// Index for field column in the schema of input. field_column_index: usize, + /// Index for the native histogram companion column on mixed inputs. + histogram_column_index: Option, ts_column_index: usize, /// Tag columns are all columns except `le`, `field` and `ts` columns. tag_columns: Vec>, partition_exprs: Vec>, - quantile: f64, + operation: HistogramFoldOperation, metric: ExecutionPlanMetricsSet, properties: Arc, } @@ -448,21 +551,22 @@ impl ExecutionPlan for HistogramFoldExec { )), options: None, }); - // add le ASC - cols.push(PhysicalSortRequirement { - expr: Arc::new(PhyCast::new( - Arc::new(PhyColumn::new( - self.input.schema().field(self.le_column_index).name(), - self.le_column_index, + if self.histogram_column_index.is_none() { + // add le ASC + cols.push(PhysicalSortRequirement { + expr: Arc::new(PhyTryCast::new( + Arc::new(PhyColumn::new( + self.input.schema().field(self.le_column_index).name(), + self.le_column_index, + )), + DataType::Float64, )), - DataType::Float64, - None, - )), - options: Some(SortOptions { - descending: false, // +INF in the last - nulls_first: false, // not nullable - }), - }); + options: Some(SortOptions { + descending: false, // numeric bounds ascending + nulls_first: false, // unparsable bounds last + }), + }); + } // Safety: `cols` is not empty let requirement = LexRequirement::new(cols).unwrap(); @@ -475,7 +579,7 @@ impl ExecutionPlan for HistogramFoldExec { } fn maintains_input_order(&self) -> Vec { - vec![true; self.children().len()] + vec![self.histogram_column_index.is_none(); self.children().len()] } fn children(&self) -> Vec<&Arc> { @@ -505,9 +609,10 @@ impl ExecutionPlan for HistogramFoldExec { ts_column_index: self.ts_column_index, tag_columns: self.tag_columns.clone(), partition_exprs: self.partition_exprs.clone(), - quantile: self.quantile, + operation: self.operation, output_schema: self.output_schema.clone(), field_column_index: self.field_column_index, + histogram_column_index: self.histogram_column_index, properties, })) } @@ -526,25 +631,31 @@ impl ExecutionPlan for HistogramFoldExec { let mut normal_indices = (0..input.schema().fields().len()).collect::>(); normal_indices.remove(&self.field_column_index); normal_indices.remove(&self.le_column_index); + if let Some(histogram_column_index) = self.histogram_column_index { + normal_indices.remove(&histogram_column_index); + } + let mode = if self.histogram_column_index.is_some() { + FoldMode::Safe + } else { + FoldMode::Optimistic + }; Ok(Box::pin(HistogramFoldStream { le_column_index: self.le_column_index, field_column_index: self.field_column_index, - quantile: self.quantile, + histogram_column_index: self.histogram_column_index, + operation: self.operation, normal_indices: normal_indices.into_iter().collect(), bucket_size: None, input_buffer: vec![], input, output_schema, input_schema: self.input.schema(), - mode: FoldMode::Optimistic, + mode, safe_group: None, metric: baseline_metric, batch_size, input_buffered_rows: 0, - output_buffer: HistogramFoldStream::empty_output_buffer( - &self.output_schema, - self.le_column_index, - )?, + output_buffer: HistogramFoldStream::empty_output_buffer(&self.input.schema())?, output_buffered_rows: 0, })) } @@ -574,9 +685,13 @@ impl DisplayAs for HistogramFoldExec { | DisplayFormatType::TreeRender => { write!( f, - "HistogramFoldExec: le=@{}, field=@{}, quantile={}", - self.le_column_index, self.field_column_index, self.quantile - ) + "HistogramFoldExec: le=@{}, field=@{}", + self.le_column_index, self.field_column_index + )?; + if let Some(histogram) = self.histogram_column_index { + write!(f, ", histogram=@{histogram}")?; + } + write!(f, ", {}", self.operation) } } } @@ -592,7 +707,8 @@ pub struct HistogramFoldStream { // internal states le_column_index: usize, field_column_index: usize, - quantile: f64, + histogram_column_index: Option, + operation: HistogramFoldOperation, /// Columns need not folding. This indices is based on input schema normal_indices: Vec, bucket_size: Option, @@ -619,6 +735,7 @@ struct SafeGroup { tag_values: Vec, buckets: Vec, counters: Vec, + native_samples: Vec<(Value, Value)>, } impl RecordBatchStream for HistogramFoldStream { @@ -686,25 +803,16 @@ impl HistogramFoldStream { self.maybe_take_output() } - /// Generate a group of empty [MutableVector]s from the output schema. - /// - /// For simplicity, this method will insert a placeholder for `le`. So that - /// the output buffers has the same schema with input. This placeholder needs - /// to be removed before returning the output batch. + /// Generate input-aligned output builders. Classic-only output drops `le` later. pub fn empty_output_buffer( schema: &SchemaRef, - le_column_index: usize, ) -> DataFusionResult>> { - let mut builders = Vec::with_capacity(schema.fields().len() + 1); + let mut builders = Vec::with_capacity(schema.fields().len()); for field in schema.fields() { let concrete_datatype = ConcreteDataType::try_from(field.data_type()).unwrap(); let mutable_vector = concrete_datatype.create_mutable_vector(0); builders.push(mutable_vector); } - builders.insert( - le_column_index, - ConcreteDataType::float64_datatype().create_mutable_vector(0), - ); Ok(builders) } @@ -803,7 +911,7 @@ impl HistogramFoldStream { counters.push(counter); } // ignore invalid data - let result = Self::evaluate_row(self.quantile, &bucket, &counters).unwrap_or(f64::NAN); + let result = Self::evaluate_row(self.operation, &bucket, &counters).unwrap_or(f64::NAN); self.output_buffer[self.field_column_index].push_value_ref(&ValueRef::from(result)); cursor += bucket_num; remaining_rows -= bucket_num; @@ -869,6 +977,13 @@ impl HistogramFoldStream { if !Self::is_positive_infinity(le_array, inf_index) { return false; } + if (cursor..=inf_index).any(|row| { + string_array_value_at_index(le_array, row) + .and_then(|value| value.parse::().ok()) + .is_none() + }) { + return false; + } for offset in 1..bucket_num { let row = cursor + offset; @@ -903,27 +1018,88 @@ impl HistogramFoldStream { self.output_buffered_rows += 1; } - fn finalize_safe_group(&mut self) -> DataFusionResult<()> { - if let Some(group) = self.safe_group.take() { - if group.tag_values.is_empty() { - return Ok(()); - } - - let has_inf = group - .buckets - .last() - .map(|v| v.is_infinite() && v.is_sign_positive()) - .unwrap_or(false); - let result = if group.buckets.len() < 2 || !has_inf { - f64::NAN - } else { - Self::evaluate_row(self.quantile, &group.buckets, &group.counters) - .unwrap_or(f64::NAN) - }; - let mut tag_value_refs = Vec::with_capacity(group.tag_values.len()); - tag_value_refs.extend(group.tag_values.iter().map(|v| v.as_value_ref())); - self.push_output_row(&tag_value_refs, result); + fn push_mixed_output_row( + &mut self, + tag_values: &[Value], + le: &Value, + result: Option, + histogram: &Value, + ) { + let histogram_column_index = self.histogram_column_index.unwrap(); + for (idx, value) in self.normal_indices.iter().zip(tag_values) { + self.output_buffer[*idx].push_value_ref(&value.as_value_ref()); } + self.output_buffer[self.le_column_index].push_value_ref(&le.as_value_ref()); + self.output_buffer[self.field_column_index] + .push_value_ref(&result.map_or(ValueRef::Null, ValueRef::from)); + self.output_buffer[histogram_column_index].push_value_ref(&histogram.as_value_ref()); + self.output_buffered_rows += 1; + } + + fn finalize_safe_group(&mut self) -> DataFusionResult<()> { + let Some(group) = self.safe_group.take() else { + return Ok(()); + }; + if group.tag_values.is_empty() { + return Ok(()); + } + + if self.histogram_column_index.is_some() { + let classic_result = if group.buckets.is_empty() { + None + } else { + let mut buckets = group + .buckets + .into_iter() + .zip(group.counters) + .collect::>(); + buckets.sort_by(|lhs, rhs| lhs.0.total_cmp(&rhs.0)); + let (bounds, counters): (Vec<_>, Vec<_>) = buckets.into_iter().unzip(); + let has_inf = bounds + .last() + .is_some_and(|value| value.is_infinite() && value.is_sign_positive()); + Some(if has_inf { + Self::evaluate_row(self.operation, &bounds, &counters).unwrap_or(f64::NAN) + } else { + f64::NAN + }) + }; + let has_null_native = group.native_samples.iter().any(|(le, _)| le.is_null()); + for (le, histogram) in group.native_samples.iter().filter(|(le, _)| !le.is_null()) { + self.push_mixed_output_row(&group.tag_values, le, None, histogram); + } + for (le, histogram) in group.native_samples.iter().filter(|(le, _)| le.is_null()) { + self.push_mixed_output_row(&group.tag_values, le, classic_result, histogram); + } + if !has_null_native && let Some(result) = classic_result { + self.push_mixed_output_row( + &group.tag_values, + &Value::Null, + Some(result), + &Value::Null, + ); + } + return Ok(()); + } + if group.buckets.is_empty() { + return Ok(()); + } + + let has_inf = group + .buckets + .last() + .is_some_and(|value| value.is_infinite() && value.is_sign_positive()); + let result = if has_inf { + Self::evaluate_row(self.operation, &group.buckets, &group.counters).unwrap_or(f64::NAN) + } else { + f64::NAN + }; + let tag_value_refs = group + .tag_values + .iter() + .map(Value::as_value_ref) + .collect::>(); + self.push_output_row(&tag_value_refs, result); Ok(()) } @@ -955,6 +1131,7 @@ impl HistogramFoldStream { tag_values: tag_values_buf.iter().cloned().map(Value::from).collect(), buckets: Vec::new(), counters: Vec::new(), + native_samples: Vec::new(), }); } @@ -962,17 +1139,28 @@ impl HistogramFoldStream { continue; }; + let mixed = self.histogram_column_index.is_some(); let bucket = string_array_value_at_index(le_array, row) - .and_then(|value| value.parse::().ok()) - .unwrap_or(f64::NAN); - let counter = if field_array.is_valid(row) { - field_array.value(row) - } else { - f64::NAN - }; - - group.buckets.push(bucket); - group.counters.push(counter); + .and_then(|value| value.parse::().ok()); + if let Some(bucket) = bucket + && (!mixed || field_array.is_valid(row)) + { + let counter = if field_array.is_valid(row) { + field_array.value(row) + } else { + f64::NAN + }; + group.buckets.push(bucket); + group.counters.push(counter); + } + if let Some(histogram_column_index) = self.histogram_column_index { + let histogram = vectors[histogram_column_index].get(row); + if !histogram.is_null() { + group + .native_samples + .push((vectors[self.le_column_index].get(row), histogram)); + } + } } Ok(()) @@ -998,14 +1186,15 @@ impl HistogramFoldStream { return Ok(None); } - let mut output_buf = Self::empty_output_buffer(&self.output_schema, self.le_column_index)?; + let mut output_buf = Self::empty_output_buffer(&self.input_schema)?; std::mem::swap(&mut self.output_buffer, &mut output_buf); let mut columns = Vec::with_capacity(output_buf.len()); for builder in output_buf.iter_mut() { columns.push(builder.to_vector().to_arrow_array()); } - // remove the placeholder column for `le` - columns.remove(self.le_column_index); + if self.histogram_column_index.is_none() { + columns.remove(self.le_column_index); + } self.output_buffered_rows = 0; RecordBatch::try_new(self.output_schema.clone(), columns) @@ -1040,12 +1229,18 @@ impl HistogramFoldStream { } /// Evaluate the field column and return the result - fn evaluate_row(quantile: f64, bucket: &[f64], counter: &[f64]) -> DataFusionResult { + fn evaluate_row( + operation: HistogramFoldOperation, + bucket: &[f64], + counter: &[f64], + ) -> DataFusionResult { // check bucket - if bucket.len() <= 1 { + if bucket.is_empty() + || matches!(operation, HistogramFoldOperation::Quantile(_)) && bucket.len() == 1 + { return Ok(f64::NAN); } - if bucket.last().unwrap().is_finite() { + if bucket.last() != Some(&f64::INFINITY) { return Err(DataFusionError::Execution( "last bucket should be +Inf".to_string(), )); @@ -1055,39 +1250,55 @@ impl HistogramFoldStream { "bucket and counter should have the same length".to_string(), )); } - // check quantile - if quantile < 0.0 { - return Ok(f64::NEG_INFINITY); - } else if quantile > 1.0 { - return Ok(f64::INFINITY); - } else if quantile.is_nan() { - return Ok(f64::NAN); + if let HistogramFoldOperation::Quantile(quantile) = operation { + let quantile = f64::from(quantile); + if quantile < 0.0 { + return Ok(f64::NEG_INFINITY); + } else if quantile > 1.0 { + return Ok(f64::INFINITY); + } else if quantile.is_nan() { + return Ok(f64::NAN); + } } // check input value if !bucket.windows(2).all(|w| w[0] <= w[1]) { return Ok(f64::NAN); } - let counter = { - let needs_fix = - counter.iter().any(|v| !v.is_finite()) || !counter.windows(2).all(|w| w[0] <= w[1]); - if !needs_fix { - Cow::Borrowed(counter) - } else { - let mut fixed = Vec::with_capacity(counter.len()); - let mut prev = 0.0; - for (idx, &v) in counter.iter().enumerate() { - let mut val = if v.is_finite() { v } else { prev }; - if idx > 0 && val < prev { - val = prev; + let counter = match operation { + HistogramFoldOperation::Quantile(_) => { + let needs_fix = counter.iter().any(|v| !v.is_finite()) + || !counter.windows(2).all(|w| w[0] <= w[1]); + if !needs_fix { + Cow::Borrowed(counter) + } else { + let mut fixed = Vec::with_capacity(counter.len()); + let mut prev = 0.0; + for (idx, &v) in counter.iter().enumerate() { + let mut val = if v.is_finite() { v } else { prev }; + if idx > 0 && val < prev { + val = prev; + } + fixed.push(val); + prev = val; } - fixed.push(val); - prev = val; + Cow::Owned(fixed) } - Cow::Owned(fixed) } + HistogramFoldOperation::Fraction { .. } => Cow::Borrowed(counter), }; + Ok(match operation { + HistogramFoldOperation::Quantile(quantile) => { + Self::evaluate_quantile(quantile.into(), bucket, &counter) + } + HistogramFoldOperation::Fraction { lower, upper } => { + Self::evaluate_fraction(lower.into(), upper.into(), bucket, &counter) + } + }) + } + + fn evaluate_quantile(quantile: f64, bucket: &[f64], counter: &[f64]) -> f64 { let total = *counter.last().unwrap(); let expected_pos = total * quantile; let mut fit_bucket_pos = 0; @@ -1095,7 +1306,7 @@ impl HistogramFoldStream { fit_bucket_pos += 1; } if fit_bucket_pos >= bucket.len() - 1 { - Ok(bucket[bucket.len() - 2]) + bucket[bucket.len() - 2] } else { let upper_bound = bucket[fit_bucket_pos]; let upper_count = counter[fit_bucket_pos]; @@ -1106,13 +1317,100 @@ impl HistogramFoldStream { lower_count = counter[fit_bucket_pos - 1]; } if (upper_count - lower_count).abs() < 1e-10 { - return Ok(f64::NAN); + return f64::NAN; } - Ok(lower_bound + lower_bound + (upper_bound - lower_bound) / (upper_count - lower_count) - * (expected_pos - lower_count)) + * (expected_pos - lower_count) } } + + fn evaluate_fraction(lower: f64, upper: f64, bucket: &[f64], counter: &[f64]) -> f64 { + let coalesced = bucket + .windows(2) + .any(|bounds| bounds[0] == bounds[1]) + .then(|| { + let mut bounds = Vec::with_capacity(bucket.len()); + let mut counts = Vec::with_capacity(counter.len()); + for (&bound, &count) in bucket.iter().zip(counter) { + if bounds.last() == Some(&bound) { + *counts.last_mut().unwrap() += count; + } else { + bounds.push(bound); + counts.push(count); + } + } + (bounds, counts) + }); + let (bucket, counter) = match &coalesced { + Some((bounds, counts)) => (bounds.as_slice(), counts.as_slice()), + None => (bucket, counter), + }; + let total = *counter.last().unwrap(); + if total == 0.0 || lower.is_nan() || upper.is_nan() { + return f64::NAN; + } + if lower >= upper { + return 0.0; + } + + let mut rank = 0.0; + let mut lower_rank = 0.0; + let mut upper_rank = 0.0; + let mut lower_set = false; + let mut upper_set = false; + let mut lower_bound = if bucket[0] > 0.0 { + 0.0 + } else { + f64::NEG_INFINITY + }; + + for (idx, (&upper_bound, &upper_count)) in bucket.iter().zip(counter).enumerate() { + if idx > 0 { + lower_bound = bucket[idx - 1]; + } + let interpolate = |value: f64| { + if lower_bound == f64::NEG_INFINITY { + upper_count + } else { + rank + (upper_count - rank) * (value - lower_bound) + / (upper_bound - lower_bound) + } + }; + + if !lower_set && lower_bound >= lower { + lower_rank = rank; + lower_set = true; + } + if !upper_set && lower_bound >= upper { + upper_rank = rank; + upper_set = true; + } + if lower_set && upper_set { + break; + } + if !lower_set && lower_bound < lower && upper_bound > lower { + lower_rank = interpolate(lower); + lower_set = true; + } + if !upper_set && lower_bound < upper && upper_bound > upper { + upper_rank = interpolate(upper); + upper_set = true; + } + if lower_set && upper_set { + break; + } + rank = upper_count; + } + + if !lower_set || lower_rank > total { + lower_rank = total; + } + if !upper_set || upper_rank > total { + upper_rank = total; + } + (upper_rank - lower_rank) / total + } } #[cfg(test)] @@ -1209,6 +1507,20 @@ mod test { schema: SchemaRef, quantile: f64, ts_column_index: usize, + ) -> Arc { + build_fold_exec_from_batches_with_operation( + batches, + schema, + HistogramFoldOperation::Quantile(quantile.into()), + ts_column_index, + ) + } + + fn build_fold_exec_from_batches_with_operation( + batches: Vec, + schema: SchemaRef, + operation: HistogramFoldOperation, + ts_column_index: usize, ) -> Arc { let input: Arc = Arc::new(DataSourceExec::new(Arc::new( MemorySourceConfig::try_new(&[batches], schema.clone(), None).unwrap(), @@ -1226,7 +1538,8 @@ mod test { Arc::new(HistogramFoldExec { le_column_index: 1, field_column_index: 2, - quantile, + histogram_column_index: None, + operation, ts_column_index, input, output_schema, @@ -1301,7 +1614,8 @@ mod test { let fold_exec = Arc::new(HistogramFoldExec { le_column_index: 1, field_column_index: 2, - quantile: 0.4, + histogram_column_index: None, + operation: HistogramFoldOperation::Quantile(0.4.into()), ts_column_index: 0, input: memory_exec, output_schema, @@ -1563,6 +1877,46 @@ mod test { assert_eq!(result_literal, expected); } + #[tokio::test] + async fn ignore_unparsable_bucket_bounds() { + let schema = Arc::new(Schema::new(vec![ + Field::new("host", DataType::Utf8, true), + Field::new("le", DataType::Utf8, true), + Field::new("val", DataType::Float64, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec!["a", "a", "a", "b"])), + Arc::new(StringArray::from(vec![ + Some("bad"), + Some("1"), + Some("+Inf"), + None, + ])), + Arc::new(Float64Array::from(vec![1.0, 2.0, 4.0, 1.0])), + ], + ) + .unwrap(); + let fold_exec = build_fold_exec_from_batches(vec![batch], schema, 0.5, 0); + + let result = + datafusion::physical_plan::collect(fold_exec, SessionContext::default().task_ctx()) + .await + .unwrap(); + + assert_eq!(result.len(), 1); + assert_eq!(result[0].num_rows(), 1); + assert_eq!( + string_array_value_at_index(result[0].column(0), 0), + Some("a") + ); + assert_eq!( + result[0].column(1).as_primitive::().value(0), + 1.0 + ); + } + #[tokio::test] async fn safe_mode_handles_misaligned_groups() { let schema = Arc::new(Schema::new(vec![ @@ -1738,8 +2092,12 @@ mod test { ]; for case in cases { - let actual = - HistogramFoldStream::evaluate_row(case.quantile, &bucket, &case.counters).unwrap(); + let actual = HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Quantile(case.quantile.into()), + &bucket, + &case.counters, + ) + .unwrap(); assert_eq!( format!("{actual}"), format!("{}", case.expected), @@ -1753,7 +2111,12 @@ mod test { fn evaluate_out_of_order_input() { let bucket = [0.0, 1.0, 2.0, 3.0, 4.0, f64::INFINITY]; let counters = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]; - let result = HistogramFoldStream::evaluate_row(0.5, &bucket, &counters).unwrap(); + let result = HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Quantile(0.5.into()), + &bucket, + &counters, + ) + .unwrap(); assert_eq!(0.0, result); } @@ -1761,7 +2124,11 @@ mod test { fn evaluate_wrong_bucket() { let bucket = [0.0, 1.0, 2.0, 3.0, 4.0, f64::INFINITY, 5.0]; let counters = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]; - let result = HistogramFoldStream::evaluate_row(0.5, &bucket, &counters); + let result = HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Quantile(0.5.into()), + &bucket, + &counters, + ); assert!(result.is_err()); } @@ -1769,7 +2136,12 @@ mod test { fn evaluate_small_fraction() { let bucket = [0.0, 2.0, 4.0, 6.0, f64::INFINITY]; let counters = [0.0, 1.0 / 300.0, 2.0 / 300.0, 0.01, 0.01]; - let result = HistogramFoldStream::evaluate_row(0.5, &bucket, &counters).unwrap(); + let result = HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Quantile(0.5.into()), + &bucket, + &counters, + ) + .unwrap(); assert_eq!(3.0, result); } @@ -1777,7 +2149,12 @@ mod test { fn evaluate_non_monotonic_counter() { let bucket = [0.0, 1.0, 2.0, 3.0, f64::INFINITY]; let counters = [0.1, 0.2, 0.4, 0.17, 0.5]; - let result = HistogramFoldStream::evaluate_row(0.5, &bucket, &counters).unwrap(); + let result = HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Quantile(0.5.into()), + &bucket, + &counters, + ) + .unwrap(); assert!((result - 1.25).abs() < 1e-10, "{result}"); } @@ -1785,10 +2162,96 @@ mod test { fn evaluate_nan_counter() { let bucket = [0.0, 1.0, 2.0, 3.0, f64::INFINITY]; let counters = [f64::NAN, 1.0, 2.0, 3.0, 3.0]; - let result = HistogramFoldStream::evaluate_row(0.5, &bucket, &counters).unwrap(); + let result = HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Quantile(0.5.into()), + &bucket, + &counters, + ) + .unwrap(); assert!((result - 1.5).abs() < 1e-10, "{result}"); } + #[test] + fn evaluate_classic_histogram_fraction() { + let buckets = [1.0, 2.0, f64::INFINITY]; + let counters = [2.0, 4.0, 4.0]; + let fraction = |lower, upper| { + HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Fraction { + lower: OrderedF64::from(lower), + upper: OrderedF64::from(upper), + }, + &buckets, + &counters, + ) + .unwrap() + }; + + assert_eq!(fraction(0.0, 1.0), 0.5); + assert_eq!(fraction(f64::NEG_INFINITY, f64::INFINITY), 1.0); + assert_eq!(fraction(2.0, 1.0), 0.0); + + assert_eq!( + HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Fraction { + lower: 0.0.into(), + upper: 1.0.into(), + }, + &[1.0, 1.0, f64::INFINITY], + &[1.0, 2.0, 4.0], + ) + .unwrap(), + 0.75 + ); + + assert_eq!( + HistogramFoldStream::evaluate_row( + HistogramFoldOperation::Fraction { + lower: f64::NEG_INFINITY.into(), + upper: f64::INFINITY.into(), + }, + &[f64::INFINITY], + &[4.0], + ) + .unwrap(), + 1.0 + ); + } + + #[tokio::test] + async fn fraction_handles_single_inf_bucket_after_safe_fallback() { + let schema = Arc::new(Schema::new(vec![ + Field::new("host", DataType::Utf8, false), + Field::new("le", DataType::Utf8, false), + Field::new("val", DataType::Float64, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec!["a", "a", "b"])), + Arc::new(StringArray::from(vec!["1", "+Inf", "+Inf"])), + Arc::new(Float64Array::from(vec![2.0, 4.0, 4.0])), + ], + ) + .unwrap(); + let fold = build_fold_exec_from_batches_with_operation( + vec![batch], + schema, + HistogramFoldOperation::Fraction { + lower: f64::NEG_INFINITY.into(), + upper: f64::INFINITY.into(), + }, + 0, + ); + + let batches = + datafusion::physical_plan::collect(fold, SessionContext::default().task_ctx()) + .await + .unwrap(); + let values = batches[0].column(1).as_primitive::(); + assert_eq!(values.values(), &[1.0, 1.0]); + } + fn build_empty_relation(schema: &Arc) -> LogicalPlan { LogicalPlan::EmptyRelation(EmptyRelation { produce_one_row: false, @@ -1812,8 +2275,21 @@ mod test { input_plan.clone(), ) .unwrap(); + let fraction_node = HistogramFold::new_with_operation( + "le".to_string(), + "val".to_string(), + "ts".to_string(), + HistogramFoldOperation::Fraction { + lower: 0.0.into(), + upper: 1.0.into(), + }, + None, + input_plan.clone(), + ) + .unwrap(); + assert!(fraction_node.serialize().is_err()); - let bytes = plan_node.serialize(); + let bytes = plan_node.serialize().unwrap(); let histogram_fold = HistogramFold::deserialize(&bytes).unwrap(); // need fix @@ -1824,7 +2300,10 @@ mod test { assert_eq!(histogram_fold.le_column, "le"); assert_eq!(histogram_fold.ts_column, "ts"); assert_eq!(histogram_fold.field_column, "val"); - assert_eq!(histogram_fold.quantile, OrderedF64::from(0.8)); + assert_eq!( + histogram_fold.operation, + HistogramFoldOperation::Quantile(OrderedF64::from(0.8)) + ); assert_eq!(histogram_fold.output_schema.fields().len(), 2); assert_eq!(histogram_fold.output_schema.field(0).name(), "ts"); assert_eq!(histogram_fold.output_schema.field(1).name(), "val"); diff --git a/src/promql/src/functions/native_histogram.rs b/src/promql/src/functions/native_histogram.rs index 18363c8f771..ffdd0062665 100644 --- a/src/promql/src/functions/native_histogram.rs +++ b/src/promql/src/functions/native_histogram.rs @@ -326,11 +326,17 @@ fn record_counter_reset_contradiction_warning( ); } -fn scalar_histogram_udf( +fn scalar_histogram_udf( name: &'static str, extra_input_types: Vec, - calc: fn(&NativeHistogram, &[ColumnarValue], usize, usize, &'static str) -> DfResult, -) -> ScalarUDF { + calc: F, +) -> ScalarUDF +where + F: Fn(&NativeHistogram, &[ColumnarValue], usize, usize, &'static str) -> DfResult + + Send + + Sync + + 'static, +{ let mut input_types = vec![native_histogram_arrow_type()]; input_types.extend(extra_input_types); create_udf( @@ -795,11 +801,28 @@ impl NativeHistogramQuantile { } pub fn scalar_udf() -> ScalarUDF { + Self::scalar_udf_with_collector(None) + } + + pub fn scalar_udf_with_collector(collector: Option) -> ScalarUDF { scalar_histogram_udf( Self::name(), vec![DataType::Float64], - |histogram, input, row, len, name| { - Ok(histogram.quantile(read_scalar_f64_arg(&input[1], row, len, name)?)) + move |histogram, input, row, len, name| { + let q = read_scalar_f64_arg(&input[1], row, len, name)?; + let (value, info) = histogram.quantile_with_info(q); + if let Some(info) = info { + let message = match info { + NativeHistogramQuantileInfo::NaNSkew => { + "input to histogram_quantile has NaN observations, result is skewed higher" + } + NativeHistogramQuantileInfo::NaNResult => { + "input to histogram_quantile has NaN observations, result is NaN" + } + }; + record_info(&collector, message); + } + Ok(value) }, ) } @@ -813,13 +836,24 @@ impl NativeHistogramFraction { } pub fn scalar_udf() -> ScalarUDF { + Self::scalar_udf_with_collector(None) + } + + pub fn scalar_udf_with_collector(collector: Option) -> ScalarUDF { scalar_histogram_udf( Self::name(), vec![DataType::Float64, DataType::Float64], - |histogram, input, row, len, name| { + move |histogram, input, row, len, name| { let lower = read_scalar_f64_arg(&input[1], row, len, name)?; let upper = read_scalar_f64_arg(&input[2], row, len, name)?; - Ok(histogram.fraction(lower, upper)) + let (value, excluded_nans) = histogram.fraction_with_info(lower, upper); + if excluded_nans { + record_info( + &collector, + "input to histogram_fraction has NaN observations, which are excluded from all fractions", + ); + } + Ok(value) }, ) } @@ -3076,6 +3110,60 @@ mod tests { infos } + #[test] + fn quantile_and_fraction_report_nan_observations() { + let histogram = sample_histogram(10.0, f64::NAN, vec![8.0]); + let histogram_arg = + || ColumnarValue::Array(build_histogram_array(&[Some(histogram.clone())])); + + let quantile_collector = PromqlAnnotationCollector::default(); + let skewed = run_scalar_udf( + NativeHistogramQuantile::scalar_udf_with_collector(Some(quantile_collector.clone())), + vec![ + histogram_arg(), + ColumnarValue::Scalar(ScalarValue::Float64(Some(0.5))), + ], + ); + assert!(skewed.is_finite()); + let nan = run_scalar_udf( + NativeHistogramQuantile::scalar_udf_with_collector(Some(quantile_collector.clone())), + vec![ + histogram_arg(), + ColumnarValue::Scalar(ScalarValue::Float64(Some(0.9))), + ], + ); + assert!(nan.is_nan()); + let infos = collected_infos(&quantile_collector); + assert!( + infos + .iter() + .any(|info| info.ends_with("result is skewed higher")) + ); + assert!(infos.iter().any(|info| info.ends_with("result is NaN"))); + + let fraction_collector = PromqlAnnotationCollector::default(); + assert_eq!( + run_scalar_udf( + NativeHistogramFraction::scalar_udf_with_collector(Some( + fraction_collector.clone(), + )), + vec![ + histogram_arg(), + ColumnarValue::Scalar(ScalarValue::Float64(Some(f64::NEG_INFINITY))), + ColumnarValue::Scalar(ScalarValue::Float64(Some(f64::INFINITY))), + ], + ), + 0.8 + ); + assert_eq!( + collected_infos(&fraction_collector), + vec![ + "input to histogram_fraction has NaN observations, which are excluded from all fractions" + .to_string() + ] + ); + } + #[test] fn mixed_ranges_follow_prometheus_sample_type_semantics() { let first = sample_histogram(1.0, 1.0, vec![1.0]); diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index 83ccc273791..c42eb0e285a 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -33,8 +33,8 @@ use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ExprSchema, JoinType, ScalarValue}; use datafusion_expr::expr::{Exists, ScalarFunction}; use datafusion_expr::{ - AggregateUDF, Expr, ExprSchemable as _, LogicalPlanBuilder, Operator, Subquery, binary_expr, - col, lit, + AggregateUDF, Expr, ExprSchemable as _, Extension, LogicalPlanBuilder, Operator, Subquery, + binary_expr, col, lit, }; use datafusion_functions::datetime::date_bin; use datafusion_functions::datetime::expr_fn::now; @@ -44,6 +44,7 @@ use datatypes::schema::{ColumnSchema, SchemaBuilder, SchemaRef}; use futures::Stream; use futures::task::{Context, Poll}; use pretty_assertions::assert_eq; +use promql::extension_plan::{HistogramFold, HistogramFoldOperation}; use regex::Regex; use store_api::data_source::DataSource; use store_api::storage::ScanRequest; @@ -300,6 +301,72 @@ fn find_merge_scan(plan: &LogicalPlan) -> Option<&MergeScanLogicalPlan> { plan.inputs().into_iter().find_map(find_merge_scan) } +#[test] +fn frontend_only_histogram_folds_stay_above_merge_scan() { + let table = TestTable::table_with_name(0, "t".to_string()); + let table_source = Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(table), + ))); + let input = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) + .unwrap() + .build() + .unwrap(); + + for (name, operation, histogram_column) in [ + ( + "mixed", + HistogramFoldOperation::Quantile(0.5.into()), + Some("pk2".to_string()), + ), + ( + "fraction", + HistogramFoldOperation::Fraction { + lower: 0.0.into(), + upper: 1.0.into(), + }, + None, + ), + ] { + let fold = HistogramFold::new_with_operation( + "pk1".to_string(), + "number".to_string(), + "ts".to_string(), + operation, + histogram_column, + input.clone(), + ) + .unwrap(); + let plan = LogicalPlan::Extension(Extension { + node: Arc::new(fold), + }); + assert!( + DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .is_err(), + "{name}" + ); + + let result = DistPlannerAnalyzer {} + .analyze(plan, &ConfigOptions::default()) + .unwrap(); + let result_text = result.to_string(); + assert!( + result_text.contains("HistogramFold:"), + "{name}\n{result_text}" + ); + + let remote_input = find_merge_scan(&result).unwrap().input(); + let remote_text = remote_input.to_string(); + assert!( + !remote_text.contains("HistogramFold:"), + "{name}\n{remote_text}" + ); + DFLogicalSubstraitConvertor + .encode(remote_input, DefaultSerializer) + .unwrap(); + } +} + #[test] fn dictionary_literals_are_unwrapped_before_pushdown() { let table = Arc::new(Table::new( diff --git a/src/query/src/dist_plan/commutativity.rs b/src/query/src/dist_plan/commutativity.rs index b380fbba05e..1a70fd8e1d9 100644 --- a/src/query/src/dist_plan/commutativity.rs +++ b/src/query/src/dist_plan/commutativity.rs @@ -315,7 +315,10 @@ impl Categorizer { | Expr::Exists(_) | Expr::InList(_) | Expr::Case(_) => Commutativity::Commutative, - Expr::ScalarFunction(_udf) => Commutativity::Commutative, + // Annotation collection must not affect distribution. Keep scalar UDFs pushdownable; + // TODO: Preserve datanode PromQL annotations through plan decoding, return them in + // Flight stream metadata, and merge them on the frontend. + Expr::ScalarFunction(_) => Commutativity::Commutative, Expr::AggregateFunction(_udaf) => Commutativity::Commutative, Expr::Like(_) @@ -376,6 +379,11 @@ mod tests { use datafusion_common::Column; use datafusion_expr::LogicalPlanBuilder; + use datafusion_expr::expr::ScalarFunction; + use datafusion_functions::core::coalesce; + use promql::functions::{ + NativeHistogramDrop, NativeHistogramFraction, NativeHistogramQuantile, + }; use super::*; @@ -396,6 +404,22 @@ mod tests { let commutativity = Categorizer::check_extension_plan(&series_divide, &partition_cols); assert!(matches!(commutativity, Commutativity::Commutative)); } + + #[test] + fn annotated_histogram_helpers_do_not_block_pushdown() { + for udf in [ + NativeHistogramQuantile::scalar_udf(), + NativeHistogramFraction::scalar_udf(), + NativeHistogramDrop::warning_bool_false_udf(String::new(), None), + ] { + let helper = Expr::ScalarFunction(ScalarFunction::new_udf(Arc::new(udf), vec![])); + let expr = Expr::ScalarFunction(ScalarFunction::new_udf(coalesce(), vec![helper])); + assert!(matches!( + Categorizer::check_expr(&expr), + Commutativity::Commutative + )); + } + } } pub type Transformer = Arc Option>; diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 5dfb1a537a8..de0181af7f6 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -41,6 +41,7 @@ use datafusion::logical_expr::{ BinaryExpr, Cast, Extension, LogicalPlan, LogicalPlanBuilder, Operator, ScalarUDF as ScalarUdfDef, WindowFrame, WindowFunctionDefinition, }; +use datafusion::optimizer::simplify_expressions::ExprSimplifier; use datafusion::prelude as df_prelude; use datafusion::prelude::{Column, Expr as DfExpr, JoinType}; use datafusion::scalar::ScalarValue; @@ -49,6 +50,7 @@ use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; use datafusion_common::{DFSchema, NullEquality}; use datafusion_expr::expr::WindowFunctionParams; use datafusion_expr::expr_fn::when; +use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::utils::{conjunction, disjunction}; use datafusion_expr::{ ExprSchemable, Literal, Projection, SortExpr, TableScan, TableSource, col, lit, @@ -59,8 +61,9 @@ use datatypes::data_type::{ConcreteDataType, DataType as GreptimeDataType}; use itertools::Itertools; use once_cell::sync::Lazy; use promql::extension_plan::{ - Absent, EmptyMetric, HistogramFold, InstantManipulate, Millisecond, RangeManipulate, - ScalarCalculate, SeriesDivide, SeriesNormalize, UnionDistinctOn, build_special_time_expr, + Absent, EmptyMetric, HistogramFold, HistogramFoldOperation, InstantManipulate, Millisecond, + RangeManipulate, ScalarCalculate, SeriesDivide, SeriesNormalize, UnionDistinctOn, + build_special_time_expr, }; use promql::functions::{ AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing, @@ -118,6 +121,8 @@ const SCALAR_FUNCTION: &str = "scalar"; const SPECIAL_ABSENT_FUNCTION: &str = "absent"; /// `histogram_quantile` function in PromQL const SPECIAL_HISTOGRAM_QUANTILE: &str = "histogram_quantile"; +/// `histogram_fraction` function in PromQL +const SPECIAL_HISTOGRAM_FRACTION: &str = "histogram_fraction"; /// `vector` function in PromQL const SPECIAL_VECTOR_FUNCTION: &str = "vector"; /// `le` column for conventional histogram. @@ -856,7 +861,11 @@ impl PromPlanner { .context(DataFusionPlanningSnafu); } - let val = Self::get_param_as_literal_expr(param, Some(*op), Some(ArrowDataType::Float64))?; + let val = Self::get_param_as_literal_expr( + param.as_deref(), + Some(*op), + Some(ArrowDataType::Float64), + )?; // convert op and value columns to window exprs. let window_exprs = self.create_window_exprs(*op, group_exprs.clone(), &input)?; @@ -2074,8 +2083,10 @@ impl PromPlanner { let Call { func, args } = call_expr; // some special functions that are not expression but a plan match func.name { - SPECIAL_HISTOGRAM_QUANTILE => { - return self.create_histogram_plan(args, query_engine_state).await; + SPECIAL_HISTOGRAM_QUANTILE | SPECIAL_HISTOGRAM_FRACTION => { + return self + .create_histogram_plan(func.name, args, query_engine_state) + .await; } SPECIAL_VECTOR_FUNCTION => return self.create_vector_plan(args).await, SCALAR_FUNCTION => return self.create_scalar_plan(args, query_engine_state).await, @@ -3221,7 +3232,7 @@ impl PromPlanner { } _ => { - let expr = Self::get_param_as_literal_expr(&Some(arg.clone()), None, None)?; + let expr = Self::get_param_as_literal_expr(Some(arg.as_ref()), None, None)?; result.literals.push(expr); } } @@ -3563,9 +3574,6 @@ impl PromPlanner { "histogram_stdvar" => { ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramStdvar::scalar_udf())) } - "histogram_fraction" => { - ScalarFunc::NativeHistogramUdf(Arc::new(NativeHistogramFraction::scalar_udf())) - } "time" => { exprs.push(build_special_time_expr( self.ctx.time_index_column.as_ref().unwrap(), @@ -4304,8 +4312,11 @@ impl PromPlanner { let expr = match op.id() { token::T_SUM => sum_udaf().call(vec![input]), token::T_QUANTILE => { - let q = - Self::get_param_as_literal_expr(param, Some(op), Some(ArrowDataType::Float64))?; + let q = Self::get_param_as_literal_expr( + param.as_deref(), + Some(op), + Some(ArrowDataType::Float64), + )?; quantile_udaf().call(vec![q, input]) } token::T_AVG => avg_udaf().call(vec![input]), @@ -4580,11 +4591,11 @@ impl PromPlanner { } fn get_param_as_literal_expr( - param: &Option>, + param: Option<&PromExpr>, op: Option, expected_type: Option, ) -> Result { - let prom_param = param.as_deref().with_context(|| { + let prom_param = param.with_context(|| { if let Some(op) = op { FunctionInvalidArgumentSnafu { fn_name: op.to_string(), @@ -4682,63 +4693,80 @@ impl PromPlanner { Ok(normalized_exprs) } - /// Try to build a [f64] from [PromExpr]. - #[deprecated( - note = "use `Self::get_param_as_literal_expr` instead. This is only for `create_histogram_plan`" - )] - fn try_build_float_literal(expr: &PromExpr) -> Option { - match expr { - PromExpr::NumberLiteral(NumberLiteral { val }) => Some(*val), - PromExpr::Paren(ParenExpr { expr }) => Self::try_build_float_literal(expr), - PromExpr::Unary(UnaryExpr { expr, .. }) => { - Self::try_build_float_literal(expr).map(|f| -f) - } - PromExpr::StringLiteral(_) - | PromExpr::Binary(_) - | PromExpr::VectorSelector(_) - | PromExpr::MatrixSelector(_) - | PromExpr::Call(_) - | PromExpr::Extension(_) - | PromExpr::Aggregate(_) - | PromExpr::Subquery(_) => None, - } - } - - /// Create a [SPECIAL_HISTOGRAM_QUANTILE] plan. + /// Create a classic, native, or mixed histogram helper plan. async fn create_histogram_plan( &mut self, + function_name: &str, args: &PromFunctionArgs, query_engine_state: &QueryEngineState, ) -> Result { - if args.args.len() != 2 { - return FunctionInvalidArgumentSnafu { - fn_name: SPECIAL_HISTOGRAM_QUANTILE.to_string(), + let float_literal = |param: &PromExpr| -> Result { + let value = (|| { + let expr = Self::get_param_as_literal_expr( + Some(param), + None, + Some(ArrowDataType::Float64), + ) + .ok()?; + let simplifier = ExprSimplifier::new(SimplifyContext::default()); + let expr = simplifier.coerce(expr, &DFSchema::empty()).ok()?; + let DfExpr::Literal(value, _) = simplifier.simplify(expr).ok()? else { + return None; + }; + let ScalarValue::Float64(Some(value)) = + value.cast_to(&ArrowDataType::Float64).ok()? + else { + return None; + }; + Some(value) + })() + .with_context(|| FunctionInvalidArgumentSnafu { + fn_name: function_name.to_string(), + })?; + Ok(value) + }; + let (function, input) = match (function_name, args.args.as_slice()) { + (SPECIAL_HISTOGRAM_QUANTILE, [quantile, input]) => ( + HistogramFoldOperation::Quantile(float_literal(quantile)?.into()), + input.as_ref().clone(), + ), + (SPECIAL_HISTOGRAM_FRACTION, [lower, upper, input]) => ( + HistogramFoldOperation::Fraction { + lower: float_literal(lower)?.into(), + upper: float_literal(upper)?.into(), + }, + input.as_ref().clone(), + ), + _ => { + return FunctionInvalidArgumentSnafu { + fn_name: function_name.to_string(), + } + .fail(); } - .fail(); - } - #[allow(deprecated)] - let phi = Self::try_build_float_literal(&args.args[0]).with_context(|| { - FunctionInvalidArgumentSnafu { - fn_name: SPECIAL_HISTOGRAM_QUANTILE.to_string(), - } - })?; + }; - let input = args.args[1].as_ref().clone(); let input_plan = self.prom_expr_to_plan(&input, query_engine_state).await?; - // `histogram_quantile` folds buckets across `le`, so `__tsid` (which includes `le`) is not - // a stable series identifier anymore. Also, HistogramFold infers label columns from the - // input schema and must not treat `__tsid` as a label column. + // Histogram helpers fold buckets across `le`, so `__tsid` (which includes `le`) is not a + // stable series identifier anymore. HistogramFold must not treat it as a label column. let input_plan = self.strip_tsid_column(input_plan)?; self.ctx.use_tsid = false; - if let Some(histogram_field) = + if let Some((float_field, histogram_field)) = Self::alternative_sample_columns(input_plan.schema(), &self.ctx.field_columns) - .map(|(_, histogram)| histogram.to_string()) + .map(|(float, histogram)| (float.to_string(), histogram.to_string())) { + if self.ctx.has_le_tag() { + return self.create_mixed_histogram_plan( + function, + input_plan, + float_field, + histogram_field, + ); + } self.ctx.field_columns = vec![histogram_field]; } if self.all_field_columns_are_native_histograms(input_plan.schema()) { - return self.create_native_histogram_quantile_plan(phi, input_plan); + return self.create_native_histogram_plan(function, input_plan); } if !self.ctx.has_le_tag() { @@ -4764,51 +4792,68 @@ impl PromPlanner { .field_columns .first() .with_context(|| FunctionInvalidArgumentSnafu { - fn_name: SPECIAL_HISTOGRAM_QUANTILE.to_string(), + fn_name: function.function_name().to_string(), })? .clone(); // remove le column from tag columns self.ctx.tag_columns.retain(|col| col != LE_COLUMN_NAME); + let fold = HistogramFold::new_with_operation( + LE_COLUMN_NAME.to_string(), + field_column, + time_index_column, + function, + None, + input_plan, + ) + .context(DataFusionPlanningSnafu)?; Ok(LogicalPlan::Extension(Extension { - node: Arc::new( - HistogramFold::new( - LE_COLUMN_NAME.to_string(), - field_column, - time_index_column, - phi, - input_plan, - ) - .context(DataFusionPlanningSnafu)?, - ), + node: Arc::new(fold), })) } - fn create_native_histogram_quantile_plan( + fn create_native_histogram_expr( + &self, + function: HistogramFoldOperation, + field_column: &str, + ) -> DfExpr { + let field = DfExpr::Column(Column::from_name(field_column)); + let (func, args) = match function { + HistogramFoldOperation::Quantile(quantile) => ( + Arc::new(NativeHistogramQuantile::scalar_udf_with_collector( + self.promql_annotations.clone(), + )), + vec![field, lit(f64::from(quantile))], + ), + HistogramFoldOperation::Fraction { lower, upper } => ( + Arc::new(NativeHistogramFraction::scalar_udf_with_collector( + self.promql_annotations.clone(), + )), + vec![field, lit(f64::from(lower)), lit(f64::from(upper))], + ), + }; + DfExpr::ScalarFunction(ScalarFunction { func, args }) + } + + fn create_native_histogram_plan( &mut self, - phi: f64, + function: HistogramFoldOperation, input_plan: LogicalPlan, ) -> Result { ensure!( self.ctx.field_columns.len() == 1, MultiFieldsNotSupportedSnafu { - operator: SPECIAL_HISTOGRAM_QUANTILE + operator: function.function_name() }, ); let field_column = self.ctx.field_columns[0].clone(); - let quantile_expr = DfExpr::ScalarFunction(ScalarFunction { - func: Arc::new(NativeHistogramQuantile::scalar_udf()), - args: vec![ - DfExpr::Column(Column::from_name(field_column)), - DfExpr::Literal(ScalarValue::Float64(Some(phi)), None), - ], - }); - let display_name = quantile_expr.schema_name().to_string(); + let function_expr = self.create_native_histogram_expr(function, &field_column); + let display_name = function_expr.schema_name().to_string(); self.ctx.field_columns = vec![display_name.clone()]; let project_exprs = std::iter::once(self.create_time_index_column_expr()?) - .chain(std::iter::once(quantile_expr.alias(display_name))) + .chain(std::iter::once(function_expr.alias(display_name))) .chain(self.create_tag_column_exprs()?) .collect::>(); @@ -4821,6 +4866,68 @@ impl PromPlanner { .context(DataFusionPlanningSnafu) } + fn create_mixed_histogram_plan( + &mut self, + function: HistogramFoldOperation, + input_plan: LogicalPlan, + float_field: String, + histogram_field: String, + ) -> Result { + let time_index_column = + self.ctx + .time_index_column + .clone() + .with_context(|| TimeIndexNotFoundSnafu { + table: self.ctx.table_name.clone().unwrap_or_default(), + })?; + let tag_columns = self.ctx.tag_columns.clone(); + let folded = HistogramFold::new_with_operation( + LE_COLUMN_NAME.to_string(), + float_field.clone(), + time_index_column.clone(), + function, + Some(histogram_field.clone()), + input_plan, + ) + .context(DataFusionPlanningSnafu)?; + let record_collision = DfExpr::ScalarFunction(ScalarFunction { + func: Arc::new(NativeHistogramDrop::warning_bool_false_udf( + "vector contains a mix of classic and native histograms".to_string(), + self.promql_annotations.clone(), + )), + args: vec![col(&float_field), col(&histogram_field)], + }); + let keep = when( + col(&float_field) + .is_not_null() + .and(col(&histogram_field).is_not_null()), + record_collision, + ) + .otherwise(lit(true)) + .context(DataFusionPlanningSnafu)?; + + let native_expr = self.create_native_histogram_expr(function, &histogram_field); + let output_field = native_expr.schema_name().to_string(); + let value = DfExpr::ScalarFunction(ScalarFunction { + func: coalesce(), + args: vec![col(&float_field), native_expr], + }); + self.ctx.field_columns = vec![output_field.clone()]; + LogicalPlanBuilder::from(LogicalPlan::Extension(Extension { + node: Arc::new(folded), + })) + .filter(keep) + .context(DataFusionPlanningSnafu)? + .project( + std::iter::once(col(&time_index_column)) + .chain(std::iter::once(value.alias(output_field))) + .chain(tag_columns.iter().map(col)), + ) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu) + } + /// Create a [SPECIAL_VECTOR_FUNCTION] plan async fn create_vector_plan(&mut self, args: &PromFunctionArgs) -> Result { if args.args.len() != 1 { @@ -4829,7 +4936,7 @@ impl PromPlanner { } .fail(); } - let lit = Self::get_param_as_literal_expr(&Some(args.args[0].clone()), None, None)?; + let lit = Self::get_param_as_literal_expr(Some(args.args[0].as_ref()), None, None)?; // reuse `SPECIAL_TIME_FUNCTION` as name of time index column self.ctx.time_index_column = Some(SPECIAL_TIME_FUNCTION.to_string()); @@ -5019,8 +5126,9 @@ impl PromPlanner { } } PromExpr::Paren(ParenExpr { expr }) => Self::try_build_literal_expr(expr), - // TODO(ruihang): support Unary operator - PromExpr::Unary(UnaryExpr { expr, .. }) => Self::try_build_literal_expr(expr), + PromExpr::Unary(UnaryExpr { expr, .. }) => Some(DfExpr::Negative(Box::new( + Self::try_build_literal_expr(expr)?, + ))), PromExpr::Binary(PromBinaryExpr { lhs, rhs, @@ -6904,6 +7012,7 @@ mod test { name: &str, table_id: u32, tag: &str, + le: Option<&str>, value: DirectOrValue, ) -> table::TableRef { let value_type = match &value { @@ -6912,12 +7021,20 @@ mod test { DirectOrValue::NativeHistogram(_) => native_histogram_value_type().clone(), DirectOrValue::Utf8(_) => ConcreteDataType::string_datatype(), }; - let schema = Arc::new(Schema::new(vec![ - ColumnSchema::new( - "tag".to_string(), + let tag_count = 1 + usize::from(le.is_some()); + let mut columns = vec![ColumnSchema::new( + "tag".to_string(), + ConcreteDataType::string_datatype(), + false, + )]; + if le.is_some() { + columns.push(ColumnSchema::new( + LE_COLUMN_NAME.to_string(), ConcreteDataType::string_datatype(), false, - ), + )); + } + columns.extend([ ColumnSchema::new( "ts".to_string(), ConcreteDataType::timestamp_millisecond_datatype(), @@ -6925,16 +7042,17 @@ mod test { ) .with_time_index(true), ColumnSchema::new("v".to_string(), value_type, true), - ])); - let batch = RecordBatch::try_new( - schema.arrow_schema().clone(), - vec![ - Arc::new(StringArray::from(vec![tag])) as Arc, - Arc::new(TimestampMillisecondArray::from(vec![1_000])), - value.array(), - ], - ) - .unwrap(); + ]); + let schema = Arc::new(Schema::new(columns)); + let mut arrays = vec![Arc::new(StringArray::from(vec![tag])) as Arc]; + if let Some(le) = le { + arrays.push(Arc::new(StringArray::from(vec![le]))); + } + arrays.extend([ + Arc::new(TimestampMillisecondArray::from(vec![1_000])) as Arc, + value.array(), + ]); + let batch = RecordBatch::try_new(schema.arrow_schema().clone(), arrays).unwrap(); let backing = GreptimeMemTable::new_with_catalog( name, GreptimeRecordBatch::from_df_record_batch(schema.clone(), batch), @@ -6942,11 +7060,12 @@ mod test { DEFAULT_CATALOG_NAME.to_string(), DEFAULT_SCHEMA_NAME.to_string(), ); + let value_index = tag_count + 1; let meta = TableMetaBuilder::empty() .schema(schema) - .primary_key_indices(vec![0]) - .value_indices(vec![2]) - .next_column_id(3) + .primary_key_indices((0..tag_count).collect()) + .value_indices(vec![value_index]) + .next_column_id((value_index + 1) as u32) .build() .unwrap(); let info = Arc::new( @@ -6967,21 +7086,37 @@ mod test { fn operator_table_provider() -> DfTableSourceProvider { let catalog = MemoryCatalogManager::with_default_setup(); let tables = [ - operator_metric_table("lf", 2_001, "a", DirectOrValue::Float64(2.0)), + operator_metric_table("lf", 2_001, "a", None, DirectOrValue::Float64(2.0)), operator_metric_table( "lh", 2_002, "b", + None, DirectOrValue::NativeHistogram(direct_or_histogram()), ), - operator_metric_table("rf", 2_003, "b", DirectOrValue::Float64(3.0)), + operator_metric_table("rf", 2_003, "b", None, DirectOrValue::Float64(3.0)), operator_metric_table( "rh", 2_004, "a", + None, + DirectOrValue::NativeHistogram(direct_or_histogram()), + ), + operator_metric_table("fallback", 2_005, "c", None, DirectOrValue::Float64(7.0)), + operator_metric_table( + "bad_classic", + 2_006, + "d", + Some("broken"), + DirectOrValue::Float64(1.0), + ), + operator_metric_table( + "bad_native", + 2_007, + "d", + None, DirectOrValue::NativeHistogram(direct_or_histogram()), ), - operator_metric_table("fallback", 2_005, "c", DirectOrValue::Float64(7.0)), ]; for table in tables { let info = table.table_info(); @@ -7621,6 +7756,11 @@ mod test { ConcreteDataType::string_datatype(), false, ), + ColumnSchema::new( + LE_COLUMN_NAME.to_string(), + ConcreteDataType::string_datatype(), + true, + ), ColumnSchema::new( "timestamp".to_string(), ConcreteDataType::timestamp_millisecond_datatype(), @@ -7636,8 +7776,8 @@ mod test { let schema = Arc::new(Schema::new(columns)); let table_meta = TableMetaBuilder::empty() .schema(schema) - .primary_key_indices(vec![0]) - .value_indices(vec![2]) + .primary_key_indices(vec![0, 1]) + .value_indices(vec![3]) .next_column_id(1024) .build() .unwrap(); @@ -7816,6 +7956,122 @@ mod test { ) } + fn classic_and_native_histogram_table_provider( + native_tag: &str, + native_le: Option<&str>, + native_histogram: NativeHistogram, + ) -> DfTableSourceProvider { + let table_name = "mixed_histogram"; + let catalog = MemoryCatalogManager::with_default_setup(); + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "tag".to_string(), + ConcreteDataType::string_datatype(), + false, + ), + ColumnSchema::new( + LE_COLUMN_NAME.to_string(), + ConcreteDataType::string_datatype(), + true, + ), + ColumnSchema::new( + "timestamp".to_string(), + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ) + .with_time_index(true), + ColumnSchema::new( + greptime_native_histogram().to_string(), + native_histogram_value_type().clone(), + true, + ), + ColumnSchema::new( + greptime_value().to_string(), + ConcreteDataType::float64_datatype(), + true, + ), + ])); + let table_meta = TableMetaBuilder::empty() + .schema(schema.clone()) + .primary_key_indices(vec![0, 1]) + .value_indices(vec![3, 4]) + .next_column_id(5) + .build() + .unwrap(); + let table_info = Arc::new( + TableInfoBuilder::default() + .name(table_name) + .meta(table_meta) + .build() + .unwrap(), + ); + let batch = RecordBatch::try_new( + schema.arrow_schema().clone(), + vec![ + Arc::new(StringArray::from(vec![ + "classic", "classic", native_tag, "classic", "classic", native_tag, + ])), + Arc::new(StringArray::from(vec![ + Some("1"), + Some("+Inf"), + native_le, + Some("1"), + Some("+Inf"), + native_le, + ])), + Arc::new(TimestampMillisecondArray::from(vec![ + 1_000, 1_000, 1_000, 2_000, 2_000, 2_000, + ])), + build_histogram_array(&[ + None, + None, + Some(native_histogram.clone()), + None, + None, + Some(native_histogram), + ]), + Arc::new(Float64Array::from(vec![ + Some(2.0), + Some(4.0), + None, + Some(2.0), + Some(4.0), + None, + ])), + ], + ) + .unwrap(); + let backing = GreptimeMemTable::new_with_catalog( + table_name, + GreptimeRecordBatch::from_df_record_batch(schema, batch), + 2_200, + DEFAULT_CATALOG_NAME.to_string(), + DEFAULT_SCHEMA_NAME.to_string(), + ); + let table = Arc::new(Table::new( + table_info, + FilterPushDownType::Unsupported, + backing.data_source(), + )); + catalog + .register_table_sync(RegisterTableRequest { + catalog: DEFAULT_CATALOG_NAME.to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + table_id: 2_200, + table, + }) + .unwrap(); + + DfTableSourceProvider::new( + catalog, + false, + QueryContext::arc(), + DummyDecoder::arc(), + false, + ) + } + async fn build_test_table_provider_with_tsid( table_name_tuples: &[(String, String)], num_tag: usize, @@ -10091,7 +10347,7 @@ mod test { let plan = native_histogram_plan("histogram_count(some_metric)").await; assert!(plan.contains("prom_native_histogram_count"), "{plan}"); - assert!(!plan.contains("PromHistogramFold"), "{plan}"); + assert!(!plan.contains("HistogramFold:"), "{plan}"); } #[tokio::test] @@ -10102,6 +10358,7 @@ mod test { "stale_histogram", 2_100, "a", + None, DirectOrValue::NativeHistogram(stale), ); let catalog = MemoryCatalogManager::with_default_setup(); @@ -10211,7 +10468,8 @@ mod test { let plan = native_histogram_plan("histogram_quantile(0.9, some_metric)").await; assert!(plan.contains("prom_native_histogram_quantile"), "{plan}"); - assert!(!plan.contains("PromHistogramFold"), "{plan}"); + assert!(!plan.contains("HistogramFold:"), "{plan}"); + assert!(plan.contains("some_metric.le"), "{plan}"); // The phi literal is threaded into the native quantile UDF as its second argument. assert!(plan.contains("Float64(0.9)"), "{plan}"); // The empty-values filter drops NULL quantile results so the output is empty @@ -10239,6 +10497,213 @@ mod test { assert!(!plan.contains("EmptyRelation"), "{plan}"); } + #[tokio::test] + async fn mixed_histogram_helpers_execute_classic_and_native_samples() { + let state = build_query_engine_state(); + for (query, expected) in [ + ( + "histogram_quantile(0.5, mixed_histogram)", + vec![("classic", 1.0), ("native", 0.0)], + ), + ( + "histogram_fraction(-Inf, +Inf, mixed_histogram)", + vec![("classic", 1.0), ("native", 1.0)], + ), + ] { + let plan = PromPlanner::stmt_to_plan( + classic_and_native_histogram_table_provider("native", None, direct_or_histogram()), + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert!(plan_text.contains("HistogramFold:"), "{plan_text}"); + assert!(plan_text.contains("prom_native_histogram_"), "{plan_text}"); + let value_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + let mut actual = batches + .iter() + .flat_map(|batch| { + let tags = batch + .column_by_name("tag") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column_by_name(&value_field) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()).map(|row| (tags.value(row), values.value(row))) + }) + .collect::>(); + actual.sort_by_key(|(tag, _)| *tag); + assert_eq!(actual, expected, "{query}"); + } + } + + #[tokio::test] + async fn mixed_histogram_helpers_report_annotations() { + let state = build_query_engine_state(); + let mut native_histogram = direct_or_histogram(); + native_histogram.count = 2.0; + native_histogram.sum = f64::NAN; + for (native_tag, expected_rows, expected_warnings, expected_infos) in [ + ( + "classic", + 0, + vec!["vector contains a mix of classic and native histograms"], + vec![], + ), + ( + "native", + 2, + vec![], + vec!["input to histogram_quantile has NaN observations, result is skewed higher"], + ), + ] { + let collector = PromqlAnnotationCollector::default(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + classic_and_native_histogram_table_provider( + native_tag, + None, + native_histogram.clone(), + ), + &operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"), + &state, + Some(collector.clone()), + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected_rows + ); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert_eq!(warnings, expected_warnings); + assert_eq!(infos, expected_infos); + } + } + + #[tokio::test] + async fn mixed_histogram_helper_preserves_native_le_and_scans_once() { + let state = build_query_engine_state(); + let mut stmt = operator_eval_stmt("histogram_quantile(0.5, mixed_histogram)"); + stmt.end = UNIX_EPOCH.checked_add(Duration::from_secs(2)).unwrap(); + let plan = PromPlanner::stmt_to_plan( + classic_and_native_histogram_table_provider( + "classic", + Some("native"), + direct_or_histogram(), + ), + &stmt, + &state, + ) + .await + .unwrap(); + let plan_text = plan.display_indent_schema().to_string(); + assert_eq!( + plan_text.matches("TableScan: mixed_histogram").count(), + 1, + "{plan_text}" + ); + + let value_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + let (_, batches) = execute(plan, &state).await; + let mut actual = batches + .iter() + .flat_map(|batch| { + let le = batch + .column_by_name(LE_COLUMN_NAME) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let timestamps = batch + .column_by_name("timestamp") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column_by_name(&value_field) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()).map(|row| { + ( + timestamps.value(row), + (!le.is_null(row)).then(|| le.value(row).to_string()), + values.value(row), + ) + }) + }) + .collect::>(); + actual.sort_by(|lhs, rhs| (lhs.0, &lhs.1).cmp(&(rhs.0, &rhs.1))); + assert_eq!( + actual, + vec![ + (1_000, None, 1.0), + (1_000, Some("native".to_string()), 0.0), + (2_000, None, 1.0), + (2_000, Some("native".to_string()), 0.0), + ] + ); + } + + #[tokio::test] + async fn nested_histogram_helpers_ignore_unparsable_bucket_labels() { + let state = build_query_engine_state(); + for native_le in [None, Some("native")] { + for query in [ + "histogram_quantile(0.5, histogram_quantile(0.5, mixed_histogram))", + "histogram_fraction(-Inf, +Inf, histogram_fraction(-Inf, +Inf, mixed_histogram))", + ] { + let plan = PromPlanner::stmt_to_plan( + classic_and_native_histogram_table_provider( + "native", + native_le, + direct_or_histogram(), + ), + &operator_eval_stmt(query), + &state, + ) + .await + .unwrap(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 0, + "native_le={native_le:?}, query={query}" + ); + } + } + } + #[tokio::test] async fn native_histogram_quantile_rejects_multi_field_input() { let table_provider = build_test_multi_histogram_table_provider("some_metric").await; @@ -10461,7 +10926,7 @@ mod test { "prom_native_histogram_stdvar", ), ( - "histogram_fraction(0, 1, some_metric)", + "histogram_fraction(-2 + 1, 2 / 2, some_metric)", "prom_native_histogram_fraction", ), ]; @@ -10469,6 +10934,9 @@ mod test { for (query, expected_udf) in cases { let plan = native_histogram_plan(query).await; assert!(plan.contains(expected_udf), "{query}\n{plan}"); + if query.starts_with("histogram_fraction") { + assert!(plan.contains("Float64(-1)"), "{query}\n{plan}"); + } } } @@ -12319,6 +12787,37 @@ Projection: count(prometheus_tsdb_head_series.greptime_value) AS my_series, prom } } + #[tokio::test] + async fn malformed_classic_bucket_does_not_drop_native_histogram() { + let state = build_query_engine_state(); + let collector = PromqlAnnotationCollector::default(); + let plan = PromPlanner::stmt_to_plan_with_annotations( + operator_table_provider(), + &operator_eval_stmt("histogram_quantile(0.5, bad_classic or bad_native)"), + &state, + Some(collector.clone()), + ) + .await + .unwrap(); + let value_field = plan + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + + let (_, batches) = execute(plan, &state).await; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!(values(&batches, &value_field), vec![0.0]); + let mut warnings = vec![]; + let mut infos = vec![]; + collector.append_to(&mut warnings, &mut infos); + assert!(warnings.is_empty()); + assert!(infos.is_empty()); + } + #[tokio::test] async fn test_mixed_binary_operator_aligns_both_alternative_inputs() { let state = build_query_engine_state(); diff --git a/tests/cases/distributed/flow-tql/tsid_on_phy.result b/tests/cases/distributed/flow-tql/tsid_on_phy.result index 68dd01374c1..0a1755b5c2b 100644 --- a/tests/cases/distributed/flow-tql/tsid_on_phy.result +++ b/tests/cases/distributed/flow-tql/tsid_on_phy.result @@ -120,7 +120,7 @@ TQL EXPLAIN ( | | ]] | | physical_plan | HistogramFoldExec: le=@0, field=@4, quantile=0.5 | | | RepartitionExec: REDACTED -| | SortExec: expr=[tag4@1 ASC NULLS LAST, tag5@2 ASC NULLS LAST, ts@3 ASC NULLS LAST, CAST(le@0 AS Float64) ASC NULLS LAST], preserve_partitioning=[false] | +| | SortExec: expr=[tag4@1 ASC NULLS LAST, tag5@2 ASC NULLS LAST, ts@3 ASC NULLS LAST, TRY_CAST(le@0 AS Float64) ASC NULLS LAST], preserve_partitioning=[false] | | | MergeScanExec: REDACTED | | | +---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ diff --git a/tests/cases/standalone/common/promql/histogram_multi_partition.result b/tests/cases/standalone/common/promql/histogram_multi_partition.result index 93168ce449f..3fba2ba8465 100644 --- a/tests/cases/standalone/common/promql/histogram_multi_partition.result +++ b/tests/cases/standalone/common/promql/histogram_multi_partition.result @@ -42,7 +42,7 @@ tql analyze (0, 10, '10s') histogram_quantile(0.5, sum by (le) (histogram_gap_bu | stage | node | plan_| +-+-+-+ | 0_| 0_|_HistogramFoldExec: le=@0, field=@2, quantile=0.5 REDACTED -|_|_|_SortExec: expr=[ts@1 ASC NULLS LAST, CAST(le@0 AS Float64) ASC NULLS LAST], preserve_partitioning=[true] REDACTED +|_|_|_SortExec: expr=[ts@1 ASC NULLS LAST, TRY_CAST(le@0 AS Float64) ASC NULLS LAST], preserve_partitioning=[true] REDACTED |_|_|_RepartitionExec: partitioning=Hash([ts@1],REDACTED |_|_|_AggregateExec: mode=FinalPartitioned, gby=[le@0 as le, ts@1 as ts], aggr=[sum(histogram_gap_bucket.val)] REDACTED |_|_|_RepartitionExec: partitioning=Hash([le@0, ts@1],REDACTED diff --git a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result index af5f3a9a4f8..e50fa43f10e 100644 --- a/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result +++ b/tests/cases/standalone/common/promql/tsid_histogram_quantile_regression.result @@ -110,7 +110,7 @@ TQL ANALYZE (0, 10, '5s') histogram_quantile(0.5, tsid_no_aggr_histogram_bucket) +-+-+-+ | 0_| 0_|_HistogramFoldExec: le=@2, field=@0, quantile=0.5 REDACTED |_|_|_RepartitionExec: partitioning=REDACTED -|_|_|_SortExec: expr=[job@1 ASC NULLS LAST, ts@3 ASC NULLS LAST, CAST(le@2 AS Float64) ASC NULLS LAST], preserve_partitioning=[false] REDACTED +|_|_|_SortExec: expr=[job@1 ASC NULLS LAST, ts@3 ASC NULLS LAST, TRY_CAST(le@2 AS Float64) ASC NULLS LAST], preserve_partitioning=[false] REDACTED |_|_|_MergeScanExec: REDACTED |_|_|_| | 1_| 0_|_ProjectionExec: expr=[val@0 as val, job@1 as job, le@2 as le, ts@4 as ts] REDACTED From 1c5eabcbbfc11f8eb75ccad6411f2a0ee055412c Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Mon, 24 Aug 2026 12:59:44 +0000 Subject: [PATCH 73/92] feat(otlp): support cumulative exponential histograms (#8900) * feat: implement exponential histogram Signed-off-by: shuiyisong * chore: remove duplicate tests Signed-off-by: shuiyisong * fix(otlp): enforce exponential histogram ingestion safety Signed-off-by: shuiyisong * chore: update rfc Signed-off-by: shuiyisong * fix: test Signed-off-by: shuiyisong * fix(otlp): remove protocol-coupled histogram checks Signed-off-by: shuiyisong * perf(otlp): reuse native histogram schema across data points Signed-off-by: shuiyisong * fix: merge repeated OTLP histogram fragments Signed-off-by: shuiyisong * fix(otlp): build rejection messages lazily Signed-off-by: shuiyisong * fix: add doc Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- config/config.md | 2 + config/frontend.example.toml | 2 + config/standalone.example.toml | 2 + docs/rfcs/2026-08-04-native-histograms.md | 109 ++- src/cmd/src/standalone.rs | 4 +- src/cmd/tests/load_config_test.rs | 49 + src/frontend/src/frontend.rs | 4 +- src/frontend/src/instance/otlp.rs | 31 +- src/frontend/src/server.rs | 7 +- src/frontend/src/service_config/otlp.rs | 13 + src/operator/src/insert.rs | 15 +- src/servers/src/http.rs | 9 +- src/servers/src/http/otlp.rs | 75 +- src/servers/src/lib.rs | 1 + src/servers/src/native_histogram.rs | 524 ++++++++++ src/servers/src/otel_arrow.rs | 108 ++- src/servers/src/otlp/metrics.rs | 1051 ++++++++++++++++++++- src/servers/src/prom_remote_write/v2.rs | 562 +---------- src/servers/src/query_handler.rs | 11 +- src/session/src/protocol_ctx.rs | 2 + src/standalone/src/options.rs | 5 +- tests-integration/src/test_util.rs | 22 +- tests-integration/tests/grpc.rs | 199 +++- tests-integration/tests/http.rs | 145 ++- 24 files changed, 2322 insertions(+), 630 deletions(-) create mode 100644 src/servers/src/native_histogram.rs diff --git a/config/config.md b/config/config.md index 1b18565cc08..91cb4fe7669 100644 --- a/config/config.md +++ b/config/config.md @@ -74,6 +74,7 @@ | `jaeger.enable` | Bool | `true` | Whether to enable Jaeger protocol in HTTP API. | | `otlp` | -- | -- | OpenTelemetry protocol options. | | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | +| `otlp.experimental_enable_exponential_histogram` | Bool | `false` | Experimental: enable cumulative OTLP exponential histogram ingestion. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | @@ -316,6 +317,7 @@ | `jaeger.enable` | Bool | `true` | Whether to enable Jaeger protocol in HTTP API. | | `otlp` | -- | -- | OpenTelemetry protocol options. | | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | +| `otlp.experimental_enable_exponential_histogram` | Bool | `false` | Experimental: enable cumulative OTLP exponential histogram ingestion. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index ec1f9c89b42..3635a63f063 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -237,6 +237,8 @@ enable = true [otlp] ## Whether to enable OpenTelemetry protocol in HTTP API. enable = true +## Experimental: enable cumulative OTLP exponential histogram ingestion. +experimental_enable_exponential_histogram = false ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 5ada605b92b..72ba992cce0 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -204,6 +204,8 @@ enable = true [otlp] ## Whether to enable OpenTelemetry protocol in HTTP API. enable = true +## Experimental: enable cumulative OTLP exponential histogram ingestion. +experimental_enable_exponential_histogram = false ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 diff --git a/docs/rfcs/2026-08-04-native-histograms.md b/docs/rfcs/2026-08-04-native-histograms.md index 1e7692bf516..c714b35351b 100644 --- a/docs/rfcs/2026-08-04-native-histograms.md +++ b/docs/rfcs/2026-08-04-native-histograms.md @@ -7,22 +7,23 @@ Author: codex # Summary -GreptimeDB stores a Prometheus native histogram in one Struct-valued field and -evaluates it as a first-class PromQL sample. This document records the supported -protocols, the storage invariant, and the compatibility decisions needed to make -that path predictable. +GreptimeDB stores native histograms in one Prometheus-compatible Struct-valued +field and evaluates them as first-class PromQL samples. Prometheus Remote Write +2.0 supplies native histograms directly. OTLP `ExponentialHistogram` is an +ingestion transport: accepted cumulative points are normalized into the same +Struct before persistence and are queried only as native histograms. -Native histograms are experimental. Prometheus Remote Write 2.0 is the only -supported native-histogram ingestion protocol. Remote Write 1.0 histogram -payloads are rejected instead of being acknowledged and dropped. +Native histograms are experimental. Prometheus Remote Write 2.0 and cumulative +OTLP exponential histograms are supported behind separate configuration gates. +Remote Write 1.0 histogram payloads are rejected instead of being acknowledged and dropped. Native-histogram Remote Read is deferred; the existing Remote Read path continues to return scalar samples only. # Goals 1. Prevent silent native-histogram data loss at protocol boundaries. -2. Keep one stable persisted representation for integer, float, exponential, - and custom-bucket native histograms. +2. Keep one stable Prometheus-compatible persisted representation across + Prometheus and OTLP ingestion. 3. Match Prometheus query behavior where it is observable and practical. 4. State intentional limitations explicitly so incomplete behavior is not mistaken for support. @@ -32,20 +33,37 @@ continues to return scalar samples only. - Supporting native histograms in Remote Write 1.0. - Returning native histograms through Prometheus Remote Read. - Persisting Remote Write metric metadata. -- Ingesting OTLP exponential histograms. - Persisting exemplars. +- Providing an OTLP-specific histogram query surface or reconstructing and + re-exporting the original OTLP point after persistence. - Removing mixed float/histogram handling from PromQL expressions. - Propagating PromQL annotations produced on datanodes back to the frontend. # Data Model and Invariants -Each native histogram is stored in the configured native-histogram field -(`greptime_native_histogram` by default), a Struct whose children preserve the -wire-level schema, zero threshold, sum, reset hint, start timestamp, custom -bounds, spans, and either the integer or float count family. -Integer bucket deltas are converted to absolute integer counts for storage; -float bucket counts are already absolute. No separate sample-kind discriminator -is stored because the populated count family identifies it. +Each accepted native histogram, whether received directly through Remote Write +2.0 or normalized from OTLP `ExponentialHistogram`, is stored in the configured +native-histogram field (`greptime_native_histogram` by default) as the same +canonical Struct. The persistence boundary admits only values accepted by the +shared native-histogram validation. + +## Prometheus Remote Write 2.0 + +Remote Write 2.0 native histograms enter the persistence path in Prometheus +format. The Struct preserves the validated schema, zero threshold, sum, reset +hint, start timestamp, custom bounds, spans, and either the integer or float +count family. Integer bucket deltas are converted to absolute integer counts +for storage; float bucket counts are already absolute. No separate sample-kind +discriminator is stored because the populated count family identifies it. + +## OTLP ExponentialHistogram + +OTLP `ExponentialHistogram` does not introduce a second persisted format. +Accepted cumulative points are normalized into a valid Prometheus native +histogram and then pass through the same validator and Struct encoder. The +Struct does not retain the raw OTLP point or a source-protocol discriminator, +and queries use only the native-histogram PromQL behavior. Detailed conversion +and rejection rules are listed under [OTLP](#otlp). Within one resolved catalog, schema, and physical-table routing context, a metric name has exactly one persisted sample kind: @@ -78,7 +96,7 @@ that send native histograms must use Remote Write 2.0. ## Remote Write 2.0 Remote Write 2.0 accepts integer and float native histograms while -`http.experimental_enable_prometheus_native_histogram` is enabled. Supported +`prom_store.experimental_enable_prometheus_native_histogram` is enabled. Supported exponential schemas are `-4` through `8`; schema `-53` represents native histograms with custom buckets. @@ -103,11 +121,46 @@ explicit if sampled responses are implemented. ## OTLP -OTLP exponential histograms are not converted in this version. The ingestion -branch intentionally remains deferred until it can map temporality, scale, -reset behavior, attributes, and rejected-point reporting into the canonical -native-histogram path. The code carries an explicit TODO rather than a partial -encoder. +OTLP exponential histograms are accepted when +`otlp.experimental_enable_exponential_histogram` is enabled. The option defaults +to false and applies to OTLP/HTTP. Disabled points are rejected rather than +silently acknowledged. OTel Arrow exponential histograms are rejected because +the current Arrow wire format omits `zero_threshold`; accepting them would +silently change the distribution. Cumulative temporality is required; delta and +unspecified exponential histograms are rejected before their points are +converted. Explicit OTLP histograms keep their existing `_bucket`, `_sum`, and +`_count` representation, including their existing delta behavior. + +OTLP scales `-4` through `8` map directly to Prometheus schemas. Higher scales +are downscaled to schema `8`: dense counts that collide are merged before the +OTLP lower-bound index is shifted by one to the Prometheus upper-bound index. +This preserves count mass but irreversibly loses distinctions between source +buckets that merge. Lower scales are rejected. OTLP counts always populate the +integer histogram family. The transmitted non-negative finite zero threshold +and zero count are preserved; the Struct start timestamp is stored in +milliseconds and the reset hint is unknown. Point timestamps and attributes +retain their existing mode-specific conversion rules. Legacy mode retains its +normalized name, attribute rules, and nanosecond row timestamp. Non-legacy mode +retains Prometheus-compatible translation and its millisecond row timestamp. +Both modes write the same canonical Struct. + +An absent sum is stored as an ordinary quiet NaN, so the sample remains +selectable while `histogram_sum` is unknown. Any NaN sum on a point without +`NoRecordedValue` is normalized to the same value, even if its payload has the +Prometheus stale-marker bits. Only an exponential-histogram point carrying OTLP +`NoRecordedValue` becomes an empty schema-0 integer histogram with the canonical +Prometheus stale-NaN sum; its attributes and timestamps remain. This +interpretation intentionally does not change gauges, sums, or explicit histograms. + +Invalid points are skipped while unrelated valid points continue. Mixed +accepted/rejected OTLP/HTTP requests return partial success; a request with only +rejected points returns `InvalidArgument`. OTel Arrow uses an `OK` batch status +for mixed batches and `INVALID_ARGUMENT` when all points are rejected. Rejection +details are bounded, and metric metadata is emitted only for a metric that +produced an accepted row. + +Minimum, maximum, and exemplars are not persisted. Delta accumulation, +zero-run span compaction, and a dedicated rejection metric remain deferred. # PromQL Compatibility Decisions @@ -208,20 +261,22 @@ boundary already recorded in # Testing and Compatibility -Compatibility coverage includes protocol rejection, kind exclusivity, +Compatibility coverage includes protocol rejection, OTLP scale conversion and +partial-success behavior, kind exclusivity, stale-marker selector semantics, synthetic-zero rates, incompatible empty layouts, overflow-safe averages, layout-sensitive equality, infinite custom midpoints, and exponential overflow indices for schemas `-4`, `0`, and `8`. Behavioral coverage also verifies frontend PromQL warning and info responses. -This work does not change GreptimeDB's persisted Struct, protobuf dependencies, -or public configuration. Existing native-histogram data remains readable. +This work does not change GreptimeDB's persisted Struct or protobuf +dependencies. It adds the disabled-by-default public OTLP gate; existing +native-histogram data remains readable. # Future Work - Native-histogram Remote Read, including exact integer round-trips and streamed chunks with start timestamps. -- OTLP exponential-histogram conversion with partial-success reporting. +- OTLP exponential-histogram delta accumulation. - Persistent Remote Write metadata and accurate help/unit updates. - Native-histogram exemplars and exemplar query APIs. - Start-timestamp overlap annotations. diff --git a/src/cmd/src/standalone.rs b/src/cmd/src/standalone.rs index 2ffcc5e729c..2b5bd5e8cec 100644 --- a/src/cmd/src/standalone.rs +++ b/src/cmd/src/standalone.rs @@ -1129,7 +1129,9 @@ mod tests { fn test_toml() { let opts = StandaloneOptions::default(); let toml_string = toml::to_string(&opts).unwrap(); - let _parsed: StandaloneOptions = toml::from_str(&toml_string).unwrap(); + assert!(toml_string.contains("experimental_enable_exponential_histogram = false")); + let parsed: StandaloneOptions = toml::from_str(&toml_string).unwrap(); + assert_eq!(parsed.otlp, opts.otlp); } #[test] diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index fd32680f4c6..70594d0653f 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -153,6 +153,12 @@ fn test_load_frontend_example_config() { let options = GreptimeOptions::::load_layered_options(example_config.to_str(), "") .unwrap(); + assert!( + !options + .component + .otlp + .experimental_enable_exponential_histogram + ); let expected = GreptimeOptions:: { component: FrontendOptions { default_timezone: Some("UTC".to_string()), @@ -329,6 +335,12 @@ fn test_load_standalone_example_config() { let options = GreptimeOptions::::load_layered_options(example_config.to_str(), "") .unwrap(); + assert!( + !options + .component + .otlp + .experimental_enable_exponential_histogram + ); let expected = GreptimeOptions:: { component: StandaloneOptions { default_timezone: Some("UTC".to_string()), @@ -380,6 +392,43 @@ fn test_load_standalone_example_config() { similar_asserts::assert_eq!(options, expected); } +#[test] +fn test_load_otlp_exponential_histogram_option() { + let config = tempfile::NamedTempFile::new().unwrap(); + std::fs::write( + config.path(), + "[otlp]\nexperimental_enable_exponential_histogram = true\n", + ) + .unwrap(); + + let frontend = + GreptimeOptions::::load_layered_options(config.path().to_str(), "") + .unwrap(); + assert!( + frontend + .component + .otlp + .experimental_enable_exponential_histogram + ); + + let standalone = + GreptimeOptions::::load_layered_options(config.path().to_str(), "") + .unwrap(); + assert!( + standalone + .component + .otlp + .experimental_enable_exponential_histogram + ); + assert!( + standalone + .component + .frontend_options() + .otlp + .experimental_enable_exponential_histogram + ); +} + #[test] fn test_load_standalone_user_provider_from_config() { let config = tempfile::NamedTempFile::new().unwrap(); diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index b8fea7b6043..18cedc0545c 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -218,7 +218,9 @@ mod tests { fn test_toml() { let opts = FrontendOptions::default(); let toml_string = toml::to_string(&opts).unwrap(); - let _parsed: FrontendOptions = toml::from_str(&toml_string).unwrap(); + assert!(toml_string.contains("experimental_enable_exponential_histogram = false")); + let parsed: FrontendOptions = toml::from_str(&toml_string).unwrap(); + assert_eq!(parsed.otlp, opts.otlp); } #[test] diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index 459a05321cb..08086198536 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -27,7 +27,7 @@ use client::Output; use common_catalog::consts::{trace_operations_table_name, trace_services_table_name}; use common_error::ext::BoxedError; use common_query::prelude::GREPTIME_PHYSICAL_TABLE; -use common_telemetry::tracing; +use common_telemetry::{tracing, warn}; use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; @@ -38,7 +38,7 @@ use servers::interceptor::{OpenTelemetryProtocolInterceptor, OpenTelemetryProtoc use servers::otlp; use servers::otlp::trace::span::TraceSpanGroup; use servers::query_handler::{ - OpenTelemetryProtocolHandler, PipelineHandlerRef, TraceIngestOutcome, + MetricsIngestOutcome, OpenTelemetryProtocolHandler, PipelineHandlerRef, TraceIngestOutcome, }; use session::context::QueryContextRef; use snafu::ResultExt; @@ -90,7 +90,7 @@ impl OpenTelemetryProtocolHandler for Instance { &self, request: ExportMetricsServiceRequest, ctx: QueryContextRef, - ) -> ServerResult { + ) -> ServerResult { self.plugins .get::() .as_ref() @@ -120,8 +120,19 @@ impl OpenTelemetryProtocolHandler for Instance { .unwrap_or_default(); metric_ctx.is_legacy = is_legacy; - let (requests, rows, semantic_index) = + let (requests, rows, semantic_index, mut outcome) = otlp::metrics::to_grpc_insert_requests(request, &mut metric_ctx)?; + if outcome.rejected_data_points > 0 { + warn!( + "Rejected {} OTLP exponential histogram data points: {}", + outcome.rejected_data_points, + outcome.error_message.as_deref().unwrap_or_default() + ); + } + if outcome.accepted_data_points == 0 { + return Ok(outcome); + } + self.check_row_insert_permission(&requests, &ctx, PermissionReq::Action(OTLP_WRITE)) .context(AuthSnafu)?; self.cache_otlp_legacy(&input_names, &ctx, is_legacy)?; @@ -143,9 +154,9 @@ impl OpenTelemetryProtocolHandler for Instance { Arc::new(c) }; - // If the user uses the legacy path, it is by default without metric engine. - if metric_ctx.is_legacy || !metric_ctx.with_metric_engine { - self.handle_row_inserts(requests, ctx, false, false) + // OTLP tables have one sample field in both the legacy and physical paths. + let output = if metric_ctx.is_legacy || !metric_ctx.with_metric_engine { + self.handle_row_inserts(requests, ctx, false, true) .await .map_err(BoxedError::new) .context(error::ExecuteGrpcQuerySnafu) @@ -154,11 +165,13 @@ impl OpenTelemetryProtocolHandler for Instance { .extension(PHYSICAL_TABLE_PARAM) .unwrap_or(GREPTIME_PHYSICAL_TABLE) .to_string(); - self.handle_metric_row_inserts(requests, ctx, physical_table.clone()) + self.handle_metric_row_inserts(requests, ctx, physical_table) .await .map_err(BoxedError::new) .context(error::ExecuteGrpcQuerySnafu) - } + }?; + outcome.write_cost = output.meta.cost; + Ok(outcome) } #[tracing::instrument(skip_all)] diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index 8fdebd8683c..f0dc226ba85 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -158,8 +158,11 @@ where } if opts.otlp.enable { - builder = builder - .with_otlp_handler(self.instance.clone(), opts.prom_store.with_metric_engine); + builder = builder.with_otlp_handler( + self.instance.clone(), + opts.prom_store.with_metric_engine, + opts.otlp.experimental_enable_exponential_histogram, + ); } if opts.jaeger.enable { diff --git a/src/frontend/src/service_config/otlp.rs b/src/frontend/src/service_config/otlp.rs index 3b01f422894..5dcab2c8ce2 100644 --- a/src/frontend/src/service_config/otlp.rs +++ b/src/frontend/src/service_config/otlp.rs @@ -20,6 +20,7 @@ const DEFAULT_TRACE_INGEST_CHUNK_SIZE: usize = 512; #[serde(default)] pub struct OtlpOptions { pub enable: bool, + pub experimental_enable_exponential_histogram: bool, /// Maximum spans per trace ingest chunk. Set to 0 to disable splitting. pub trace_ingest_chunk_size: usize, } @@ -28,6 +29,7 @@ impl Default for OtlpOptions { fn default() -> Self { Self { enable: true, + experimental_enable_exponential_histogram: false, trace_ingest_chunk_size: DEFAULT_TRACE_INGEST_CHUNK_SIZE, } } @@ -41,10 +43,12 @@ mod tests { fn test_otlp_options() { let default = OtlpOptions::default(); assert!(default.enable); + assert!(!default.experimental_enable_exponential_histogram); assert_eq!(default.trace_ingest_chunk_size, 512); let options: OtlpOptions = toml::from_str("enable = false").unwrap(); assert!(!options.enable); + assert!(!options.experimental_enable_exponential_histogram); assert_eq!( options.trace_ingest_chunk_size, DEFAULT_TRACE_INGEST_CHUNK_SIZE @@ -52,6 +56,15 @@ mod tests { let options: OtlpOptions = toml::from_str("trace_ingest_chunk_size = 0").unwrap(); assert!(options.enable); + assert!(!options.experimental_enable_exponential_histogram); assert_eq!(options.trace_ingest_chunk_size, 0); + + let options: OtlpOptions = + toml::from_str("experimental_enable_exponential_histogram = true").unwrap(); + assert!(options.experimental_enable_exponential_histogram); + + let serialized = toml::to_string(&options).unwrap(); + assert!(serialized.contains("experimental_enable_exponential_histogram = true")); + assert_eq!(toml::from_str::(&serialized).unwrap(), options); } } diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index fabc1b3778b..ead517ab80f 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -982,8 +982,8 @@ impl Inserter { /// When `accommodate_existing_schema` is true, it may modify the input `req` to /// accommodate it with existing schema. See [`create_or_alter_tables_on_demand`](Self::create_or_alter_tables_on_demand) /// for more details. - /// When `accommodate_existing_schema` is true and `is_single_value` is true, it also consider fields when modifying the - /// input `req`. + /// When `is_single_value` is true, it also rejects native-histogram/float kind changes. + /// When both options are true, it considers fields when modifying the input `req`. fn get_alter_table_expr_on_demand( &self, req: &mut RowInsertRequest, @@ -1018,8 +1018,7 @@ impl Inserter { return Ok(None); }; - // If accommodate_existing_schema is true, update request schema for Timestamp/Field columns - if accommodate_existing_schema { + if is_single_value { let request_is_native_histogram = request_is_native_histogram(request_schema); let table_is_native_histogram = table_is_native_histogram(table); ensure!( @@ -1030,6 +1029,10 @@ impl Inserter { ), } ); + } + + // If accommodate_existing_schema is true, update request schema for Timestamp/Field columns + if accommodate_existing_schema { let table_schema = table.schema(); // Find timestamp column name let ts_col_name = table_schema.timestamp_column().map(|c| c.name.clone()); @@ -1638,7 +1641,7 @@ mod tests { }), }; let error = inserter - .get_alter_table_expr_on_demand(&mut histogram_req, &table, &ctx, true, true, true) + .get_alter_table_expr_on_demand(&mut histogram_req, &table, &ctx, false, true, true) .unwrap_err(); assert!( error @@ -1666,7 +1669,7 @@ mod tests { &mut sample_req, &histogram_table, &ctx, - true, + false, true, true, ) diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index bf7478eacfd..c0f6af81843 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -724,11 +724,16 @@ impl HttpServerBuilder { self, handler: OpenTelemetryProtocolHandlerRef, with_metric_engine: bool, + experimental_enable_exponential_histogram: bool, ) -> Self { Self { router: self.router.nest( &format!("/{HTTP_API_VERSION}/otlp"), - HttpServer::route_otlp(handler, with_metric_engine), + HttpServer::route_otlp( + handler, + with_metric_engine, + experimental_enable_exponential_histogram, + ), ), ..self } @@ -1386,6 +1391,7 @@ impl HttpServer { fn route_otlp( otlp_handler: OpenTelemetryProtocolHandlerRef, with_metric_engine: bool, + experimental_enable_exponential_histogram: bool, ) -> Router { Router::new() .route("/v1/metrics", routing::post(otlp::metrics)) @@ -1397,6 +1403,7 @@ impl HttpServer { ) .with_state(OtlpState { with_metric_engine, + experimental_enable_exponential_histogram, handler: otlp_handler, }) } diff --git a/src/servers/src/http/otlp.rs b/src/servers/src/http/otlp.rs index 6540f6d4c48..d04022cdf88 100644 --- a/src/servers/src/http/otlp.rs +++ b/src/servers/src/http/otlp.rs @@ -27,7 +27,9 @@ use mime_guess::mime; use opentelemetry_proto::tonic::collector::logs::v1::{ ExportLogsServiceRequest, ExportLogsServiceResponse, }; -use opentelemetry_proto::tonic::collector::metrics::v1::ExportMetricsServiceResponse; +use opentelemetry_proto::tonic::collector::metrics::v1::{ + ExportMetricsPartialSuccess, ExportMetricsServiceResponse, +}; use opentelemetry_proto::tonic::collector::trace::v1::{ ExportTracePartialSuccess, ExportTraceServiceRequest, ExportTraceServiceResponse, }; @@ -44,7 +46,9 @@ use crate::http::extractor::{ }; use crate::http::header::{CONTENT_TYPE_PROTOBUF, write_cost_header_map}; use crate::metrics::METRIC_HTTP_OPENTELEMETRY_LOGS_ELAPSED; -use crate::query_handler::{OpenTelemetryProtocolHandlerRef, PipelineHandler, TraceIngestOutcome}; +use crate::query_handler::{ + MetricsIngestOutcome, OpenTelemetryProtocolHandlerRef, PipelineHandler, TraceIngestOutcome, +}; #[derive(Clone, prost::Message)] pub struct GoogleRpcStatus { @@ -73,6 +77,7 @@ fn content_type_to_string(content_type: Option<&TypedHeader>) -> St #[derive(Clone)] pub struct OtlpState { pub with_metric_engine: bool, + pub experimental_enable_exponential_histogram: bool, pub handler: OpenTelemetryProtocolHandlerRef, } @@ -84,7 +89,7 @@ pub async fn metrics( http_opts: OtlpMetricOptions, content_type: Option>, bytes: Bytes, -) -> Result> { +) -> Result { if is_json_content_type(content_type.as_ref().map(|h| &h.0)) { return error::UnsupportedJsonContentTypeSnafu {}.fail(); } @@ -103,6 +108,7 @@ pub async fn metrics( let OtlpState { with_metric_engine, + experimental_enable_exponential_histogram, handler, } = state; @@ -111,6 +117,7 @@ pub async fn metrics( resource_attrs: http_opts.resource_attrs, promote_scope_attrs: http_opts.promote_scope_attrs, with_metric_engine, + experimental_enable_exponential_histogram, // set is_legacy later is_legacy: false, metric_type: MetricType::Init, @@ -118,15 +125,15 @@ pub async fn metrics( })); let query_ctx = Arc::new(query_ctx); - handler - .metrics(request, query_ctx) - .await - .map(|o| OtlpResponse { - resp_body: ExportMetricsServiceResponse { - partial_success: None, - }, - write_cost: o.meta.cost, - }) + handler.metrics(request, query_ctx).await.map(|outcome| { + if outcome.accepted_data_points == 0 && outcome.rejected_data_points > 0 { + OtlpMetricsResponse::Failure(outcome) + } else if outcome.rejected_data_points > 0 { + OtlpMetricsResponse::PartialSuccess(outcome) + } else { + OtlpMetricsResponse::FullSuccess(outcome) + } + }) } #[axum_macros::debug_handler] @@ -259,6 +266,50 @@ pub struct OtlpResponse { write_cost: usize, } +pub enum OtlpMetricsResponse { + FullSuccess(MetricsIngestOutcome), + PartialSuccess(MetricsIngestOutcome), + Failure(MetricsIngestOutcome), +} + +impl IntoResponse for OtlpMetricsResponse { + fn into_response(self) -> axum::response::Response { + match self { + OtlpMetricsResponse::FullSuccess(outcome) => { + let mut header_map = write_cost_header_map(outcome.write_cost); + header_map.insert(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.clone()); + let body = ExportMetricsServiceResponse { + partial_success: None, + }; + (header_map, body.encode_to_vec()).into_response() + } + OtlpMetricsResponse::PartialSuccess(outcome) => { + let mut header_map = write_cost_header_map(outcome.write_cost); + header_map.insert(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.clone()); + let body = ExportMetricsServiceResponse { + partial_success: Some(ExportMetricsPartialSuccess { + rejected_data_points: outcome.rejected_data_points, + error_message: outcome.error_message.unwrap_or_default(), + }), + }; + (header_map, body.encode_to_vec()).into_response() + } + OtlpMetricsResponse::Failure(outcome) => { + let status = GoogleRpcStatus { + code: tonic::Code::InvalidArgument as i32, + message: outcome.error_message.unwrap_or_default(), + }; + ( + StatusCode::BAD_REQUEST, + [(header::CONTENT_TYPE, CONTENT_TYPE_PROTOBUF.as_ref())], + status.encode_to_vec(), + ) + .into_response() + } + } + } +} + impl IntoResponse for OtlpResponse { fn into_response(self) -> axum::response::Response { let mut header_map = write_cost_header_map(self.write_cost); diff --git a/src/servers/src/lib.rs b/src/servers/src/lib.rs index 2a3796aab21..79fd392a1a9 100644 --- a/src/servers/src/lib.rs +++ b/src/servers/src/lib.rs @@ -34,6 +34,7 @@ pub mod interceptor; pub mod metrics; pub mod metrics_handler; pub mod mysql; +pub(crate) mod native_histogram; pub mod opentsdb; pub mod otel_arrow; pub mod otlp; diff --git a/src/servers/src/native_histogram.rs b/src/servers/src/native_histogram.rs new file mode 100644 index 00000000000..f1dae526993 --- /dev/null +++ b/src/servers/src/native_histogram.rs @@ -0,0 +1,524 @@ +// Copyright 2023 Greptime Team +// +// 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. + +use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; +use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram}; +use api::helper::ColumnDataTypeWrapper; +use api::v1::value::ValueData; +use api::v1::{ColumnSchema, ListValue, SemanticType, Value}; +use common_query::native_histogram::*; +use common_query::prelude::greptime_native_histogram; +use snafu::{Snafu, ensure}; + +const MAX_NATIVE_HISTOGRAM_SCHEMA: i32 = 8; +const MAX_REDUCIBLE_NATIVE_HISTOGRAM_SCHEMA: i32 = 52; + +#[derive(Debug, Snafu)] +#[snafu(display("{message}"))] +pub(crate) struct NativeHistogramError { + message: String, +} + +type Result = std::result::Result; + +/// Returns the canonical column schema for a native histogram value. +pub(crate) fn native_histogram_column_schema() -> Result { + let (datatype, datatype_extension) = + ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone()) + .map_err(|error| NativeHistogramError { + message: format!("native histogram type cannot be encoded: {error}"), + })? + .into_parts(); + + Ok(ColumnSchema { + column_name: greptime_native_histogram().to_string(), + datatype: datatype as i32, + semantic_type: SemanticType::Field as i32, + datatype_extension, + options: None, + }) +} + +/// Validates and encodes a Prometheus histogram into the canonical Struct value. +pub(crate) fn encode_native_histogram(histogram: &Histogram) -> Result { + let uses_float_counts = native_histogram_uses_float_counts(histogram)?; + validate_native_histogram(histogram, uses_float_counts)?; + + let mut items = Vec::with_capacity(NATIVE_HISTOGRAM_FIELD_NAMES.len()); + let positive_span_lengths = i32_span_lengths("positive", &histogram.positive_spans)?; + let negative_span_lengths = i32_span_lengths("negative", &histogram.negative_spans)?; + items.extend([ + pb_value(ValueData::I32Value(histogram.schema)), + pb_value(ValueData::F64Value(histogram.zero_threshold)), + pb_value(ValueData::F64Value(histogram.sum)), + pb_value(ValueData::I32Value(histogram.reset_hint)), + optional_pb_value((histogram.start_timestamp != 0).then_some( + ValueData::TimestampMillisecondValue(histogram.start_timestamp), + )), + f64_list_value(histogram.custom_values.iter().copied()), + i32_list_value(histogram.positive_spans.iter().map(|span| span.offset)), + i32_list_value(positive_span_lengths), + i32_list_value(histogram.negative_spans.iter().map(|span| span.offset)), + i32_list_value(negative_span_lengths), + ]); + + if uses_float_counts { + validate_float_native_histogram_counts(histogram)?; + let count = match histogram.count.as_ref() { + Some(Count::CountFloat(count)) => *count, + _ => 0.0, + }; + let zero_count = match histogram.zero_count.as_ref() { + Some(ZeroCount::ZeroCountFloat(zero_count)) => *zero_count, + _ => 0.0, + }; + items.extend([ + null_pb_value(), + null_pb_value(), + i64_list_value(std::iter::empty()), + i64_list_value(std::iter::empty()), + pb_value(ValueData::F64Value(count)), + pb_value(ValueData::F64Value(zero_count)), + f64_list_value(histogram.positive_counts.iter().copied()), + f64_list_value(histogram.negative_counts.iter().copied()), + ]); + } else { + let count = match histogram.count.as_ref() { + Some(Count::CountInt(count)) => *count, + _ => 0, + }; + let zero_count = match histogram.zero_count.as_ref() { + Some(ZeroCount::ZeroCountInt(zero_count)) => *zero_count, + _ => 0, + }; + let positive_buckets = bucket_counts_from_deltas(&histogram.positive_deltas)?; + let negative_buckets = bucket_counts_from_deltas(&histogram.negative_deltas)?; + validate_integer_native_histogram_counts(histogram, &positive_buckets, &negative_buckets)?; + let count = i64::try_from(count).map_err(|_| NativeHistogramError { + message: format!("native histogram integer count {count} overflows i64"), + })?; + let zero_count = i64::try_from(zero_count).map_err(|_| NativeHistogramError { + message: format!("native histogram integer zero_count {zero_count} overflows i64"), + })?; + items.extend([ + pb_value(ValueData::I64Value(count)), + pb_value(ValueData::I64Value(zero_count)), + i64_list_value(positive_buckets), + i64_list_value(negative_buckets), + null_pb_value(), + null_pb_value(), + f64_list_value(std::iter::empty()), + f64_list_value(std::iter::empty()), + ]); + } + + Ok(ValueData::StructValue(api::v1::StructValue { items })) +} + +fn validate_native_histogram(histogram: &Histogram, uses_float_counts: bool) -> Result<()> { + let exponential_overflow_index = validate_native_histogram_schema(histogram.schema)?; + validate_native_histogram_custom_values(histogram)?; + + if histogram.schema == CUSTOM_BUCKETS_SCHEMA { + ensure!( + histogram.zero_threshold == 0.0 && native_histogram_zero_count_is_zero(histogram), + NativeHistogramSnafu { + message: "custom native histogram must not use a zero bucket" + } + ); + ensure!( + histogram.negative_spans.is_empty() + && histogram.negative_deltas.is_empty() + && histogram.negative_counts.is_empty(), + NativeHistogramSnafu { + message: "custom native histogram must not use negative buckets" + } + ); + } + + let (positive_buckets, negative_buckets) = if uses_float_counts { + ( + histogram.positive_counts.len(), + histogram.negative_counts.len(), + ) + } else { + ( + histogram.positive_deltas.len(), + histogram.negative_deltas.len(), + ) + }; + let bucket_index_range = if let Some(overflow_index) = exponential_overflow_index { + (i32::MIN, overflow_index) + } else { + ( + 0, + i32::try_from(histogram.custom_values.len()).map_err(|_| NativeHistogramError { + message: "custom native histogram has too many custom_values".to_string(), + })?, + ) + }; + validate_native_histogram_spans( + "positive", + &histogram.positive_spans, + positive_buckets, + bucket_index_range, + )?; + validate_native_histogram_spans( + "negative", + &histogram.negative_spans, + negative_buckets, + bucket_index_range, + )?; + + Ok(()) +} + +fn validate_native_histogram_schema(schema: i32) -> Result> { + if schema == CUSTOM_BUCKETS_SCHEMA { + return Ok(None); + } + + if let Some(overflow_index) = exponential_overflow_bucket_index(schema) { + return Ok(Some(overflow_index)); + } + + if (MAX_NATIVE_HISTOGRAM_SCHEMA + 1..=MAX_REDUCIBLE_NATIVE_HISTOGRAM_SCHEMA).contains(&schema) { + Err(NativeHistogramError { + message: format!("native histogram schema {schema} must be reduced before ingestion"), + }) + } else { + Err(NativeHistogramError { + message: format!("native histogram schema {schema} is unsupported"), + }) + } +} + +fn validate_native_histogram_custom_values(histogram: &Histogram) -> Result<()> { + if histogram.schema != CUSTOM_BUCKETS_SCHEMA { + ensure!( + histogram.custom_values.is_empty(), + NativeHistogramSnafu { + message: "standard native histogram must not use custom_values" + } + ); + return Ok(()); + } + + for value in &histogram.custom_values { + ensure!( + !value.is_nan() && *value != f64::INFINITY, + NativeHistogramSnafu { + message: "custom native histogram custom_values must not contain +Inf or NaN" + } + ); + } + for values in histogram.custom_values.windows(2) { + ensure!( + values[0] < values[1], + NativeHistogramSnafu { + message: "custom native histogram custom_values must be sorted" + } + ); + } + + Ok(()) +} + +fn validate_native_histogram_spans( + name: &str, + spans: &[BucketSpan], + bucket_count: usize, + bucket_index_range: (i32, i32), +) -> Result<()> { + let span_len = spans.iter().try_fold(0usize, |sum, span| { + let length = usize::try_from(span.length).map_err(|_| NativeHistogramError { + message: format!("native histogram {name} span length exceeds usize"), + })?; + sum.checked_add(length).ok_or_else(|| NativeHistogramError { + message: format!("native histogram {name} spans overflow"), + }) + })?; + ensure!( + span_len == bucket_count, + NativeHistogramSnafu { + message: format!( + "native histogram {name} spans describe {span_len} buckets, found {bucket_count}" + ) + } + ); + + let mut current_index = 0i32; + for (span_index, span) in spans.iter().enumerate() { + ensure!( + span.offset >= 0 || (span_index == 0 && bucket_index_range.0 == i32::MIN), + NativeHistogramSnafu { + message: format!( + "native histogram {name} span {} has negative offset {}", + span_index + 1, + span.offset + ) + } + ); + current_index = if span_index == 0 { + span.offset + } else { + current_index + .checked_add(span.offset) + .ok_or_else(|| NativeHistogramError { + message: format!("native histogram {name} span index overflows i32"), + })? + }; + + for _ in 0..span.length { + ensure!( + (bucket_index_range.0..=bucket_index_range.1).contains(¤t_index), + NativeHistogramSnafu { + message: format!( + "native histogram {name} bucket index {current_index} is out of range" + ) + } + ); + current_index = current_index + .checked_add(1) + .ok_or_else(|| NativeHistogramError { + message: format!("native histogram {name} span index overflows i32"), + })?; + } + } + + Ok(()) +} + +fn validate_float_native_histogram_counts(histogram: &Histogram) -> Result<()> { + let count = match histogram.count.as_ref() { + Some(Count::CountFloat(count)) => *count, + _ => 0.0, + }; + ensure!( + count >= 0.0 || count.is_nan(), + NativeHistogramSnafu { + message: "native histogram float count must not be negative" + } + ); + + let zero_count = match histogram.zero_count.as_ref() { + Some(ZeroCount::ZeroCountFloat(zero_count)) => *zero_count, + _ => 0.0, + }; + ensure!( + zero_count >= 0.0 || zero_count.is_nan(), + NativeHistogramSnafu { + message: "native histogram float zero_count must not be negative" + } + ); + + for (name, counts) in [ + ("positive", &histogram.positive_counts), + ("negative", &histogram.negative_counts), + ] { + for (index, count) in counts.iter().enumerate() { + ensure!( + *count >= 0.0 || count.is_nan(), + NativeHistogramSnafu { + message: format!( + "native histogram {name} bucket {} count must not be negative", + index + 1 + ) + } + ); + } + } + + Ok(()) +} + +fn validate_integer_native_histogram_counts( + histogram: &Histogram, + positive_buckets: &[i64], + negative_buckets: &[i64], +) -> Result<()> { + let count = match histogram.count.as_ref() { + Some(Count::CountInt(count)) => *count, + _ => 0, + }; + let zero_count = match histogram.zero_count.as_ref() { + Some(ZeroCount::ZeroCountInt(zero_count)) => *zero_count, + _ => 0, + }; + let bucket_count = + positive_buckets + .iter() + .chain(negative_buckets) + .try_fold(zero_count, |total, bucket| { + let bucket = u64::try_from(*bucket).map_err(|_| NativeHistogramError { + message: "native histogram bucket count is negative".to_string(), + })?; + total + .checked_add(bucket) + .ok_or_else(|| NativeHistogramError { + message: "native histogram bucket total overflows u64".to_string(), + }) + })?; + ensure!( + if histogram.sum.is_nan() { + bucket_count <= count + } else { + bucket_count == count + }, + NativeHistogramSnafu { + message: format!( + "native histogram has {bucket_count} observations in buckets, count is {count}" + ) + } + ); + + Ok(()) +} + +fn native_histogram_zero_count_is_zero(histogram: &Histogram) -> bool { + match histogram.zero_count.as_ref() { + Some(ZeroCount::ZeroCountInt(zero_count)) => *zero_count == 0, + Some(ZeroCount::ZeroCountFloat(zero_count)) => *zero_count == 0.0, + None => true, + } +} + +fn native_histogram_uses_float_counts(histogram: &Histogram) -> Result { + let uses_float_count = matches!(histogram.count, Some(Count::CountFloat(_))) + || matches!(histogram.zero_count, Some(ZeroCount::ZeroCountFloat(_))); + let uses_int_count = matches!(histogram.count, Some(Count::CountInt(_))) + || matches!(histogram.zero_count, Some(ZeroCount::ZeroCountInt(_))); + let uses_float_buckets = + !histogram.positive_counts.is_empty() || !histogram.negative_counts.is_empty(); + let uses_int_buckets = + !histogram.positive_deltas.is_empty() || !histogram.negative_deltas.is_empty(); + + ensure!( + !matches!( + (&histogram.count, &histogram.zero_count), + (Some(Count::CountInt(_)), Some(ZeroCount::ZeroCountFloat(_))) + | (Some(Count::CountFloat(_)), Some(ZeroCount::ZeroCountInt(_))) + ), + NativeHistogramSnafu { + message: "native histogram count and zero_count must use the same integer or float family" + } + ); + ensure!( + !(uses_float_buckets && uses_int_buckets), + NativeHistogramSnafu { + message: "native histogram bucket counts must use either integer deltas or float counts" + } + ); + ensure!( + !(uses_float_count && uses_int_buckets), + NativeHistogramSnafu { + message: "float native histogram must not use integer bucket deltas" + } + ); + ensure!( + !(uses_int_count && uses_float_buckets), + NativeHistogramSnafu { + message: "integer native histogram must not use float bucket counts" + } + ); + + Ok(uses_float_count || uses_float_buckets) +} + +fn pb_value(value_data: ValueData) -> Value { + optional_pb_value(Some(value_data)) +} + +fn null_pb_value() -> Value { + optional_pb_value(None) +} + +fn optional_pb_value(value_data: Option) -> Value { + Value { value_data } +} + +fn list_value(values: impl IntoIterator) -> Value { + pb_value(ValueData::ListValue(ListValue { + items: values.into_iter().map(pb_value).collect(), + })) +} + +fn i32_list_value(values: impl IntoIterator) -> Value { + list_value(values.into_iter().map(ValueData::I32Value)) +} + +fn i32_span_lengths(name: &str, spans: &[BucketSpan]) -> Result> { + spans + .iter() + .map(|span| { + i32::try_from(span.length).map_err(|_| NativeHistogramError { + message: format!( + "native histogram {name} span length {} overflows i32", + span.length + ), + }) + }) + .collect() +} + +fn i64_list_value(values: impl IntoIterator) -> Value { + list_value(values.into_iter().map(ValueData::I64Value)) +} + +fn f64_list_value(values: impl IntoIterator) -> Value { + list_value(values.into_iter().map(ValueData::F64Value)) +} + +fn bucket_counts_from_deltas(deltas: &[i64]) -> Result> { + let mut count = 0_i64; + let mut buckets = Vec::with_capacity(deltas.len()); + + for delta in deltas { + count = count + .checked_add(*delta) + .ok_or_else(|| NativeHistogramError { + message: "native histogram bucket count overflows i64".to_string(), + })?; + ensure!( + count >= 0, + NativeHistogramSnafu { + message: "native histogram bucket count is negative" + } + ); + buckets.push(count); + } + + Ok(buckets) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn shared_errors_are_protocol_neutral() { + let error = encode_native_histogram(&Histogram { + schema: 9, + ..Default::default() + }) + .unwrap_err(); + + assert_eq!( + error.to_string(), + "native histogram schema 9 must be reduced before ingestion" + ); + assert!(!error.to_string().contains("remote write")); + assert!(!error.to_string().contains("OTLP")); + } +} diff --git a/src/servers/src/otel_arrow.rs b/src/servers/src/otel_arrow.rs index f905b39f5bc..320ef87dafe 100644 --- a/src/servers/src/otel_arrow.rs +++ b/src/servers/src/otel_arrow.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; + use auth::UserProviderRef; use common_error::ext::ErrorExt; use common_error::status_code::status_to_tonic_code; @@ -19,14 +21,20 @@ use common_telemetry::error; use futures::SinkExt; use otel_arrow_rust::Consumer; use otel_arrow_rust::proto::opentelemetry::arrow::v1::arrow_metrics_service_server::ArrowMetricsService; -use otel_arrow_rust::proto::opentelemetry::arrow::v1::{BatchArrowRecords, BatchStatus}; +use otel_arrow_rust::proto::opentelemetry::arrow::v1::{ + BatchArrowRecords, BatchStatus, StatusCode as ArrowStatusCode, +}; +use otel_arrow_rust::proto::opentelemetry::metrics::v1::metric; +use session::protocol_ctx::{OtlpMetricCtx, ProtocolCtx}; use tonic::metadata::{Entry, MetadataValue}; use tonic::service::Interceptor; use tonic::{Request, Response, Status, Streaming}; use crate::error; use crate::grpc::context_auth; -use crate::query_handler::OpenTelemetryProtocolHandlerRef; +use crate::query_handler::{MetricsIngestOutcome, OpenTelemetryProtocolHandlerRef}; + +const EXPONENTIAL_HISTOGRAM_UNSUPPORTED: &str = "OTel Arrow exponential histograms are unsupported because the Arrow wire format omits zero_threshold"; pub struct OtelArrowServiceHandler { handler: T, @@ -42,6 +50,31 @@ impl OtelArrowServiceHandler { } } +fn batch_status( + batch_id: i64, + outcome: MetricsIngestOutcome, + has_exponential_histogram_data_points: bool, +) -> BatchStatus { + let status_code = if outcome.accepted_data_points == 0 && outcome.rejected_data_points > 0 { + ArrowStatusCode::InvalidArgument + } else { + ArrowStatusCode::Ok + }; + let status_message = match outcome.error_message { + // Arrow keeps the feature gate off, so these fail before per-point validation. + Some(_) if has_exponential_histogram_data_points => { + EXPONENTIAL_HISTOGRAM_UNSUPPORTED.to_string() + } + Some(message) => message, + None => String::new(), + }; + BatchStatus { + batch_id, + status_code: status_code as i32, + status_message, + } +} + #[async_trait::async_trait] impl ArrowMetricsService for OtelArrowServiceHandler { type ArrowMetricsStream = futures::channel::mpsc::Receiver>; @@ -55,6 +88,11 @@ impl ArrowMetricsService for OtelArrowServiceHandler outcome, + Err(e) => { + let _ = sender + .send(Err(Status::new( + status_to_tonic_code(e.status_code()), + e.to_string(), + ))) + .await; + error!(e; "Failed to ingest metrics from otel-arrow"); + return; + } + }; + let batch_status = + batch_status(batch_id, outcome, has_exponential_histogram_data_points); let _ = sender.send(Ok(batch_status)).await; } }); @@ -131,3 +181,25 @@ impl Interceptor for HeaderInterceptor { Ok(request) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn batch_status_explains_arrow_exponential_histogram_limit() { + let status = batch_status( + 7, + MetricsIngestOutcome { + rejected_data_points: 1, + error_message: Some("internal OTLP rejection detail".to_string()), + ..Default::default() + }, + true, + ); + + assert_eq!(7, status.batch_id); + assert_eq!(ArrowStatusCode::InvalidArgument as i32, status.status_code); + assert_eq!(EXPONENTIAL_HISTOGRAM_UNSUPPORTED, status.status_message); + } +} diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index 6d94b3a139f..332fa46f1d1 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -13,9 +13,16 @@ // limitations under the License. use ahash::HashSet; -use api::v1::{RowInsertRequests, Value}; +use api::greptime_proto::io::prometheus::write::v2::histogram::{ + Count as PromCount, ResetHint, ZeroCount as PromZeroCount, +}; +use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram as PromHistogram}; +use api::v1::value::ValueData; +use api::v1::{RowInsertRequests, SemanticType, Value}; use common_grpc::precision::Precision; +use common_query::native_histogram::native_histogram_value_type; use common_query::prelude::{GREPTIME_COUNT, greptime_timestamp, greptime_value}; +use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; use lazy_static::lazy_static; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, KeyValue, any_value}; @@ -26,8 +33,10 @@ use table::requests::{ SEMANTIC_METRIC_TEMPORALITY, SEMANTIC_METRIC_TYPE, SEMANTIC_METRIC_UNIT, }; -use crate::error::Result; +use crate::error::{self, Result}; +use crate::native_histogram::{encode_native_histogram, native_histogram_column_schema}; use crate::otlp::trace::{KEY_SERVICE_INSTANCE_ID, KEY_SERVICE_NAME}; +use crate::query_handler::MetricsIngestOutcome; use crate::row_writer::{self, MultiTableData, TableData}; pub use crate::semantic::SemanticIndex; use crate::semantic::{ @@ -82,6 +91,9 @@ lazy_static! { const OTEL_SCOPE_NAME: &str = "name"; const OTEL_SCOPE_VERSION: &str = "version"; const OTEL_SCOPE_SCHEMA_URL: &str = "schema_url"; +const MIN_EXPONENTIAL_HISTOGRAM_SCALE: i32 = -4; +const MAX_EXPONENTIAL_HISTOGRAM_SCALE: i32 = 8; +const MAX_REJECTION_MESSAGE_BYTES: usize = 512; /// Convert OpenTelemetry metrics to GreptimeDB insert requests /// @@ -94,9 +106,15 @@ const OTEL_SCOPE_SCHEMA_URL: &str = "schema_url"; pub fn to_grpc_insert_requests( request: ExportMetricsServiceRequest, metric_ctx: &mut OtlpMetricCtx, -) -> Result<(RowInsertRequests, usize, SemanticIndex)> { +) -> Result<( + RowInsertRequests, + usize, + SemanticIndex, + MetricsIngestOutcome, +)> { let mut table_writer = MultiTableData::default(); let mut semantic_index = SemanticIndex::default(); + let mut outcome = MetricsIngestOutcome::default(); for resource in &request.resource_metrics { let resource_attrs = resource.resource.as_ref().map(|r| { @@ -123,13 +141,46 @@ pub fn to_grpc_insert_requests( scope_attrs.as_ref(), metric_ctx, &mut semantic_index, + &mut outcome, )?; } } } let (requests, rows) = table_writer.into_row_insert_requests(); - Ok((requests, rows, semantic_index)) + validate_sample_kinds(&requests)?; + Ok((requests, rows, semantic_index, outcome)) +} + +fn validate_sample_kinds(requests: &RowInsertRequests) -> Result<()> { + for request in &requests.inserts { + let Some(rows) = &request.rows else { + continue; + }; + let field_count = rows + .schema + .iter() + .filter(|column| column.semantic_type == SemanticType::Field as i32) + .count(); + let has_native_histogram = rows.schema.iter().any(|column| { + column.semantic_type == SemanticType::Field as i32 + && api::helper::is_column_type_value_eq( + column.datatype, + column.datatype_extension.clone(), + native_histogram_value_type(), + ) + }); + if has_native_histogram && field_count != 1 { + return Err(error::InvalidParameterSnafu { + reason: format!( + "OTLP metric `{}` cannot mix native histogram and float sample fields", + request.table_name + ), + } + .build()); + } + } + Ok(()) } /// The tables a metric emits and their per-table `metric.type`. Histogram fans @@ -159,8 +210,10 @@ fn emitted_semantic_tables( (format!("{base}{COUNT_TABLE_SUFFIX}"), METRIC_TYPE_COUNTER), (format!("{base}{SUM_TABLE_SUFFIX}"), METRIC_TYPE_COUNTER), ], - // ExponentialHistogram is a no-op today; Init never reaches encoding. - MetricType::ExponentialHistogram | MetricType::Init => vec![], + MetricType::ExponentialHistogram => { + vec![(base.to_string(), METRIC_TYPE_HISTOGRAM)] + } + MetricType::Init => vec![], } } @@ -170,6 +223,7 @@ fn temporality_value(data: &metric::Data) -> Option<&'static str> { let raw = match data { metric::Data::Sum(sum) => sum.aggregation_temporality, metric::Data::Histogram(hist) => hist.aggregation_temporality, + metric::Data::ExponentialHistogram(hist) => hist.aggregation_temporality, _ => return None, }; match AggregationTemporality::try_from(raw) { @@ -311,6 +365,7 @@ fn encode_metrics( scope_attrs: Option<&Vec>, metric_ctx: &OtlpMetricCtx, semantic_index: &mut SemanticIndex, + outcome: &mut MetricsIngestOutcome, ) -> Result<()> { let name = if metric_ctx.is_legacy { legacy_normalize_otlp_name(&metric.name) @@ -322,12 +377,7 @@ fn encode_metrics( ) }; - // Stamp semantic metadata against the same table name(s) the data is written - // to below. `unit` is captured here (it is otherwise discarded by the row - // encoders) along with the declared type/temporality. - record_metric_semantics(semantic_index, metric, &name, metric_ctx); - - if let Some(data) = &metric.data { + let emitted = if let Some(data) = &metric.data { match data { metric::Data::Gauge(gauge) => { encode_gauge( @@ -338,6 +388,8 @@ fn encode_metrics( scope_attrs, metric_ctx, )?; + add_accepted_data_points(outcome, gauge.data_points.len())?; + !gauge.data_points.is_empty() } metric::Data::Sum(sum) => { encode_sum( @@ -348,6 +400,8 @@ fn encode_metrics( scope_attrs, metric_ctx, )?; + add_accepted_data_points(outcome, sum.data_points.len())?; + !sum.data_points.is_empty() } metric::Data::Summary(summary) => { encode_summary( @@ -358,6 +412,8 @@ fn encode_metrics( scope_attrs, metric_ctx, )?; + add_accepted_data_points(outcome, summary.data_points.len())?; + !summary.data_points.is_empty() } metric::Data::Histogram(hist) => { encode_histogram( @@ -368,16 +424,384 @@ fn encode_metrics( scope_attrs, metric_ctx, )?; + add_accepted_data_points(outcome, hist.data_points.len())?; + !hist.data_points.is_empty() } - // TODO: Convert OTLP exponential histograms into the canonical native-histogram - // Struct. See docs/rfcs/2026-08-04-native-histograms.md. - metric::Data::ExponentialHistogram(_hist) => {} + metric::Data::ExponentialHistogram(hist) => encode_exponential_histogram( + table_writer, + &name, + hist, + resource_attrs, + scope_attrs, + metric_ctx, + outcome, + )?, } + } else { + false + }; + + if emitted { + // Stamp semantic metadata only after at least one row was accepted. + record_metric_semantics(semantic_index, metric, &name, metric_ctx); } Ok(()) } +fn add_accepted_data_points(outcome: &mut MetricsIngestOutcome, count: usize) -> Result<()> { + let count = i64::try_from(count).map_err(|_| { + error::InvalidParameterSnafu { + reason: "OTLP metrics data-point count exceeds i64", + } + .build() + })?; + outcome.accepted_data_points = + outcome + .accepted_data_points + .checked_add(count) + .ok_or_else(|| { + error::InvalidParameterSnafu { + reason: "OTLP accepted data-point count overflows i64", + } + .build() + })?; + Ok(()) +} + +fn reject_data_points( + outcome: &mut MetricsIngestOutcome, + count: usize, + reason: impl FnOnce() -> String, +) -> Result<()> { + if count == 0 { + return Ok(()); + } + let count = i64::try_from(count).map_err(|_| { + error::InvalidParameterSnafu { + reason: "OTLP rejected data-point count exceeds i64", + } + .build() + })?; + outcome.rejected_data_points = + outcome + .rejected_data_points + .checked_add(count) + .ok_or_else(|| { + error::InvalidParameterSnafu { + reason: "OTLP rejected data-point count overflows i64", + } + .build() + })?; + append_rejection_message(&mut outcome.error_message, reason); + Ok(()) +} + +fn append_rejection_message(message: &mut Option, reason: impl FnOnce() -> String) { + let message = message.get_or_insert_with(String::new); + let separator = if message.is_empty() { "" } else { "; " }; + let Some(available) = MAX_REJECTION_MESSAGE_BYTES.checked_sub(message.len()) else { + return; + }; + if available <= separator.len() { + return; + } + let reason = reason(); + message.push_str(separator); + + let available = MAX_REJECTION_MESSAGE_BYTES - message.len(); + if reason.len() <= available { + message.push_str(&reason); + return; + } + + const ELLIPSIS: &str = "..."; + let mut end = available.saturating_sub(ELLIPSIS.len()); + while !reason.is_char_boundary(end) { + end -= 1; + } + message.push_str(&reason[..end]); + if available >= ELLIPSIS.len() { + message.push_str(ELLIPSIS); + } +} + +fn encode_exponential_histogram( + table_writer: &mut MultiTableData, + name: &str, + histogram: &ExponentialHistogram, + resource_attrs: Option<&Vec>, + scope_attrs: Option<&Vec>, + metric_ctx: &OtlpMetricCtx, + outcome: &mut MetricsIngestOutcome, +) -> Result { + if !metric_ctx.experimental_enable_exponential_histogram { + reject_data_points(outcome, histogram.data_points.len(), || { + format!( + "metric `{name}` uses OTLP exponential histograms; set otlp.experimental_enable_exponential_histogram = true to enable ingestion" + ) + })?; + return Ok(false); + } + + match AggregationTemporality::try_from(histogram.aggregation_temporality) { + Ok(AggregationTemporality::Cumulative) => {} + Ok(AggregationTemporality::Delta) => { + reject_data_points(outcome, histogram.data_points.len(), || { + format!( + "metric `{name}` uses delta OTLP exponential histograms; only cumulative temporality is supported" + ) + })?; + return Ok(false); + } + _ => { + reject_data_points(outcome, histogram.data_points.len(), || { + format!( + "metric `{name}` has unspecified OTLP exponential histogram temporality; cumulative temporality is required" + ) + })?; + return Ok(false); + } + } + + let column_schema = native_histogram_column_schema().map_err(|error| { + error::InternalSnafu { + err_msg: error.to_string(), + } + .build() + })?; + let mut emitted = false; + for (index, data_point) in histogram.data_points.iter().enumerate() { + let (value, timestamp_nanos) = match exponential_histogram_value(data_point) { + Ok(value) => value, + Err(reason) => { + reject_data_points(outcome, 1, || { + format!("metric `{name}` data point {index}: {reason}") + })?; + continue; + } + }; + + let table = table_writer.get_or_default_table_data( + name, + APPROXIMATE_COLUMN_COUNT, + histogram.data_points.len(), + ); + let mut row = table.alloc_one_row(); + write_tags_and_timestamp( + table, + &mut row, + resource_attrs, + scope_attrs, + Some(data_point.attributes.as_ref()), + timestamp_nanos, + metric_ctx, + )?; + row_writer::write_by_schema( + table, + std::iter::once((column_schema.clone(), Some(value))), + &mut row, + )?; + table.add_row(row); + add_accepted_data_points(outcome, 1)?; + emitted = true; + } + + Ok(emitted) +} + +fn exponential_histogram_value( + data_point: &ExponentialHistogramDataPoint, +) -> std::result::Result<(ValueData, i64), String> { + if data_point.start_time_unix_nano > data_point.time_unix_nano { + return Err(format!( + "start_time_unix_nano {} exceeds time_unix_nano {}", + data_point.start_time_unix_nano, data_point.time_unix_nano + )); + } + + let timestamp_nanos = i64::try_from(data_point.time_unix_nano) + .map_err(|_| format!("time_unix_nano {} overflows i64", data_point.time_unix_nano))?; + let timestamp = timestamp_millis(data_point.time_unix_nano, "time_unix_nano")?; + let start_timestamp = + timestamp_millis(data_point.start_time_unix_nano, "start_time_unix_nano")?; + + let no_recorded_value = data_point.flags & DataPointFlags::NoRecordedValueMask as u32 != 0; + let histogram = if no_recorded_value { + PromHistogram { + sum: f64::from_bits(PROMETHEUS_STALE_NAN_BITS), + schema: 0, + zero_threshold: 0.0, + reset_hint: ResetHint::Unspecified as i32, + timestamp, + start_timestamp, + count: Some(PromCount::CountInt(0)), + zero_count: Some(PromZeroCount::ZeroCountInt(0)), + ..Default::default() + } + } else { + if data_point.scale < MIN_EXPONENTIAL_HISTOGRAM_SCALE { + return Err(format!( + "scale {} is unsupported; minimum supported scale is {}", + data_point.scale, MIN_EXPONENTIAL_HISTOGRAM_SCALE + )); + } + if !data_point.zero_threshold.is_finite() || data_point.zero_threshold < 0.0 { + return Err(format!( + "zero_threshold {} must be finite and non-negative", + data_point.zero_threshold + )); + } + + let downscale_shift = if data_point.scale > MAX_EXPONENTIAL_HISTOGRAM_SCALE { + let shift = data_point + .scale + .checked_sub(MAX_EXPONENTIAL_HISTOGRAM_SCALE) + .ok_or_else(|| "downscale shift overflows i32".to_string())?; + u32::try_from(shift).map_err(|_| "downscale shift exceeds u32".to_string())? + } else { + 0 + }; + let (positive_spans, positive_deltas, positive_count) = + convert_bucket_range("positive", data_point.positive.as_ref(), downscale_shift)?; + let (negative_spans, negative_deltas, negative_count) = + convert_bucket_range("negative", data_point.negative.as_ref(), downscale_shift)?; + let bucket_count = data_point + .zero_count + .checked_add(positive_count) + .and_then(|count| count.checked_add(negative_count)) + .ok_or_else(|| "bucket observation total overflows u64".to_string())?; + if bucket_count != data_point.count { + return Err(format!( + "buckets contain {bucket_count} observations, declared count is {}", + data_point.count + )); + } + + i64::try_from(data_point.count) + .map_err(|_| format!("count {} overflows i64", data_point.count))?; + i64::try_from(data_point.zero_count) + .map_err(|_| format!("zero_count {} overflows i64", data_point.zero_count))?; + + let sum = match data_point.sum { + Some(sum) if sum.is_nan() => f64::NAN, + Some(sum) => sum, + None => f64::NAN, + }; + PromHistogram { + sum, + schema: data_point.scale.min(MAX_EXPONENTIAL_HISTOGRAM_SCALE), + zero_threshold: data_point.zero_threshold, + negative_spans, + negative_deltas, + positive_spans, + positive_deltas, + reset_hint: ResetHint::Unspecified as i32, + timestamp, + start_timestamp, + count: Some(PromCount::CountInt(data_point.count)), + zero_count: Some(PromZeroCount::ZeroCountInt(data_point.zero_count)), + ..Default::default() + } + }; + + encode_native_histogram(&histogram) + .map(|value| (value, timestamp_nanos)) + .map_err(|error| format!("OTLP exponential histogram cannot be encoded: {error}")) +} + +fn timestamp_millis(timestamp_nanos: u64, name: &str) -> std::result::Result { + i64::try_from(timestamp_nanos / 1_000_000) + .map_err(|_| format!("{name} {timestamp_nanos} milliseconds overflow i64")) +} + +fn convert_bucket_range( + name: &str, + buckets: Option<&exponential_histogram_data_point::Buckets>, + downscale_shift: u32, +) -> std::result::Result<(Vec, Vec, u64), String> { + let Some(buckets) = buckets else { + return Ok((Vec::new(), Vec::new(), 0)); + }; + if buckets.bucket_counts.is_empty() { + return Ok((Vec::new(), Vec::new(), 0)); + } + + let mut merged = Vec::<(i32, u64)>::with_capacity(buckets.bucket_counts.len()); + let mut total = 0u64; + for (position, count) in buckets.bucket_counts.iter().copied().enumerate() { + let position = + i32::try_from(position).map_err(|_| format!("{name} bucket length exceeds i32"))?; + let source_index = buckets.offset.checked_add(position).ok_or_else(|| { + format!( + "{name} bucket index overflows i32 at offset {} and position {position}", + buckets.offset + ) + })?; + let target_index = downscale_bucket_index(source_index, downscale_shift)? + .checked_add(1) + .ok_or_else(|| format!("{name} shifted bucket index overflows i32"))?; + if let Some((last_index, last_count)) = merged.last_mut() { + if *last_index == target_index { + *last_count = last_count + .checked_add(count) + .ok_or_else(|| format!("{name} merged bucket count overflows u64"))?; + total = total + .checked_add(count) + .ok_or_else(|| format!("{name} bucket count total overflows u64"))?; + continue; + } + let next_index = last_index + .checked_add(1) + .ok_or_else(|| format!("{name} target bucket index overflows i32"))?; + if target_index != next_index { + return Err(format!( + "{name} bucket indexes are not contiguous after downscaling" + )); + } + } + total = total + .checked_add(count) + .ok_or_else(|| format!("{name} bucket count total overflows u64"))?; + merged.push((target_index, count)); + } + + let length = u32::try_from(merged.len()) + .map_err(|_| format!("{name} bucket span length exceeds u32"))?; + let span = BucketSpan { + offset: merged[0].0, + length, + }; + let mut deltas = Vec::with_capacity(merged.len()); + let mut previous = 0i64; + for (_, count) in merged { + let count = i64::try_from(count) + .map_err(|_| format!("{name} bucket count {count} overflows i64"))?; + let delta = count + .checked_sub(previous) + .ok_or_else(|| format!("{name} bucket delta overflows i64"))?; + deltas.push(delta); + previous = count; + } + + Ok((vec![span], deltas, total)) +} + +fn downscale_bucket_index(index: i32, downscale_shift: u32) -> std::result::Result { + if downscale_shift == 0 { + return Ok(index); + } + if downscale_shift >= i32::BITS { + return Ok(if index < 0 { -1 } else { 0 }); + } + + let divisor = 1i64 + .checked_shl(downscale_shift) + .ok_or_else(|| format!("downscale shift {downscale_shift} is invalid"))?; + i32::try_from(i64::from(index).div_euclid(divisor)) + .map_err(|_| format!("downscaled bucket index {index} overflows i32")) +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum AttributeType { Resource, @@ -620,17 +1044,17 @@ fn encode_histogram( let count_table_name = format!("{}{}", normalized_name, COUNT_TABLE_SUFFIX); let data_points_len = hist.data_points.len(); - // Note that the row and columns number here is approximate - let mut bucket_table = TableData::new(APPROXIMATE_COLUMN_COUNT, data_points_len * 3); - let mut sum_table = TableData::new(APPROXIMATE_COLUMN_COUNT, data_points_len); - let mut count_table = TableData::new(APPROXIMATE_COLUMN_COUNT, data_points_len); - for data_point in &hist.data_points { + let bucket_table = table_writer.get_or_default_table_data( + &bucket_table_name, + APPROXIMATE_COLUMN_COUNT, + data_points_len * 3, + ); let mut accumulated_count = 0; for (idx, count) in data_point.bucket_counts.iter().enumerate() { let mut bucket_row = bucket_table.alloc_one_row(); write_tags_and_timestamp( - &mut bucket_table, + bucket_table, &mut bucket_row, resource_attrs, scope_attrs, @@ -641,7 +1065,7 @@ fn encode_histogram( if let Some(upper_bounds) = data_point.explicit_bounds.get(idx) { row_writer::write_tag( - &mut bucket_table, + bucket_table, HISTOGRAM_LE_COLUMN, upper_bounds, &mut bucket_row, @@ -649,7 +1073,7 @@ fn encode_histogram( } else if idx == data_point.explicit_bounds.len() { // The last bucket row_writer::write_tag( - &mut bucket_table, + bucket_table, HISTOGRAM_LE_COLUMN, f64::INFINITY, &mut bucket_row, @@ -658,7 +1082,7 @@ fn encode_histogram( accumulated_count += count; row_writer::write_f64( - &mut bucket_table, + bucket_table, greptime_value(), accumulated_count as f64, &mut bucket_row, @@ -668,9 +1092,14 @@ fn encode_histogram( } if let Some(sum) = data_point.sum { + let sum_table = table_writer.get_or_default_table_data( + &sum_table_name, + APPROXIMATE_COLUMN_COUNT, + data_points_len, + ); let mut sum_row = sum_table.alloc_one_row(); write_tags_and_timestamp( - &mut sum_table, + sum_table, &mut sum_row, resource_attrs, scope_attrs, @@ -679,13 +1108,18 @@ fn encode_histogram( metric_ctx, )?; - row_writer::write_f64(&mut sum_table, greptime_value(), sum, &mut sum_row)?; + row_writer::write_f64(sum_table, greptime_value(), sum, &mut sum_row)?; sum_table.add_row(sum_row); } + let count_table = table_writer.get_or_default_table_data( + &count_table_name, + APPROXIMATE_COLUMN_COUNT, + data_points_len, + ); let mut count_row = count_table.alloc_one_row(); write_tags_and_timestamp( - &mut count_table, + count_table, &mut count_row, resource_attrs, scope_attrs, @@ -695,7 +1129,7 @@ fn encode_histogram( )?; row_writer::write_f64( - &mut count_table, + count_table, greptime_value(), data_point.count as f64, &mut count_row, @@ -703,10 +1137,6 @@ fn encode_histogram( count_table.add_row(count_row); } - table_writer.add_table_data(bucket_table_name, bucket_table); - table_writer.add_table_data(sum_table_name, sum_table); - table_writer.add_table_data(count_table_name, count_table); - Ok(()) } @@ -851,6 +1281,7 @@ mod tests { use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ AggregationTemporality, HistogramDataPoint, NumberDataPoint, SummaryDataPoint, }; + use otel_arrow_rust::proto::opentelemetry::resource::v1::Resource; use super::*; @@ -1361,4 +1792,560 @@ mod tests { ); } } + + fn exponential_buckets( + offset: i32, + bucket_counts: Vec, + ) -> exponential_histogram_data_point::Buckets { + exponential_histogram_data_point::Buckets { + offset, + bucket_counts, + } + } + + fn exponential_point() -> ExponentialHistogramDataPoint { + ExponentialHistogramDataPoint { + start_time_unix_nano: 1_000_000, + time_unix_nano: 2_000_000, + count: 28, + sum: None, + scale: 9, + zero_count: 7, + positive: Some(exponential_buckets(-3, vec![1, 2, 3, 4, 5, 6])), + zero_threshold: 0.0, + ..Default::default() + } + } + + fn native_field(value: &ValueData, name: &str) -> Option { + let ValueData::StructValue(value) = value else { + panic!("expected native histogram Struct value"); + }; + let index = common_query::native_histogram::NATIVE_HISTOGRAM_FIELD_NAMES + .iter() + .position(|field| *field == name) + .unwrap(); + value.items[index].value_data.clone() + } + + fn i32_list(value: Option) -> Vec { + let Some(ValueData::ListValue(value)) = value else { + panic!("expected i32 list"); + }; + value + .items + .into_iter() + .map(|item| match item.value_data { + Some(ValueData::I32Value(value)) => value, + _ => panic!("expected i32 value"), + }) + .collect() + } + + fn i64_list(value: Option) -> Vec { + let Some(ValueData::ListValue(value)) = value else { + panic!("expected i64 list"); + }; + value + .items + .into_iter() + .map(|item| match item.value_data { + Some(ValueData::I64Value(value)) => value, + _ => panic!("expected i64 value"), + }) + .collect() + } + + #[test] + fn test_downscale_bucket_index_uses_signed_floor_division() { + assert_eq!(downscale_bucket_index(-3, 1).unwrap(), -2); + assert_eq!(downscale_bucket_index(-2, 1).unwrap(), -1); + assert_eq!(downscale_bucket_index(-1, 1).unwrap(), -1); + assert_eq!(downscale_bucket_index(0, 1).unwrap(), 0); + assert_eq!(downscale_bucket_index(1, 1).unwrap(), 0); + assert_eq!(downscale_bucket_index(2, 1).unwrap(), 1); + assert_eq!(downscale_bucket_index(i32::MIN, 32).unwrap(), -1); + assert_eq!(downscale_bucket_index(i32::MAX, 32).unwrap(), 0); + } + + #[test] + fn test_convert_bucket_range_downscales_before_prometheus_shift() { + let buckets = exponential_buckets(-3, vec![1, 2, 3, 4, 5, 6]); + let (spans, deltas, total) = convert_bucket_range("positive", Some(&buckets), 1).unwrap(); + + assert_eq!( + spans, + vec![BucketSpan { + offset: -1, + length: 4 + }] + ); + assert_eq!(deltas, vec![1, 4, 4, -3]); + assert_eq!(total, 21); + } + + #[test] + fn test_exponential_histogram_value_uses_integer_family() { + use common_query::native_histogram::{ + COUNT_F64_FIELD, COUNT_I64_FIELD, POSITIVE_BUCKETS_I64_FIELD, + POSITIVE_SPAN_LENGTHS_FIELD, POSITIVE_SPAN_OFFSETS_FIELD, SCHEMA_FIELD, SUM_FIELD, + ZERO_COUNT_I64_FIELD, + }; + + let (value, timestamp_nanos) = exponential_histogram_value(&exponential_point()).unwrap(); + + assert_eq!(timestamp_nanos, 2_000_000); + assert_eq!( + native_field(&value, SCHEMA_FIELD), + Some(ValueData::I32Value(8)) + ); + assert_eq!( + native_field(&value, COUNT_I64_FIELD), + Some(ValueData::I64Value(28)) + ); + assert_eq!( + native_field(&value, ZERO_COUNT_I64_FIELD), + Some(ValueData::I64Value(7)) + ); + assert_eq!(native_field(&value, COUNT_F64_FIELD), None); + assert_eq!( + i32_list(native_field(&value, POSITIVE_SPAN_OFFSETS_FIELD)), + vec![-1] + ); + assert_eq!( + i32_list(native_field(&value, POSITIVE_SPAN_LENGTHS_FIELD)), + vec![4] + ); + assert_eq!( + i64_list(native_field(&value, POSITIVE_BUCKETS_I64_FIELD)), + vec![1, 5, 9, 6] + ); + let Some(ValueData::F64Value(sum)) = native_field(&value, SUM_FIELD) else { + panic!("expected histogram sum"); + }; + assert_eq!(sum.to_bits(), f64::NAN.to_bits()); + assert_ne!(sum.to_bits(), PROMETHEUS_STALE_NAN_BITS); + } + + #[test] + fn test_no_recorded_value_ignores_other_value_fields() { + use common_query::native_histogram::{ + COUNT_I64_FIELD, POSITIVE_BUCKETS_I64_FIELD, SCHEMA_FIELD, SUM_FIELD, + ZERO_COUNT_I64_FIELD, ZERO_THRESHOLD_FIELD, + }; + + let point = ExponentialHistogramDataPoint { + start_time_unix_nano: 1_000_000, + time_unix_nano: 2_000_000, + count: u64::MAX, + sum: Some(1.0), + scale: i32::MIN, + zero_count: u64::MAX, + positive: Some(exponential_buckets(i32::MAX, vec![u64::MAX])), + flags: DataPointFlags::NoRecordedValueMask as u32, + zero_threshold: f64::NAN, + ..Default::default() + }; + let (value, _) = exponential_histogram_value(&point).unwrap(); + + assert_eq!( + native_field(&value, SCHEMA_FIELD), + Some(ValueData::I32Value(0)) + ); + assert_eq!( + native_field(&value, ZERO_THRESHOLD_FIELD), + Some(ValueData::F64Value(0.0)) + ); + assert_eq!( + native_field(&value, COUNT_I64_FIELD), + Some(ValueData::I64Value(0)) + ); + assert_eq!( + native_field(&value, ZERO_COUNT_I64_FIELD), + Some(ValueData::I64Value(0)) + ); + assert!(i64_list(native_field(&value, POSITIVE_BUCKETS_I64_FIELD)).is_empty()); + let Some(ValueData::F64Value(sum)) = native_field(&value, SUM_FIELD) else { + panic!("expected histogram sum"); + }; + assert_eq!(sum.to_bits(), PROMETHEUS_STALE_NAN_BITS); + } + + #[test] + fn test_non_flag_nan_sum_is_normalized() { + use common_query::native_histogram::SUM_FIELD; + + let point = ExponentialHistogramDataPoint { + sum: Some(f64::from_bits(PROMETHEUS_STALE_NAN_BITS)), + ..Default::default() + }; + let (value, _) = exponential_histogram_value(&point).unwrap(); + let Some(ValueData::F64Value(sum)) = native_field(&value, SUM_FIELD) else { + panic!("expected histogram sum"); + }; + assert_eq!(sum.to_bits(), f64::NAN.to_bits()); + assert_ne!(sum.to_bits(), PROMETHEUS_STALE_NAN_BITS); + } + + #[test] + fn test_exponential_histogram_rejects_invalid_values() { + let mut cases = Vec::new(); + + let mut point = exponential_point(); + point.scale = -5; + cases.push((point, "scale -5 is unsupported")); + + let mut point = exponential_point(); + point.zero_threshold = f64::INFINITY; + cases.push((point, "must be finite and non-negative")); + + let mut point = exponential_point(); + point.start_time_unix_nano = point.time_unix_nano + 1; + cases.push((point, "start_time_unix_nano")); + + let mut point = exponential_point(); + point.count = 27; + cases.push((point, "declared count is 27")); + + let point = ExponentialHistogramDataPoint { + count: u64::MAX, + zero_count: u64::MAX, + ..Default::default() + }; + cases.push((point, "count 18446744073709551615 overflows i64")); + + let point = ExponentialHistogramDataPoint { + count: 1, + scale: 8, + positive: Some(exponential_buckets(i32::MAX, vec![1])), + ..Default::default() + }; + cases.push((point, "shifted bucket index overflows i32")); + + for (point, expected) in cases { + let error = exponential_histogram_value(&point).unwrap_err(); + assert!( + error.contains(expected), + "expected {expected:?}, got {error}" + ); + } + + let buckets = exponential_buckets(0, vec![u64::MAX, 1]); + let error = convert_bucket_range("positive", Some(&buckets), 1).unwrap_err(); + assert!( + error.contains("merged bucket count overflows u64"), + "{error}" + ); + + let buckets = exponential_buckets(i32::MAX, vec![1, 1]); + let error = convert_bucket_range("positive", Some(&buckets), 1).unwrap_err(); + assert!(error.contains("bucket index overflows i32"), "{error}"); + } + + fn metrics_request(metrics: Vec) -> ExportMetricsServiceRequest { + ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics, + ..Default::default() + }], + ..Default::default() + }], + } + } + + fn exponential_metric( + name: impl Into, + data_points: Vec, + temporality: AggregationTemporality, + ) -> Metric { + Metric { + name: name.into(), + data: Some(metric::Data::ExponentialHistogram(ExponentialHistogram { + data_points, + aggregation_temporality: temporality as i32, + })), + ..Default::default() + } + } + + fn histogram_metric(name: impl Into) -> Metric { + Metric { + name: name.into(), + data: Some(metric::Data::Histogram(Histogram { + data_points: vec![HistogramDataPoint { + start_time_unix_nano: 1_000_000, + time_unix_nano: 2_000_000, + count: 1, + sum: Some(1.0), + bucket_counts: vec![1], + ..Default::default() + }], + aggregation_temporality: AggregationTemporality::Cumulative as i32, + })), + ..Default::default() + } + } + + #[test] + fn test_exponential_histogram_gate_and_partial_outcome() { + let request = metrics_request(vec![ + Metric { + name: "temperature".to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint::default()], + })), + ..Default::default() + }, + exponential_metric( + "latency", + vec![exponential_point()], + AggregationTemporality::Cumulative, + ), + ]); + let (requests, _, semantic_index, outcome) = + to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + + assert_eq!(outcome.accepted_data_points, 1); + assert_eq!(outcome.rejected_data_points, 1); + assert!( + outcome + .error_message + .as_deref() + .unwrap() + .contains("otlp.experimental_enable_exponential_histogram") + ); + assert_eq!(requests.inserts.len(), 1); + assert_eq!(requests.inserts[0].table_name, "temperature"); + let semantics = decode(&semantic_index); + assert!(semantics.contains_key("temperature")); + assert!(!semantics.contains_key("latency")); + + let empty = metrics_request(vec![exponential_metric( + "empty", + vec![], + AggregationTemporality::Cumulative, + )]); + let (_, _, _, outcome) = + to_grpc_insert_requests(empty, &mut OtlpMetricCtx::default()).unwrap(); + assert_eq!(outcome.rejected_data_points, 0); + assert_eq!(outcome.error_message, None); + } + + #[test] + fn test_exponential_histogram_cannot_share_table_with_scalar_metric() { + let request = metrics_request(vec![ + Metric { + name: "latency".to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + value: Some(number_data_point::Value::AsDouble(1.0)), + ..Default::default() + }], + })), + ..Default::default() + }, + exponential_metric( + "latency", + vec![exponential_point()], + AggregationTemporality::Cumulative, + ), + ]); + let mut ctx = OtlpMetricCtx { + experimental_enable_exponential_histogram: true, + ..Default::default() + }; + + let error = to_grpc_insert_requests(request, &mut ctx).unwrap_err(); + assert!( + error + .to_string() + .contains("cannot mix native histogram and float sample fields") + ); + } + + #[test] + fn test_histogram_cannot_replace_exponential_histogram_table() { + let request = metrics_request(vec![ + exponential_metric( + "latency_bucket", + vec![exponential_point()], + AggregationTemporality::Cumulative, + ), + histogram_metric("latency"), + ]); + let mut ctx = OtlpMetricCtx { + experimental_enable_exponential_histogram: true, + ..Default::default() + }; + + let error = to_grpc_insert_requests(request, &mut ctx).unwrap_err(); + assert!( + error + .to_string() + .contains("cannot mix native histogram and float sample fields") + ); + } + + #[test] + fn test_histograms_with_same_name_across_resources_are_merged() { + let metric = histogram_metric("latency"); + let request = ExportMetricsServiceRequest { + resource_metrics: ["service-a", "service-b"] + .into_iter() + .map(|service| ResourceMetrics { + resource: Some(Resource { + attributes: vec![keyvalue("service.name", service)], + ..Default::default() + }), + scope_metrics: vec![ScopeMetrics { + metrics: vec![metric.clone()], + ..Default::default() + }], + ..Default::default() + }) + .collect(), + }; + + let (requests, rows, _, outcome) = + to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + + assert_eq!(outcome.accepted_data_points, 2); + assert_eq!(rows, 6); + assert_eq!(requests.inserts.len(), 3); + for request in requests.inserts { + assert_eq!( + request.rows.unwrap().rows.len(), + 2, + "{}", + request.table_name + ); + } + } + + #[test] + fn test_exponential_histogram_rejects_temporality_before_stale_point() { + let stale = ExponentialHistogramDataPoint { + flags: DataPointFlags::NoRecordedValueMask as u32, + ..Default::default() + }; + for temporality in [ + AggregationTemporality::Delta, + AggregationTemporality::Unspecified, + ] { + let request = metrics_request(vec![exponential_metric( + "latency", + vec![stale.clone()], + temporality, + )]); + let mut ctx = OtlpMetricCtx { + experimental_enable_exponential_histogram: true, + ..Default::default() + }; + let (requests, rows, semantic_index, outcome) = + to_grpc_insert_requests(request, &mut ctx).unwrap(); + + assert_eq!(outcome.accepted_data_points, 0); + assert_eq!(outcome.rejected_data_points, 1); + assert_eq!(rows, 0); + assert!(requests.inserts.is_empty()); + assert!(semantic_index.is_empty()); + } + } + + #[test] + fn test_exponential_histogram_legacy_and_new_modes_share_struct() { + use common_query::prelude::greptime_native_histogram; + + let mut point = exponential_point(); + point.sum = Some(42.0); + let request = metrics_request(vec![exponential_metric( + "request.duration", + vec![point], + AggregationTemporality::Cumulative, + )]); + let mut new_ctx = OtlpMetricCtx { + experimental_enable_exponential_histogram: true, + ..Default::default() + }; + let (new_requests, _, _, _) = + to_grpc_insert_requests(request.clone(), &mut new_ctx).unwrap(); + let mut legacy_ctx = OtlpMetricCtx { + experimental_enable_exponential_histogram: true, + is_legacy: true, + ..Default::default() + }; + let (legacy_requests, _, _, _) = to_grpc_insert_requests(request, &mut legacy_ctx).unwrap(); + + let new_insert = &new_requests.inserts[0]; + let legacy_insert = &legacy_requests.inserts[0]; + assert_eq!(new_insert.table_name, "request_duration"); + assert_eq!(legacy_insert.table_name, "request_duration"); + let new_rows = new_insert.rows.as_ref().unwrap(); + let legacy_rows = legacy_insert.rows.as_ref().unwrap(); + let field = greptime_native_histogram(); + let new_histogram = new_rows.rows[0].values[new_rows + .schema + .iter() + .position(|column| column.column_name == field) + .unwrap()] + .clone(); + let legacy_histogram = legacy_rows.rows[0].values[legacy_rows + .schema + .iter() + .position(|column| column.column_name == field) + .unwrap()] + .clone(); + assert_eq!(new_histogram, legacy_histogram); + assert!(matches!( + new_rows.rows[0].values[new_rows + .schema + .iter() + .position(|column| column.column_name == greptime_timestamp()) + .unwrap()] + .value_data, + Some(ValueData::TimestampMillisecondValue(2)) + )); + assert!(matches!( + legacy_rows.rows[0].values[legacy_rows + .schema + .iter() + .position(|column| column.column_name == greptime_timestamp()) + .unwrap()] + .value_data, + Some(ValueData::TimestampNanosecondValue(2_000_000)) + )); + } + + #[test] + fn test_rejection_message_is_bounded() { + let request = metrics_request(vec![exponential_metric( + "x".repeat(1_000), + vec![ExponentialHistogramDataPoint::default()], + AggregationTemporality::Cumulative, + )]); + let (_, _, _, outcome) = + to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + + assert_eq!(outcome.rejected_data_points, 1); + assert!(outcome.error_message.unwrap().len() <= MAX_REJECTION_MESSAGE_BYTES); + } + + #[test] + fn test_rejection_message_reason_is_lazy_after_cap() { + let mut outcome = MetricsIngestOutcome { + error_message: Some("x".repeat(MAX_REJECTION_MESSAGE_BYTES)), + ..Default::default() + }; + let mut reason_built = false; + + reject_data_points(&mut outcome, 1, || { + reason_built = true; + "unused".to_string() + }) + .unwrap(); + + assert_eq!(outcome.rejected_data_points, 1); + assert!(!reason_built); + } } diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 1bcec98d5f5..846ef3ee6be 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -15,19 +15,23 @@ use std::collections::hash_map::Entry; use ahash::{HashMap, HashMapExt, HashSet, HashSetExt}; +#[cfg(test)] +use api::greptime_proto::io::prometheus::write::v2::BucketSpan; +#[cfg(test)] use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; use api::greptime_proto::io::prometheus::write::v2::{ - BucketSpan, Exemplar, Histogram, Metadata, Sample, metadata, + Exemplar, Histogram, Metadata, Sample, metadata, }; #[cfg(test)] use api::greptime_proto::io::prometheus::write::v2::{Request, TimeSeries}; -use api::helper::ColumnDataTypeWrapper; +#[cfg(test)] +use api::v1::ColumnSchema; use api::v1::value::ValueData; -use api::v1::{ - ColumnDataType, ColumnSchema, ListValue, RowInsertRequest, Rows, SemanticType, Value, -}; +use api::v1::{ColumnDataType, RowInsertRequest, Rows, SemanticType, Value}; use bytes::{Buf, Bytes}; use common_grpc::precision::Precision; +use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD; +#[cfg(test)] use common_query::native_histogram::*; use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; use pipeline::{ContextOpt, ContextReq}; @@ -42,6 +46,7 @@ use table::requests::{ }; use crate::error::{self, Result}; +use crate::native_histogram::{encode_native_histogram, native_histogram_column_schema}; use crate::prom_remote_write::row_builder::PromCtx; use crate::prom_remote_write::validation::validate_label_name; use crate::prom_remote_write::{REMOTE_WRITE_V2_VERSION, try_decompress}; @@ -59,8 +64,6 @@ use crate::semantic::{ type PromTags<'a> = Vec<(&'a str, String)>; type ResolvedSeriesLabels<'a> = (PromCtx, String, PromTags<'a>); -const MAX_REMOTE_WRITE_V2_SCHEMA: i32 = 8; -const MAX_REDUCIBLE_REMOTE_WRITE_V2_SCHEMA: i32 = 52; const TIME_SERIES_LABELS_REFS_TAG: u32 = 1; const TIME_SERIES_SAMPLES_TAG: u32 = 2; const TIME_SERIES_HISTOGRAMS_TAG: u32 = 3; @@ -597,6 +600,19 @@ fn write_native_histogram<'a>( histogram: &Histogram, tags: impl Iterator, ) -> Result<()> { + let value = encode_native_histogram(histogram).map_err(|error| { + error::InvalidPromRemoteRequestSnafu { + msg: format!("remote write v2 {error}"), + } + .build() + })?; + let column_schema = native_histogram_column_schema().map_err(|error| { + error::InvalidPromRemoteRequestSnafu { + msg: format!("remote write v2 {error}"), + } + .build() + })?; + // Persist both int and float families into the logical table schema. Only one // family is populated per row; the other is written as NULL so PromQL can // infer the original histogram flavor without a separate type column. @@ -608,8 +624,11 @@ fn write_native_histogram<'a>( Precision::Millisecond, &mut row, )?; - - write_native_histogram_value(table_data, histogram, &mut row)?; + row_writer::write_by_schema( + table_data, + std::iter::once((column_schema, Some(value))), + &mut row, + )?; row_writer::write_tags(table_data, tags, &mut row)?; table_data.add_row(row); @@ -617,512 +636,6 @@ fn write_native_histogram<'a>( Ok(()) } -fn write_native_histogram_value( - table_data: &mut TableData, - histogram: &Histogram, - row: &mut Vec, -) -> Result<()> { - let column_schema = native_histogram_column_schema(); - let value = native_histogram_struct_value(histogram)?; - - row_writer::write_by_schema( - table_data, - std::iter::once((column_schema, Some(value))), - row, - ) -} - -fn native_histogram_column_schema() -> ColumnSchema { - let (datatype, datatype_extension) = - ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone()) - .expect("native histogram type is convertible to protobuf") - .into_parts(); - - ColumnSchema { - column_name: greptime_native_histogram().to_string(), - datatype: datatype as i32, - semantic_type: SemanticType::Field as i32, - datatype_extension, - options: None, - } -} - -fn native_histogram_struct_value(histogram: &Histogram) -> Result { - let uses_float_counts = native_histogram_uses_float_counts(histogram)?; - validate_native_histogram(histogram, uses_float_counts)?; - - let mut items = Vec::with_capacity(NATIVE_HISTOGRAM_FIELD_NAMES.len()); - let positive_span_lengths = i32_span_lengths("positive", &histogram.positive_spans)?; - let negative_span_lengths = i32_span_lengths("negative", &histogram.negative_spans)?; - items.extend([ - pb_value(ValueData::I32Value(histogram.schema)), - pb_value(ValueData::F64Value(histogram.zero_threshold)), - pb_value(ValueData::F64Value(histogram.sum)), - pb_value(ValueData::I32Value(histogram.reset_hint)), - optional_pb_value((histogram.start_timestamp != 0).then_some( - ValueData::TimestampMillisecondValue(histogram.start_timestamp), - )), - f64_list_value(histogram.custom_values.iter().copied()), - i32_list_value(histogram.positive_spans.iter().map(|span| span.offset)), - i32_list_value(positive_span_lengths.iter().copied()), - i32_list_value(histogram.negative_spans.iter().map(|span| span.offset)), - i32_list_value(negative_span_lengths.iter().copied()), - ]); - - if uses_float_counts { - validate_float_native_histogram_counts(histogram)?; - let count = match histogram.count.as_ref() { - Some(Count::CountFloat(count)) => *count, - _ => 0.0, - }; - let zero_count = match histogram.zero_count.as_ref() { - Some(ZeroCount::ZeroCountFloat(zero_count)) => *zero_count, - _ => 0.0, - }; - items.extend([ - null_pb_value(), - null_pb_value(), - i64_list_value(std::iter::empty()), - i64_list_value(std::iter::empty()), - pb_value(ValueData::F64Value(count)), - pb_value(ValueData::F64Value(zero_count)), - f64_list_value(histogram.positive_counts.iter().copied()), - f64_list_value(histogram.negative_counts.iter().copied()), - ]); - } else { - let count = match histogram.count.as_ref() { - Some(Count::CountInt(count)) => *count, - _ => 0, - }; - let zero_count = match histogram.zero_count.as_ref() { - Some(ZeroCount::ZeroCountInt(zero_count)) => *zero_count, - _ => 0, - }; - let positive_buckets = bucket_counts_from_deltas(&histogram.positive_deltas)?; - let negative_buckets = bucket_counts_from_deltas(&histogram.negative_deltas)?; - validate_integer_native_histogram_counts(histogram, &positive_buckets, &negative_buckets)?; - let count = i64::try_from(count) - .ok() - .context(error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram integer count {count} overflows i64" - ), - })?; - let zero_count = i64::try_from(zero_count).ok().context( - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram integer zero_count {zero_count} overflows i64" - ), - }, - )?; - items.extend([ - pb_value(ValueData::I64Value(count)), - pb_value(ValueData::I64Value(zero_count)), - i64_list_value(positive_buckets.iter().copied()), - i64_list_value(negative_buckets.iter().copied()), - null_pb_value(), - null_pb_value(), - f64_list_value(std::iter::empty()), - f64_list_value(std::iter::empty()), - ]); - } - - Ok(ValueData::StructValue(api::v1::StructValue { items })) -} - -fn validate_native_histogram(histogram: &Histogram, uses_float_counts: bool) -> Result<()> { - let exponential_overflow_index = validate_native_histogram_schema(histogram.schema)?; - validate_native_histogram_custom_values(histogram)?; - - if histogram.schema == CUSTOM_BUCKETS_SCHEMA { - ensure!( - histogram.zero_threshold == 0.0 && native_histogram_zero_count_is_zero(histogram), - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 custom native histogram must not use a zero bucket" - .to_string(), - } - ); - ensure!( - histogram.negative_spans.is_empty() - && histogram.negative_deltas.is_empty() - && histogram.negative_counts.is_empty(), - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 custom native histogram must not use negative buckets" - .to_string(), - } - ); - } - - let (positive_buckets, negative_buckets) = if uses_float_counts { - ( - histogram.positive_counts.len(), - histogram.negative_counts.len(), - ) - } else { - ( - histogram.positive_deltas.len(), - histogram.negative_deltas.len(), - ) - }; - let bucket_index_range = if let Some(overflow_index) = exponential_overflow_index { - (i32::MIN, overflow_index) - } else { - ( - 0, - i32::try_from(histogram.custom_values.len()).ok().context( - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 custom native histogram has too many custom_values" - .to_string(), - }, - )?, - ) - }; - validate_native_histogram_spans( - "positive", - &histogram.positive_spans, - positive_buckets, - bucket_index_range, - )?; - validate_native_histogram_spans( - "negative", - &histogram.negative_spans, - negative_buckets, - bucket_index_range, - )?; - - Ok(()) -} - -fn validate_native_histogram_schema(schema: i32) -> Result> { - if schema == CUSTOM_BUCKETS_SCHEMA { - return Ok(None); - } - - if let Some(overflow_index) = exponential_overflow_bucket_index(schema) { - return Ok(Some(overflow_index)); - } - - if (MAX_REMOTE_WRITE_V2_SCHEMA + 1..=MAX_REDUCIBLE_REMOTE_WRITE_V2_SCHEMA).contains(&schema) { - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram schema {schema} must be reduced before ingestion" - ), - } - .fail() - } else { - error::InvalidPromRemoteRequestSnafu { - msg: format!("remote write v2 native histogram schema {schema} is unsupported"), - } - .fail() - } -} - -fn validate_native_histogram_custom_values(histogram: &Histogram) -> Result<()> { - if histogram.schema != CUSTOM_BUCKETS_SCHEMA { - ensure!( - histogram.custom_values.is_empty(), - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 standard native histogram must not use custom_values" - .to_string(), - } - ); - return Ok(()); - } - - for value in &histogram.custom_values { - ensure!( - !value.is_nan() && *value != f64::INFINITY, - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 custom native histogram custom_values must not contain +Inf or NaN" - .to_string(), - } - ); - } - for values in histogram.custom_values.windows(2) { - ensure!( - values[0] < values[1], - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 custom native histogram custom_values must be sorted" - .to_string(), - } - ); - } - - Ok(()) -} - -fn validate_native_histogram_spans( - name: &str, - spans: &[BucketSpan], - bucket_count: usize, - bucket_index_range: (i32, i32), -) -> Result<()> { - let span_len = spans - .iter() - .try_fold(0usize, |sum, span| sum.checked_add(span.length as usize)) - .with_context(|| error::InvalidPromRemoteRequestSnafu { - msg: format!("remote write v2 native histogram {name} spans overflow"), - })?; - ensure!( - span_len == bucket_count, - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram {name} spans describe {span_len} buckets, found {bucket_count}" - ), - } - ); - - let mut current_index = 0i32; - for (span_index, span) in spans.iter().enumerate() { - ensure!( - span.offset >= 0 || (span_index == 0 && bucket_index_range.0 == i32::MIN), - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram {name} span {} has negative offset {}", - span_index + 1, - span.offset - ), - } - ); - current_index = if span_index == 0 { - span.offset - } else { - current_index.checked_add(span.offset).with_context(|| { - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram {name} span index overflows i32" - ), - } - })? - }; - - for _ in 0..span.length { - ensure!( - (bucket_index_range.0..=bucket_index_range.1).contains(¤t_index), - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram {name} bucket index {current_index} is out of range" - ), - } - ); - current_index = - current_index - .checked_add(1) - .context(error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram {name} span index overflows i32" - ), - })?; - } - } - - Ok(()) -} - -fn validate_float_native_histogram_counts(histogram: &Histogram) -> Result<()> { - let count = match histogram.count.as_ref() { - Some(Count::CountFloat(count)) => *count, - _ => 0.0, - }; - ensure!( - count >= 0.0 || count.is_nan(), - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 native histogram float count must not be negative".to_string(), - } - ); - - let zero_count = match histogram.zero_count.as_ref() { - Some(ZeroCount::ZeroCountFloat(zero_count)) => *zero_count, - _ => 0.0, - }; - ensure!( - zero_count >= 0.0 || zero_count.is_nan(), - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 native histogram float zero_count must not be negative" - .to_string(), - } - ); - - for (name, counts) in [ - ("positive", &histogram.positive_counts), - ("negative", &histogram.negative_counts), - ] { - for (index, count) in counts.iter().enumerate() { - ensure!( - *count >= 0.0 || count.is_nan(), - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram {name} bucket {} count must not be negative", - index + 1 - ), - } - ); - } - } - - Ok(()) -} - -fn validate_integer_native_histogram_counts( - histogram: &Histogram, - positive_buckets: &[i64], - negative_buckets: &[i64], -) -> Result<()> { - let count = match histogram.count.as_ref() { - Some(Count::CountInt(count)) => *count, - _ => 0, - }; - let zero_count = match histogram.zero_count.as_ref() { - Some(ZeroCount::ZeroCountInt(zero_count)) => *zero_count, - _ => 0, - }; - let bucket_count = positive_buckets - .iter() - .chain(negative_buckets) - .try_fold(zero_count, |total, bucket| { - total.checked_add(*bucket as u64) - }) - .with_context(|| error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 native histogram bucket total overflows u64".to_string(), - })?; - ensure!( - if histogram.sum.is_nan() { - bucket_count <= count - } else { - bucket_count == count - }, - error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram has {bucket_count} observations in buckets, count is {count}" - ), - } - ); - - Ok(()) -} - -fn native_histogram_zero_count_is_zero(histogram: &Histogram) -> bool { - match histogram.zero_count.as_ref() { - Some(ZeroCount::ZeroCountInt(zero_count)) => *zero_count == 0, - Some(ZeroCount::ZeroCountFloat(zero_count)) => *zero_count == 0.0, - None => true, - } -} - -fn native_histogram_uses_float_counts(histogram: &Histogram) -> Result { - let uses_float_count = matches!(histogram.count, Some(Count::CountFloat(_))) - || matches!(histogram.zero_count, Some(ZeroCount::ZeroCountFloat(_))); - let uses_int_count = matches!(histogram.count, Some(Count::CountInt(_))) - || matches!(histogram.zero_count, Some(ZeroCount::ZeroCountInt(_))); - let uses_float_buckets = - !histogram.positive_counts.is_empty() || !histogram.negative_counts.is_empty(); - let uses_int_buckets = - !histogram.positive_deltas.is_empty() || !histogram.negative_deltas.is_empty(); - - if matches!( - (&histogram.count, &histogram.zero_count), - (Some(Count::CountInt(_)), Some(ZeroCount::ZeroCountFloat(_))) - | (Some(Count::CountFloat(_)), Some(ZeroCount::ZeroCountInt(_))) - ) { - return error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 native histogram count and zero_count must use the same integer or float family".to_string(), - } - .fail(); - } - - ensure!( - !(uses_float_buckets && uses_int_buckets), - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 native histogram bucket counts must use either integer deltas or float counts".to_string(), - } - ); - ensure!( - !(uses_float_count && uses_int_buckets), - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 float native histogram must not use integer bucket deltas" - .to_string(), - } - ); - ensure!( - !(uses_int_count && uses_float_buckets), - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 integer native histogram must not use float bucket counts" - .to_string(), - } - ); - - Ok(uses_float_count || uses_float_buckets) -} - -fn pb_value(value_data: ValueData) -> Value { - optional_pb_value(Some(value_data)) -} - -fn null_pb_value() -> Value { - optional_pb_value(None) -} - -fn optional_pb_value(value_data: Option) -> Value { - Value { value_data } -} - -fn list_value(values: impl IntoIterator) -> Value { - pb_value(ValueData::ListValue(ListValue { - items: values.into_iter().map(pb_value).collect(), - })) -} - -fn i32_list_value(values: impl IntoIterator) -> Value { - list_value(values.into_iter().map(ValueData::I32Value)) -} - -fn i32_span_lengths(name: &str, spans: &[BucketSpan]) -> Result> { - spans - .iter() - .map(|span| { - i32::try_from(span.length) - .ok() - .context(error::InvalidPromRemoteRequestSnafu { - msg: format!( - "remote write v2 native histogram {name} span length {} overflows i32", - span.length - ), - }) - }) - .collect() -} - -fn i64_list_value(values: impl IntoIterator) -> Value { - list_value(values.into_iter().map(ValueData::I64Value)) -} - -fn f64_list_value(values: impl IntoIterator) -> Value { - list_value(values.into_iter().map(ValueData::F64Value)) -} - -fn bucket_counts_from_deltas(deltas: &[i64]) -> Result> { - let mut count = 0_i64; - let mut buckets = Vec::with_capacity(deltas.len()); - - for delta in deltas { - count = - count - .checked_add(*delta) - .with_context(|| error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 native histogram bucket count overflows i64".to_string(), - })?; - ensure!( - count >= 0, - error::InvalidPromRemoteRequestSnafu { - msg: "remote write v2 native histogram bucket count is negative".to_string(), - } - ); - buckets.push(count); - } - - Ok(buckets) -} - fn ensure_no_internal_histogram_labels(tags: &PromTags<'_>) -> Result<()> { // The histogram field column is generated from the protobuf payload. for (name, _) in tags { @@ -2371,6 +1884,25 @@ mod tests { ); } + #[test] + fn test_into_context_req_preserves_exponential_zero_threshold() { + for zero_threshold in [-1.0, f64::NAN] { + let ctx_req = decode_test_request(request_with_histogram(Histogram { + zero_threshold, + ..Default::default() + })) + .unwrap(); + let rows = ctx_req.histograms.all_req().next().unwrap().rows.unwrap(); + let Some(ValueData::F64Value(actual)) = + histogram_field_value(&rows, 0, ZERO_THRESHOLD_FIELD) + else { + panic!("expected zero threshold"); + }; + + assert_eq!(zero_threshold.to_bits(), actual.to_bits()); + } + } + #[test] fn test_into_context_req_rejects_internal_histogram_labels() { let mut request = test_util::request_with_labels_and_samples( diff --git a/src/servers/src/query_handler.rs b/src/servers/src/query_handler.rs index 96f412ac1e1..35ec1cc2ff5 100644 --- a/src/servers/src/query_handler.rs +++ b/src/servers/src/query_handler.rs @@ -70,6 +70,15 @@ pub struct TraceIngestOutcome { pub error_message: Option, } +/// Result of ingesting one OTLP metrics request or Arrow batch. +#[derive(Debug, Default, Clone)] +pub struct MetricsIngestOutcome { + pub write_cost: usize, + pub accepted_data_points: i64, + pub rejected_data_points: i64, + pub error_message: Option, +} + #[async_trait] pub trait InfluxdbLineProtocolHandler { /// A successful request will not return a response. @@ -133,7 +142,7 @@ pub trait OpenTelemetryProtocolHandler: PipelineHandler { &self, request: ExportMetricsServiceRequest, ctx: QueryContextRef, - ) -> Result; + ) -> Result; /// Handling opentelemetry traces request async fn traces( diff --git a/src/session/src/protocol_ctx.rs b/src/session/src/protocol_ctx.rs index 6e84761c2d9..537dec5c414 100644 --- a/src/session/src/protocol_ctx.rs +++ b/src/session/src/protocol_ctx.rs @@ -42,6 +42,7 @@ impl ProtocolCtx { /// If true, all scope attributes will be promoted to the final table schema. /// Along with the scope name, scope version and scope schema URL. /// - `with_metric_engine` +/// - `experimental_enable_exponential_histogram` /// - `is_legacy` /// If the user uses OTLP metrics ingestion before v0.16, it uses the old path. /// So we call this path 'legacy'. @@ -53,6 +54,7 @@ pub struct OtlpMetricCtx { pub resource_attrs: HashSet, pub promote_scope_attrs: bool, pub with_metric_engine: bool, + pub experimental_enable_exponential_histogram: bool, pub is_legacy: bool, pub metric_type: MetricType, pub metric_translation_strategy: OtlpMetricTranslationStrategy, diff --git a/src/standalone/src/options.rs b/src/standalone/src/options.rs index bedf9eebc70..4f31ac2ebdf 100644 --- a/src/standalone/src/options.rs +++ b/src/standalone/src/options.rs @@ -24,7 +24,7 @@ use file_engine::config::EngineConfig as FileEngineConfig; use flow::FlowConfig; use frontend::frontend::FrontendOptions; use frontend::service_config::{ - InfluxdbOptions, JaegerOptions, MysqlOptions, OpentsdbOptions, PostgresOptions, + InfluxdbOptions, JaegerOptions, MysqlOptions, OpentsdbOptions, OtlpOptions, PostgresOptions, PromStoreOptions, }; use mito2::config::MitoConfig; @@ -56,6 +56,7 @@ pub struct StandaloneOptions { pub opentsdb: OpentsdbOptions, pub influxdb: InfluxdbOptions, pub jaeger: JaegerOptions, + pub otlp: OtlpOptions, pub prom_store: PromStoreOptions, pub wal: DatanodeWalConfig, pub storage: StorageConfig, @@ -94,6 +95,7 @@ impl Default for StandaloneOptions { opentsdb: OpentsdbOptions::default(), influxdb: InfluxdbOptions::default(), jaeger: JaegerOptions::default(), + otlp: OtlpOptions::default(), prom_store: PromStoreOptions::default(), wal: DatanodeWalConfig::default(), storage: StorageConfig::default(), @@ -153,6 +155,7 @@ impl StandaloneOptions { opentsdb: cloned_opts.opentsdb, influxdb: cloned_opts.influxdb, jaeger: cloned_opts.jaeger, + otlp: cloned_opts.otlp, prom_store: cloned_opts.prom_store, meta_client: None, logging: cloned_opts.logging, diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 30b9a07348d..cb958a84e5c 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -490,7 +490,7 @@ pub async fn setup_test_http_app_with_frontend_and_slow_query_threshold( .with_log_ingest_handler(instance.fe_instance().clone(), None, None) .with_logs_handler(instance.fe_instance().clone()) .with_influxdb_handler(instance.fe_instance().clone()) - .with_otlp_handler(instance.fe_instance().clone(), true) + .with_otlp_handler(instance.fe_instance().clone(), true, false) .with_jaeger_handler(instance.fe_instance().clone()) .with_greptime_config_options(instance.opts.to_toml().unwrap()) .build(); @@ -510,6 +510,18 @@ pub async fn setup_test_http_app_with_frontend_and_user_provider( user_provider, None, None, + false, + ) + .await +} + +pub async fn setup_test_http_app_with_otlp_exponential_histogram( + store_type: StorageType, + name: &str, + enabled: bool, +) -> (Router, TestGuard) { + setup_test_http_app_with_frontend_and_custom_options( + store_type, name, None, None, None, enabled, ) .await } @@ -520,6 +532,7 @@ pub async fn setup_test_http_app_with_frontend_and_custom_options( user_provider: Option, http_opts: Option, memory_limiter: Option, + experimental_enable_exponential_histogram: bool, ) -> (Router, TestGuard) { let plugins = Plugins::new(); if let Some(user_provider) = user_provider.clone() { @@ -541,7 +554,12 @@ pub async fn setup_test_http_app_with_frontend_and_custom_options( .with_log_ingest_handler(instance.fe_instance().clone(), None, None) .with_logs_handler(instance.fe_instance().clone()) .with_influxdb_handler(instance.fe_instance().clone()) - .with_otlp_handler(instance.fe_instance().clone(), true) + .with_otlp_handler( + instance.fe_instance().clone(), + true, + experimental_enable_exponential_histogram, + ) + .with_prometheus_handler(instance.fe_instance().clone()) .with_jaeger_handler(instance.fe_instance().clone()) .with_dashboard_handler(instance.fe_instance().clone()) .with_greptime_config_options(instance.opts.to_toml().unwrap()); diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index 23862870dba..9d0e13dddc9 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; + use api::v1::alter_table_expr::Kind; use api::v1::promql_request::Promql; use api::v1::value::ValueData; @@ -32,8 +34,20 @@ use common_recordbatch::RecordBatches; use common_runtime::Runtime; use common_runtime::runtime::{BuilderBuild, RuntimeTrait}; use common_test_util::find_workspace_path; -use otel_arrow_rust::proto::opentelemetry::arrow::v1::BatchArrowRecords; +use datatypes::arrow::array::{ + Array, ArrayRef, Float64Array, Int32Array, ListBuilder, StringArray, StructArray, + TimestampNanosecondArray, UInt8Array, UInt16Array, UInt32Array, UInt64Array, UInt64Builder, +}; +use datatypes::arrow::datatypes::{DataType, Field}; +use datatypes::arrow::ipc::writer::StreamWriter; +use datatypes::arrow::record_batch::RecordBatch as ArrowRecordBatch; +use otel_arrow_rust::otlp::metrics::MetricType as ArrowMetricType; use otel_arrow_rust::proto::opentelemetry::arrow::v1::arrow_metrics_service_client::ArrowMetricsServiceClient; +use otel_arrow_rust::proto::opentelemetry::arrow::v1::{ + ArrowPayload, ArrowPayloadType, BatchArrowRecords, StatusCode as ArrowStatusCode, +}; +use otel_arrow_rust::proto::opentelemetry::metrics::v1::AggregationTemporality; +use otel_arrow_rust::schema::consts as arrow_consts; use servers::grpc::GrpcServerConfig; use servers::grpc::builder::GrpcServerBuilder; use servers::http::prometheus::{ @@ -85,6 +99,7 @@ macro_rules! grpc_tests { test_auto_create_table_with_hints, test_auto_create_table_disabled_by_config, test_otel_arrow_auth, + test_otel_arrow_exponential_histogram, test_insert_and_select, test_dbname, test_grpc_message_size_ok, @@ -378,6 +393,188 @@ pub async fn test_otel_arrow_auth(store_type: StorageType) { let _ = fe_grpc_server.shutdown().await; } +// The pinned otel-arrow Producer cannot hash List-typed bucket schemas yet, so +// serialize these test-only record batches directly into the same Arrow stream format. +fn serialize_arrow_record_batch(record_batch: &ArrowRecordBatch) -> Vec { + let mut bytes = Vec::new(); + let mut writer = StreamWriter::try_new(&mut bytes, record_batch.schema_ref()).unwrap(); + writer.write(record_batch).unwrap(); + writer.finish().unwrap(); + drop(writer); + bytes +} + +fn exponential_histogram_arrow_batch(batch_id: i64, scales: &[i32]) -> BatchArrowRecords { + let resource = StructArray::from(vec![( + Arc::new(Field::new(arrow_consts::ID, DataType::UInt16, true)), + Arc::new(UInt16Array::from(vec![0_u16])) as ArrayRef, + )]); + let scope = StructArray::from(vec![( + Arc::new(Field::new(arrow_consts::ID, DataType::UInt16, true)), + Arc::new(UInt16Array::from(vec![0_u16])) as ArrayRef, + )]); + let metrics = ArrowRecordBatch::try_from_iter(vec![ + ( + arrow_consts::ID, + Arc::new(UInt16Array::from(vec![0_u16])) as ArrayRef, + ), + (arrow_consts::RESOURCE, Arc::new(resource) as ArrayRef), + (arrow_consts::SCOPE, Arc::new(scope) as ArrayRef), + ( + arrow_consts::METRIC_TYPE, + Arc::new(UInt8Array::from(vec![ + ArrowMetricType::ExponentialHistogram as u8, + ])) as ArrayRef, + ), + ( + arrow_consts::NAME, + Arc::new(StringArray::from(vec!["otel.arrow.exponential.latency"])) as ArrayRef, + ), + ( + arrow_consts::AGGREGATION_TEMPORALITY, + Arc::new(Int32Array::from(vec![ + AggregationTemporality::Cumulative as i32, + ])) as ArrayRef, + ), + ]) + .unwrap(); + + let point_count = scales.len(); + let mut positive_counts = ListBuilder::new(UInt64Builder::new()); + let mut negative_counts = ListBuilder::new(UInt64Builder::new()); + for _ in scales { + positive_counts.values().append_slice(&[1, 2]); + positive_counts.append(true); + negative_counts.append(true); + } + let positive_counts = positive_counts.finish(); + let negative_counts = negative_counts.finish(); + let positive = StructArray::from(vec![ + ( + Arc::new(Field::new( + arrow_consts::EXP_HISTOGRAM_OFFSET, + DataType::Int32, + true, + )), + Arc::new(Int32Array::from(vec![-1; point_count])) as ArrayRef, + ), + ( + Arc::new(Field::new( + arrow_consts::EXP_HISTOGRAM_BUCKET_COUNTS, + positive_counts.data_type().clone(), + true, + )), + Arc::new(positive_counts) as ArrayRef, + ), + ]); + let negative = StructArray::from(vec![ + ( + Arc::new(Field::new( + arrow_consts::EXP_HISTOGRAM_OFFSET, + DataType::Int32, + true, + )), + Arc::new(Int32Array::from(vec![0; point_count])) as ArrayRef, + ), + ( + Arc::new(Field::new( + arrow_consts::EXP_HISTOGRAM_BUCKET_COUNTS, + negative_counts.data_type().clone(), + true, + )), + Arc::new(negative_counts) as ArrayRef, + ), + ]); + let data_points = ArrowRecordBatch::try_from_iter(vec![ + ( + arrow_consts::ID, + Arc::new(UInt32Array::from_iter_values( + (0..point_count).map(|id| u32::try_from(id).unwrap()), + )) as ArrayRef, + ), + ( + arrow_consts::PARENT_ID, + Arc::new(UInt16Array::from(vec![0_u16; point_count])) as ArrayRef, + ), + ( + arrow_consts::START_TIME_UNIX_NANO, + Arc::new(TimestampNanosecondArray::from(vec![ + 1_000_000_000; + point_count + ])) as ArrayRef, + ), + ( + arrow_consts::TIME_UNIX_NANO, + Arc::new(TimestampNanosecondArray::from(vec![ + 3_000_000_000; + point_count + ])) as ArrayRef, + ), + ( + arrow_consts::HISTOGRAM_COUNT, + Arc::new(UInt64Array::from(vec![4_u64; point_count])) as ArrayRef, + ), + ( + arrow_consts::HISTOGRAM_SUM, + Arc::new(Float64Array::from(vec![8.0; point_count])) as ArrayRef, + ), + ( + arrow_consts::EXP_HISTOGRAM_SCALE, + Arc::new(Int32Array::from(scales.to_vec())) as ArrayRef, + ), + ( + arrow_consts::EXP_HISTOGRAM_ZERO_COUNT, + Arc::new(UInt64Array::from(vec![1_u64; point_count])) as ArrayRef, + ), + ( + arrow_consts::EXP_HISTOGRAM_POSITIVE, + Arc::new(positive) as ArrayRef, + ), + ( + arrow_consts::EXP_HISTOGRAM_NEGATIVE, + Arc::new(negative) as ArrayRef, + ), + ( + arrow_consts::FLAGS, + Arc::new(UInt32Array::from(vec![0_u32; point_count])) as ArrayRef, + ), + ]) + .unwrap(); + BatchArrowRecords { + batch_id, + arrow_payloads: vec![ + ArrowPayload { + schema_id: format!("metrics-{batch_id}"), + r#type: ArrowPayloadType::UnivariateMetrics as i32, + record: serialize_arrow_record_batch(&metrics), + }, + ArrowPayload { + schema_id: format!("exp-histogram-{batch_id}"), + r#type: ArrowPayloadType::ExpHistogramDataPoints as i32, + record: serialize_arrow_record_batch(&data_points), + }, + ], + headers: vec![], + } +} + +pub async fn test_otel_arrow_exponential_histogram(store_type: StorageType) { + let (_instance, server) = + setup_grpc_server(store_type, "test_otel_arrow_exponential_histogram").await; + let addr = server.bind_addr().unwrap().to_string(); + let mut client = ArrowMetricsServiceClient::connect(format!("http://{addr}")) + .await + .unwrap(); + let batch = exponential_histogram_arrow_batch(0, &[0]); + let request = Request::new(futures::stream::once(async { batch })); + let mut response = client.arrow_metrics(request).await.unwrap().into_inner(); + let status = response.message().await.unwrap().unwrap(); + assert_eq!(0, status.batch_id); + assert_eq!(ArrowStatusCode::InvalidArgument as i32, status.status_code); + assert!(status.status_message.contains("omits zero_threshold")); + let _ = server.shutdown().await; +} + fn basic_auth(username: &str, password: &str) -> String { format!("Basic {}", basic_auth_credentials(username, password)) } diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index f89d75e9ef0..2a66c4353e7 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -50,7 +50,9 @@ use log_query::{AggFunc, Context, Limit, LogExpr, LogQuery, TimeFilter}; use loki_proto::logproto::{EntryAdapter, LabelPairAdapter, PushRequest, StreamAdapter}; use loki_proto::prost_types::Timestamp; use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; -use opentelemetry_proto::tonic::collector::metrics::v1::ExportMetricsServiceRequest; +use opentelemetry_proto::tonic::collector::metrics::v1::{ + ExportMetricsServiceRequest, ExportMetricsServiceResponse, +}; use opentelemetry_proto::tonic::collector::trace::v1::{ ExportTraceServiceRequest, ExportTraceServiceResponse, }; @@ -160,6 +162,7 @@ macro_rules! http_tests { test_pipeline_index_options, test_otlp_metrics_new, + test_otlp_exponential_histogram, test_otlp_metric_translation_strategies, test_otlp_traces_v0, test_otlp_traces_v1, @@ -2201,6 +2204,11 @@ default_merge_mode = "last_non_null" [jaeger] enable = true +[otlp] +enable = true +experimental_enable_exponential_histogram = false +trace_ingest_chunk_size = 512 + [prom_store] enable = true with_metric_engine = true @@ -6709,6 +6717,140 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_otlp_exponential_histogram(store_type: StorageType) { + use opentelemetry_proto::tonic::metrics::v1::{ + AggregationTemporality, ExponentialHistogram, ExponentialHistogramDataPoint, Metric, + ResourceMetrics, ScopeMetrics, exponential_histogram_data_point, metric, + }; + use tests_integration::test_util::setup_test_http_app_with_otlp_exponential_histogram; + + common_telemetry::init_default_ut_logging(); + let req = ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + name: "otlp.exponential.latency".to_string(), + data: Some(metric::Data::ExponentialHistogram(ExponentialHistogram { + data_points: vec![ + ExponentialHistogramDataPoint { + start_time_unix_nano: 1_000_000_000, + time_unix_nano: 3_000_000_000, + count: 4, + sum: Some(8.0), + scale: 0, + zero_count: 1, + positive: Some(exponential_histogram_data_point::Buckets { + offset: -1, + bucket_counts: vec![1, 2], + }), + ..Default::default() + }, + ExponentialHistogramDataPoint { + start_time_unix_nano: 1_000_000_000, + time_unix_nano: 4_000_000_000, + scale: -5, + ..Default::default() + }, + ], + aggregation_temporality: AggregationTemporality::Cumulative as i32, + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + }; + let body = req.encode_to_vec(); + let headers = || { + vec![( + HeaderName::from_static("content-type"), + HeaderValue::from_static("application/x-protobuf"), + )] + }; + + let (app, mut guard) = setup_test_http_app_with_otlp_exponential_histogram( + store_type, + "test_otlp_exponential_histogram_disabled", + false, + ) + .await; + let client = TestClient::new(app).await; + let res = send_req( + &client, + headers(), + "/v1/otlp/v1/metrics", + body.clone(), + false, + ) + .await; + assert_eq!(StatusCode::BAD_REQUEST, res.status()); + let status = GoogleRpcStatus::decode(res.bytes().await.as_ref()).unwrap(); + assert_eq!(3, status.code); + assert!( + status + .message + .contains("otlp.experimental_enable_exponential_histogram") + ); + validate_data( + "otlp_exponential_histogram_disabled_no_table", + &client, + "select count(*) from information_schema.tables where table_name = 'otlp_exponential_latency';", + "[[0]]", + ) + .await; + guard.remove_all().await; + + let (app, mut guard) = setup_test_http_app_with_otlp_exponential_histogram( + store_type, + "test_otlp_exponential_histogram_enabled", + true, + ) + .await; + let client = TestClient::new(app).await; + let res = send_req(&client, headers(), "/v1/otlp/v1/metrics", body, false).await; + assert_eq!(StatusCode::OK, res.status()); + let response = ExportMetricsServiceResponse::decode(res.bytes().await).unwrap(); + let partial_success = response.partial_success.unwrap(); + assert_eq!(1, partial_success.rejected_data_points); + assert!(partial_success.error_message.contains("scale -5")); + + validate_data( + "otlp_exponential_histogram_row", + &client, + "select greptime_timestamp, greptime_native_histogram from otlp_exponential_latency;", + "[[3000,{\"count_f64\":null,\"count_i64\":4,\"custom_values\":[],\"negative_buckets_f64\":[],\"negative_buckets_i64\":[],\"negative_span_lengths\":[],\"negative_span_offsets\":[],\"positive_buckets_f64\":[],\"positive_buckets_i64\":[1,2],\"positive_span_lengths\":[2],\"positive_span_offsets\":[0],\"reset_hint\":0,\"schema\":0,\"start_timestamp\":1000,\"sum\":8.0,\"zero_count_f64\":null,\"zero_count_i64\":1,\"zero_threshold\":0.0}]]", + ) + .await; + validate_data( + "otlp_exponential_histogram_metadata", + &client, + "select count(*) from information_schema.tables where table_name = 'otlp_exponential_latency' and create_options like '%greptime.semantic.metric.type=histogram%' and create_options like '%greptime.semantic.metric.temporality=cumulative%';", + "[[1]]", + ) + .await; + + for query in [ + "histogram_count(otlp_exponential_latency)", + "histogram_sum(otlp_exponential_latency)", + "histogram_quantile(0.5,otlp_exponential_latency)", + "histogram_count(rate(otlp_exponential_latency[5s]))", + ] { + let res = client + .get(&format!( + "/v1/prometheus/api/v1/query?query={}&time=3", + encode(query) + )) + .send() + .await; + assert_eq!(StatusCode::OK, res.status(), "query: {query}"); + let body = res.text().await; + assert!(body.contains("otlp_exponential_latency"), "{query}: {body}"); + } + + guard.remove_all().await; +} + pub async fn test_otlp_metric_translation_strategies(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = @@ -10646,6 +10788,7 @@ pub async fn test_http_memory_limit(store_type: StorageType) { None, Some(http_opts), Some(memory_limiter), + false, ) .await; From 0cc83c45708e442d929c22084fc0e6dda9c76f4a Mon Sep 17 00:00:00 2001 From: Yingwen Date: Mon, 24 Aug 2026 13:34:30 +0000 Subject: [PATCH 74/92] feat(cmd): add parquet development tools (#8939) * feat(cmd): add parquet metadata development tool Signed-off-by: evenyag * feat(cmd): add parquet rewrite development tool Signed-off-by: evenyag * feat(cmd): add SST replacement development tool Signed-off-by: evenyag * feat(cmd): support local parquet files in parquetbench Signed-off-by: evenyag * fix(cmd): clean up parquet development tools Signed-off-by: evenyag * refactor(cmd): share datanode tool utilities Signed-off-by: evenyag * fix(cmd): satisfy parquet tool lints Signed-off-by: evenyag * docs: document parquet development tools Signed-off-by: evenyag * refactor(cmd): rename SST replacement module Signed-off-by: evenyag * fix(cmd): validate parquet rewrite options Signed-off-by: evenyag --------- Signed-off-by: evenyag --- docs/parquet-tools.md | 253 +++++++++++ src/cmd/Cargo.toml | 2 +- src/cmd/src/bin/greptime.rs | 6 + src/cmd/src/datanode.rs | 46 ++ src/cmd/src/datanode/objbench.rs | 107 +---- src/cmd/src/datanode/parquet_meta.rs | 340 ++++++++++++++ src/cmd/src/datanode/parquet_rewrite.rs | 579 ++++++++++++++++++++++++ src/cmd/src/datanode/parquetbench.rs | 399 +++++++++++----- src/cmd/src/datanode/scanbench.rs | 60 +-- src/cmd/src/datanode/sst_replace.rs | 361 +++++++++++++++ src/cmd/src/datanode/tool_util.rs | 236 ++++++++++ 11 files changed, 2126 insertions(+), 263 deletions(-) create mode 100644 docs/parquet-tools.md create mode 100644 src/cmd/src/datanode/parquet_meta.rs create mode 100644 src/cmd/src/datanode/parquet_rewrite.rs create mode 100644 src/cmd/src/datanode/sst_replace.rs create mode 100644 src/cmd/src/datanode/tool_util.rs diff --git a/docs/parquet-tools.md b/docs/parquet-tools.md new file mode 100644 index 00000000000..8528f969431 --- /dev/null +++ b/docs/parquet-tools.md @@ -0,0 +1,253 @@ +# Parquet Development Tools + +GreptimeDB provides the following datanode CLI commands for inspecting, +rewriting, benchmarking, and replacing Parquet SST files: + +- `parquet-meta`: inspect a local Parquet file. +- `parquet-rewrite`: rewrite a local Parquet file with different writer + properties. +- `parquetbench`: benchmark reads from a local GreptimeDB SST or an SST in a + configured object store. +- `sst-replace`: replace an existing Mito SST object and update its manifest + metadata. + +These are development and recovery tools, not stable user-facing interfaces. +Always test a workflow on disposable data before using it on an important +region. + +## Build + +`parquet-meta`, `parquet-rewrite`, and `sst-replace` require the `dev-tools` +feature. The same build also includes `parquetbench`: + +```bash +cargo build -p cmd --bin greptime --features dev-tools +``` + +The examples below use: + +```bash +GREPTIME=./target/debug/greptime +``` + +Run `$GREPTIME datanode --help` for the authoritative argument list. + +## Inspect a Parquet file + +`parquet-meta` reads a local Parquet footer and optional page indexes. It reports +file, row-group, and column metadata, including compression, encodings, sizes, +page offsets, statistics presence, and bloom-filter/index offsets. + +Text output: + +```bash +$GREPTIME datanode parquet-meta \ + --input /tmp/source.parquet +``` + +Machine-readable JSON output: + +```bash +$GREPTIME datanode parquet-meta \ + --input /tmp/source.parquet \ + --format json > /tmp/source-meta.json +``` + +`--format` accepts `text` (the default) or `json`. This command can inspect a +general Parquet file; it does not require GreptimeDB SST key-value metadata. + +## Rewrite a Parquet file + +`parquet-rewrite` has two modes. First, dump a TOML properties file inferred +from an existing file: + +```bash +$GREPTIME datanode parquet-rewrite \ + --input /tmp/source.parquet \ + --dump-properties /tmp/writer-properties.toml +``` + +Review and edit the generated file. Its shape is: + +```toml +[writer] +dictionary_enabled = true +compression = "zstd" +max_row_group_row_count = 8192 + +[[columns]] +path = ["host"] +dictionary_enabled = true +compression = "zstd" +encoding = "plain" +``` + +Supported compression names are `uncompressed`, `snappy`, `gzip`, `lzo`, +`brotli`, `lz4`, `zstd`, and `lz4-raw`. Supported encodings are `plain`, +`delta-binary-packed`, `delta-length-byte-array`, `delta-byte-array`, and +`byte-stream-split`. + +The `[writer]` table also accepts `compression_level`, +`data_page_size_limit`, `data_page_row_count_limit`, and +`dictionary_page_size_limit`. A `[[columns]]` entry overrides dictionary, +compression, compression level, or encoding for its column path. Unknown TOML +fields are rejected. Compression levels are supported for `gzip`, `brotli`, and +`zstd`. A column-level `compression_level` without `compression` inherits the +writer compression codec; it is rejected when there is no writer codec to +inherit. + +Rewrite the data using the edited properties: + +```bash +$GREPTIME datanode parquet-rewrite \ + --input /tmp/source.parquet \ + --properties /tmp/writer-properties.toml \ + --output /tmp/rewritten.parquet +``` + +Use `--batch-size ` to control reader batch size. Output and dumped +properties files are not replaced unless `--overwrite` is passed. The output or +dump path must not be exactly the same path as the input, even with +`--overwrite`. This check does not resolve symbolic links, hard links, or other +spellings of the same path. + +The rewrite decodes and writes the Arrow record batches and copies the Parquet +key-value metadata, but it creates a new physical Parquet layout. Inspect and +validate the result before using it as an SST: + +```bash +$GREPTIME datanode parquet-meta \ + --input /tmp/rewritten.parquet +``` + +The dumped properties are inferred primarily from the first row group. Review +them when the source uses different properties across row groups. + +## Benchmark an SST + +`parquetbench` expects GreptimeDB region metadata embedded in the SST. It can +read a local file with the direct reader: + +```bash +$GREPTIME datanode parquetbench \ + --file-path /tmp/rewritten.parquet \ + --reader direct \ + --iterations 5 \ + --batch-size 8192 +``` + +Local-file mode cannot be combined with `--config`, `--region-id`, +`--table-dir`, or `--file-id`, and it does not support the `flat-prune` reader. + +To benchmark an SST in the object store configured for a datanode or standalone +deployment: + +```bash +$GREPTIME datanode parquetbench \ + --config /path/to/datanode.toml \ + --region-id 1024:0 \ + --table-dir data/greptime/public/1024 \ + --file-id 00020380-009c-426d-953e-b4e34c15af34 \ + --path-type bare \ + --reader flat-prune \ + --iterations 5 +``` + +Region mode requires all four of `--config`, `--region-id`, `--table-dir`, and +`--file-id`. `--region-id` accepts either the packed unsigned integer or +`:`. `--path-type` accepts `bare`, `data`, or +`metadata` and defaults to `bare`. + +An optional scan configuration selects columns and row groups: + +```json +{ + "projection_names": ["host", "value", "ts"], + "row_groups": [0, 2] +} +``` + +```bash +$GREPTIME datanode parquetbench \ + --file-path /tmp/rewritten.parquet \ + --scan-config /tmp/parquet-scan.json \ + --iterations 5 +``` + +Use `--pk-as-binary` to expose `__primary_key` as binary with the direct reader. +On Unix, `--pprof-file ` writes a flamegraph. Add +`--pprof-after-warmup` and use at least two iterations to exclude the first +iteration from profiling. + +## Replace an existing region SST + +> **Warning:** `sst-replace` mutates both an SST object and its region manifest. +> Stop the datanode that owns the region and back up the target SST and manifest +> before using `--confirm`. The SST write happens before the manifest update, so +> an interrupted or failed operation may require restoring the backup. + +`sst-replace` replaces the contents of an existing SST file ID. It does not add +a new file ID to a region. The command requires the replacement to have the same +row count and row-group count recorded in the manifest when those manifest +values are nonzero. It does not validate schema or row contents. The replacement +is loaded into memory in full, so ensure the machine has enough memory for the +SST. + +Start with the default dry run using a local replacement file: + +```bash +$GREPTIME datanode sst-replace \ + --config /path/to/datanode.toml \ + --region-id 1024:0 \ + --table-dir data/greptime/public/1024 \ + --file-id 00020380-009c-426d-953e-b4e34c15af34 \ + --replacement-file /tmp/rewritten.parquet +``` + +The dry run locates the manifest and target SST, reads and validates the +replacement footer, and prints the old and new sizes without writing anything. +`--path-type` defaults to `auto`, which probes `bare`, `data`, and `metadata`. +Specify the path type if the file ID is visible in more than one manifest. + +The replacement can instead be read from the configured object store: + +```bash +$GREPTIME datanode sst-replace \ + --config /path/to/datanode.toml \ + --region-id 1024:0 \ + --table-dir data/greptime/public/1024 \ + --file-id 00020380-009c-426d-953e-b4e34c15af34 \ + --replacement-object staging/rewritten.parquet \ + --path-type bare +``` + +After reviewing the dry-run output and confirming that the datanode is stopped, +repeat the exact command with `--confirm`: + +```bash +$GREPTIME datanode sst-replace \ + --config /path/to/datanode.toml \ + --region-id 1024:0 \ + --table-dir data/greptime/public/1024 \ + --file-id 00020380-009c-426d-953e-b4e34c15af34 \ + --replacement-file /tmp/rewritten.parquet \ + --path-type bare \ + --confirm +``` + +The confirmed operation overwrites the existing SST object, recalculates file +size and row-group statistics, and appends a manifest edit for the existing file +ID. Restart the datanode and validate queries against the region before removing +the backup. + +## Recommended rewrite and replacement workflow + +1. Back up the source SST and its region manifest. +2. Inspect the source with `parquet-meta`. +3. Dump and edit writer properties with `parquet-rewrite`. +4. Rewrite to a new local file; never rewrite directly over the source SST. +5. Inspect the output and benchmark it with `parquetbench`. +6. Stop the owning datanode. +7. Run `sst-replace` without `--confirm` and review its resolved target and + statistics. +8. Repeat with `--confirm`, restart the datanode, and validate the region. diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index 10865b48a10..579e5b1454c 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -27,7 +27,7 @@ default = [ "meta-srv/mysql_kvbackend", ] enterprise = ["common-meta/enterprise", "frontend/enterprise", "meta-srv/enterprise"] -# Developer-only helper binary gate for `query_perf_fixture`. +# Developer-only helper binaries and diagnostic datanode commands. # Kept out of `default` so normal/release builds don't compile them. dev-tools = [] mysql-object-store = ["object-store/mysql-object-store"] diff --git a/src/cmd/src/bin/greptime.rs b/src/cmd/src/bin/greptime.rs index 957e227355a..f5f4a590487 100644 --- a/src/cmd/src/bin/greptime.rs +++ b/src/cmd/src/bin/greptime.rs @@ -118,6 +118,12 @@ async fn start(cli: Command) -> Result<()> { datanode::SubCommand::Objbench(ref bench) => bench.run().await, datanode::SubCommand::Scanbench(ref bench) => bench.run().await, datanode::SubCommand::Parquetbench(ref bench) => bench.run().await, + #[cfg(feature = "dev-tools")] + datanode::SubCommand::ParquetMeta(ref cmd) => cmd.run().await, + #[cfg(feature = "dev-tools")] + datanode::SubCommand::ParquetRewrite(ref cmd) => cmd.run().await, + #[cfg(feature = "dev-tools")] + datanode::SubCommand::SstReplace(ref cmd) => cmd.run().await, }, SubCommand::Flownode(cmd) => { cmd.build(cmd.load_options(&cli.global_options)?) diff --git a/src/cmd/src/datanode.rs b/src/cmd/src/datanode.rs index c33be6c6f33..9cb8a8c66c1 100644 --- a/src/cmd/src/datanode.rs +++ b/src/cmd/src/datanode.rs @@ -15,10 +15,20 @@ pub mod builder; #[allow(clippy::print_stdout)] pub(crate) mod objbench; +#[cfg(feature = "dev-tools")] +#[allow(clippy::print_stdout)] +mod parquet_meta; +#[cfg(feature = "dev-tools")] +#[allow(clippy::print_stdout)] +mod parquet_rewrite; #[allow(clippy::print_stdout)] pub mod parquetbench; #[allow(clippy::print_stdout)] pub mod scanbench; +#[cfg(feature = "dev-tools")] +#[allow(clippy::print_stdout)] +mod sst_replace; +mod tool_util; use std::path::Path; use std::time::Duration; @@ -39,8 +49,14 @@ use tracing_appender::non_blocking::WorkerGuard; use crate::App; use crate::datanode::builder::InstanceBuilder; use crate::datanode::objbench::ObjbenchCommand; +#[cfg(feature = "dev-tools")] +use crate::datanode::parquet_meta::ParquetMetaCommand; +#[cfg(feature = "dev-tools")] +use crate::datanode::parquet_rewrite::ParquetRewriteCommand; use crate::datanode::parquetbench::ParquetbenchCommand; use crate::datanode::scanbench::ScanbenchCommand; +#[cfg(feature = "dev-tools")] +use crate::datanode::sst_replace::SstReplaceCommand; use crate::error::{ LoadLayeredConfigSnafu, MissingConfigSnafu, Result, ShutdownDatanodeSnafu, StartDatanodeSnafu, }; @@ -114,6 +130,12 @@ impl Command { // Bench commands are standalone utilities and don't need to load DatanodeOptions. SubCommand::Objbench(_) | SubCommand::Scanbench(_) => Self::default_bench_options(), SubCommand::Parquetbench(_) => Self::default_bench_options(), + #[cfg(feature = "dev-tools")] + SubCommand::ParquetMeta(_) => Self::default_bench_options(), + #[cfg(feature = "dev-tools")] + SubCommand::ParquetRewrite(_) => Self::default_bench_options(), + #[cfg(feature = "dev-tools")] + SubCommand::SstReplace(_) => Self::default_bench_options(), } } @@ -139,6 +161,15 @@ pub enum SubCommand { Scanbench(ScanbenchCommand), /// Benchmark scanning a single parquet SST. Parquetbench(ParquetbenchCommand), + /// Display metadata of a parquet file. + #[cfg(feature = "dev-tools")] + ParquetMeta(ParquetMetaCommand), + /// Rewrite a parquet file with different writer properties. + #[cfg(feature = "dev-tools")] + ParquetRewrite(ParquetRewriteCommand), + /// Replace a Mito region SST and update its manifest metadata. + #[cfg(feature = "dev-tools")] + SstReplace(SstReplaceCommand), } impl SubCommand { @@ -160,6 +191,21 @@ impl SubCommand { cmd.run().await?; std::process::exit(0); } + #[cfg(feature = "dev-tools")] + SubCommand::ParquetMeta(cmd) => { + cmd.run().await?; + std::process::exit(0); + } + #[cfg(feature = "dev-tools")] + SubCommand::ParquetRewrite(cmd) => { + cmd.run().await?; + std::process::exit(0); + } + #[cfg(feature = "dev-tools")] + SubCommand::SstReplace(cmd) => { + cmd.run().await?; + std::process::exit(0); + } } } } diff --git a/src/cmd/src/datanode/objbench.rs b/src/cmd/src/datanode/objbench.rs index f33734687dc..a01cad8468a 100644 --- a/src/cmd/src/datanode/objbench.rs +++ b/src/cmd/src/datanode/objbench.rs @@ -18,8 +18,6 @@ use std::time::Instant; use clap::Parser; use colored::Colorize; -use datanode::config::RegionEngineConfig; -use datanode::store; use futures::stream; use mito2::access_layer::{ AccessLayer, AccessLayerRef, Metrics, OperationType, SstWriteRequest, WriteType, @@ -32,19 +30,20 @@ use mito2::sst::file::{FileHandle, FileMeta}; use mito2::sst::file_purger::{FilePurger, FilePurgerRef}; use mito2::sst::index::intermediate::IntermediateManager; use mito2::sst::index::puffin_manager::PuffinManagerFactory; +use mito2::sst::parquet::WriteOptions; use mito2::sst::parquet::reader::ParquetReaderBuilder; -use mito2::sst::parquet::{PARQUET_METADATA_KEY, WriteOptions}; use mito2::worker::write_cache_from_config; use object_store::ObjectStore; -use parquet::file::metadata::{FooterTail, KeyValue}; +use parquet::file::metadata::FooterTail; use regex::Regex; use snafu::OptionExt; -use store_api::metadata::{RegionMetadata, RegionMetadataRef}; use store_api::path_utils::region_name; use store_api::region_request::PathType; use store_api::storage::FileId; -use crate::datanode::{StorageConfig, StorageConfigWrapper}; +use crate::datanode::tool_util::{ + build_object_store, extract_region_metadata, max_row_group_uncompressed_size, parse_config, +}; use crate::error; /// Object storage benchmark command @@ -67,46 +66,6 @@ pub struct ObjbenchCommand { pub pprof_file: Option, } -pub(crate) fn parse_config( - config_path: &PathBuf, -) -> error::Result<( - StorageConfig, - MitoConfig, - common_wal::config::DatanodeWalConfig, -)> { - let cfg_str = std::fs::read_to_string(config_path).map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("failed to read config {}: {e}", config_path.display()), - } - .build() - })?; - - let store_cfg: StorageConfigWrapper = toml::from_str(&cfg_str).map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("failed to parse config {}: {e}", config_path.display()), - } - .build() - })?; - - let wal_config = store_cfg.wal; - let storage_config = store_cfg.storage; - let mito_engine_config = store_cfg - .region_engine - .into_iter() - .filter_map(|c| { - if let RegionEngineConfig::Mito(mito) = c { - Some(mito) - } else { - None - } - }) - .next() - .with_context(|| error::IllegalConfigSnafu { - msg: format!("Engine config not found in {:?}", config_path), - })?; - Ok((storage_config, mito_engine_config, wal_config)) -} - impl ObjbenchCommand { pub async fn run(&self) -> error::Result<()> { if self.verbose { @@ -154,17 +113,7 @@ impl ObjbenchCommand { let region_meta = extract_region_metadata(&self.source, &parquet_meta)?; let num_rows = parquet_meta.file_metadata().num_rows() as u64; let num_row_groups = parquet_meta.num_row_groups() as u64; - let max_row_group_uncompressed_size: u64 = parquet_meta - .row_groups() - .iter() - .map(|rg| { - rg.columns() - .iter() - .map(|c| c.uncompressed_size() as u64) - .sum::() - }) - .max() - .unwrap_or(0); + let max_row_group_uncompressed_size = max_row_group_uncompressed_size(&parquet_meta); println!( "{} Metadata loaded - rows: {}, size: {} bytes", @@ -485,47 +434,6 @@ fn parse_file_dir_components(path: &str) -> error::Result { }) } -pub(crate) fn extract_region_metadata( - file_path: &str, - meta: &parquet::file::metadata::ParquetMetaData, -) -> error::Result { - let kvs: Option<&Vec> = meta.file_metadata().key_value_metadata(); - let Some(kvs) = kvs else { - return Err(error::IllegalConfigSnafu { - msg: format!("{file_path}: missing parquet key_value metadata"), - } - .build()); - }; - let json = kvs - .iter() - .find(|kv| kv.key == PARQUET_METADATA_KEY) - .and_then(|kv| kv.value.as_ref()) - .ok_or_else(|| { - error::IllegalConfigSnafu { - msg: format!("{file_path}: key {PARQUET_METADATA_KEY} not found or empty"), - } - .build() - })?; - let region: RegionMetadata = RegionMetadata::from_json(json).map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid region metadata json: {e}"), - } - .build() - })?; - Ok(Arc::new(region)) -} - -pub(crate) async fn build_object_store(sc: &StorageConfig) -> error::Result { - store::new_object_store(sc.store.clone(), &sc.data_home) - .await - .map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("Failed to build object store: {e:?}"), - } - .build() - }) -} - async fn build_access_layer_simple( components: &FileDirComponents, object_store: ObjectStore, @@ -666,7 +574,8 @@ mod tests { use common_base::readable_size::ReadableSize; use store_api::region_request::PathType; - use crate::datanode::objbench::{parse_config, parse_file_dir_components}; + use crate::datanode::objbench::parse_file_dir_components; + use crate::datanode::tool_util::parse_config; #[test] fn test_parse_dir() { diff --git a/src/cmd/src/datanode/parquet_meta.rs b/src/cmd/src/datanode/parquet_meta.rs new file mode 100644 index 00000000000..425da45e3c0 --- /dev/null +++ b/src/cmd/src/datanode/parquet_meta.rs @@ -0,0 +1,340 @@ +// Copyright 2023 Greptime Team +// +// 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. + +use std::collections::BTreeMap; +use std::path::{Path, PathBuf}; + +use clap::{Parser, ValueEnum}; +use parquet::file::metadata::ParquetMetaData; +use serde::Serialize; +use snafu::ResultExt; + +use crate::datanode::tool_util::{compression_name, load_local_parquet_metadata}; +use crate::error; + +/// Display parquet file metadata. +#[derive(Debug, Parser)] +pub struct ParquetMetaCommand { + /// Path to input parquet file. + #[clap(long, value_name = "FILE")] + input: PathBuf, + + /// Output format. + #[clap(long, value_enum, default_value = "text")] + format: MetaOutputFormat, +} + +#[derive(Debug, Clone, Copy, ValueEnum)] +enum MetaOutputFormat { + Text, + Json, +} + +#[derive(Debug, Serialize)] +struct FileMetaView { + input: String, + num_rows: i64, + num_row_groups: usize, + num_columns: usize, + key_value_metadata: BTreeMap>, + row_groups: Vec, +} + +#[derive(Debug, Serialize)] +struct RowGroupMetaView { + index: usize, + num_rows: i64, + uncompressed_size: i64, + compressed_size: i64, + compression_ratio: Option, + data_pages: Option, + dictionary_page_bytes: i64, + columns: Vec, +} + +#[derive(Debug, Serialize)] +struct ColumnChunkMetaView { + index: usize, + path: String, + physical_type: String, + encodings: Vec, + compression: String, + num_values: i64, + uncompressed_size: i64, + compressed_size: i64, + compression_ratio: Option, + data_page_offset: i64, + dictionary_page_offset: Option, + dictionary_page_bytes: Option, + data_pages: Option, + has_statistics: bool, + column_index_offset: Option, + column_index_length: Option, + offset_index_offset: Option, + offset_index_length: Option, + bloom_filter_offset: Option, + bloom_filter_length: Option, +} + +impl ParquetMetaCommand { + pub async fn run(&self) -> error::Result<()> { + let metadata = load_local_parquet_metadata(&self.input)?; + let view = build_file_meta_view(&self.input, &metadata); + + match self.format { + MetaOutputFormat::Text => print_meta_text(&view), + MetaOutputFormat::Json => { + let json = serde_json::to_string_pretty(&view).context(error::SerdeJsonSnafu)?; + println!("{json}"); + } + } + + Ok(()) + } +} + +fn build_file_meta_view(path: &Path, metadata: &ParquetMetaData) -> FileMetaView { + let key_value_metadata = metadata + .file_metadata() + .key_value_metadata() + .into_iter() + .flatten() + .map(|kv| (kv.key.clone(), kv.value.clone())) + .collect(); + let offset_index = metadata.offset_index(); + let row_groups = metadata + .row_groups() + .iter() + .enumerate() + .map(|(row_group_idx, row_group)| { + let columns: Vec<_> = row_group + .columns() + .iter() + .enumerate() + .map(|(column_idx, column)| { + let data_pages = offset_index + .and_then(|index| index.get(row_group_idx)) + .and_then(|columns| columns.get(column_idx)) + .map(|index| index.page_locations().len()); + let dictionary_page_bytes = dictionary_page_bytes( + column.dictionary_page_offset(), + column.data_page_offset(), + ); + + ColumnChunkMetaView { + index: column_idx, + path: column.column_path().string(), + physical_type: format!("{:?}", column.column_type()), + encodings: column.encodings().map(|enc| format!("{enc:?}")).collect(), + compression: compression_name(column.compression()).to_string(), + num_values: column.num_values(), + uncompressed_size: column.uncompressed_size(), + compressed_size: column.compressed_size(), + compression_ratio: compression_ratio( + column.uncompressed_size(), + column.compressed_size(), + ), + data_page_offset: column.data_page_offset(), + dictionary_page_offset: column.dictionary_page_offset(), + dictionary_page_bytes, + data_pages, + has_statistics: column.statistics().is_some(), + column_index_offset: column.column_index_offset(), + column_index_length: column.column_index_length(), + offset_index_offset: column.offset_index_offset(), + offset_index_length: column.offset_index_length(), + bloom_filter_offset: column.bloom_filter_offset(), + bloom_filter_length: column.bloom_filter_length(), + } + }) + .collect(); + let data_pages = if columns.iter().all(|column| column.data_pages.is_some()) { + Some( + columns + .iter() + .map(|column| column.data_pages.unwrap_or_default()) + .sum(), + ) + } else { + None + }; + let dictionary_page_bytes = columns + .iter() + .map(|column| column.dictionary_page_bytes.unwrap_or_default()) + .sum(); + + RowGroupMetaView { + index: row_group_idx, + num_rows: row_group.num_rows(), + uncompressed_size: row_group.total_byte_size(), + compressed_size: row_group.compressed_size(), + compression_ratio: compression_ratio( + row_group.total_byte_size(), + row_group.compressed_size(), + ), + data_pages, + dictionary_page_bytes, + columns, + } + }) + .collect(); + + FileMetaView { + input: path.display().to_string(), + num_rows: metadata.file_metadata().num_rows(), + num_row_groups: metadata.num_row_groups(), + num_columns: metadata.file_metadata().schema_descr().num_columns(), + key_value_metadata, + row_groups, + } +} + +fn print_meta_text(view: &FileMetaView) { + println!("file: {}", view.input); + println!("rows: {}", view.num_rows); + println!("row_groups: {}", view.num_row_groups); + println!("columns: {}", view.num_columns); + println!("key_value_metadata: {}", view.key_value_metadata.len()); + for row_group in &view.row_groups { + println!( + "row_group[{}]: rows={}, uncompressed_size={}, compressed_size={}, compression_ratio={}, data_pages={}, dictionary_page_bytes={}", + row_group.index, + row_group.num_rows, + row_group.uncompressed_size, + row_group.compressed_size, + format_ratio(row_group.compression_ratio), + format_optional_usize(row_group.data_pages), + row_group.dictionary_page_bytes, + ); + for column in &row_group.columns { + println!( + " column[{}] path={}: type={}, compression={}, encodings=[{}], values={}, uncompressed_size={}, compressed_size={}, compression_ratio={}, data_pages={}, dictionary_page_offset={}, dictionary_page_bytes={}, data_page_offset={}, statistics={}, column_index={}/{}, offset_index={}/{}, bloom_filter={}/{}", + column.index, + column.path, + column.physical_type, + column.compression, + column.encodings.join(","), + column.num_values, + column.uncompressed_size, + column.compressed_size, + format_ratio(column.compression_ratio), + format_optional_usize(column.data_pages), + format_optional_i64(column.dictionary_page_offset), + column + .dictionary_page_bytes + .map(|value| value.to_string()) + .unwrap_or_else(|| "none".to_string()), + column.data_page_offset, + column.has_statistics, + format_optional_i64(column.column_index_offset), + format_optional_i32(column.column_index_length), + format_optional_i64(column.offset_index_offset), + format_optional_i32(column.offset_index_length), + format_optional_i64(column.bloom_filter_offset), + format_optional_i32(column.bloom_filter_length), + ); + } + } +} + +fn dictionary_page_bytes( + dictionary_page_offset: Option, + data_page_offset: i64, +) -> Option { + dictionary_page_offset.map(|offset| data_page_offset.saturating_sub(offset)) +} + +fn compression_ratio(uncompressed: i64, compressed: i64) -> Option { + (uncompressed > 0).then(|| compressed as f64 / uncompressed as f64) +} + +fn format_ratio(ratio: Option) -> String { + ratio + .map(|value| format!("{value:.4}")) + .unwrap_or_else(|| "unknown".to_string()) +} + +fn format_optional_usize(value: Option) -> String { + value + .map(|value| value.to_string()) + .unwrap_or_else(|| "unknown".to_string()) +} + +fn format_optional_i64(value: Option) -> String { + value + .map(|value| value.to_string()) + .unwrap_or_else(|| "none".to_string()) +} + +fn format_optional_i32(value: Option) -> String { + value + .map(|value| value.to_string()) + .unwrap_or_else(|| "none".to_string()) +} + +#[cfg(test)] +mod tests { + use std::fs::File; + use std::sync::Arc; + + use datatypes::arrow::array::{Int32Array, RecordBatch, StringArray}; + use datatypes::arrow::datatypes::{DataType, Field, Schema}; + use parquet::arrow::ArrowWriter; + use parquet::file::metadata::KeyValue; + use parquet::file::properties::WriterProperties; + use tempfile::tempdir; + + use super::*; + + #[test] + fn test_meta_view_unknown_page_count() { + let dir = tempdir().unwrap(); + let path = dir.path().join("input.parquet"); + write_test_parquet(&path); + + let metadata = load_local_parquet_metadata(&path).unwrap(); + let view = build_file_meta_view(&path, &metadata); + + assert_eq!(view.num_row_groups, 1); + assert!(view.row_groups[0].uncompressed_size > 0); + assert!(view.row_groups[0].compressed_size > 0); + assert!(view.row_groups[0].compression_ratio.is_some()); + assert_eq!(view.row_groups[0].columns.len(), 2); + } + + fn write_test_parquet(path: &Path) { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("host", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4])), + Arc::new(StringArray::from(vec!["a", "a", "b", "b"])), + ], + ) + .unwrap(); + let props = WriterProperties::builder() + .set_key_value_metadata(Some(vec![KeyValue::new( + "greptime:test".to_string(), + "value".to_string(), + )])) + .build(); + let mut writer = + ArrowWriter::try_new(File::create(path).unwrap(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } +} diff --git a/src/cmd/src/datanode/parquet_rewrite.rs b/src/cmd/src/datanode/parquet_rewrite.rs new file mode 100644 index 00000000000..47346aac889 --- /dev/null +++ b/src/cmd/src/datanode/parquet_rewrite.rs @@ -0,0 +1,579 @@ +// Copyright 2023 Greptime Team +// +// 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. + +use std::fs::File; +use std::path::{Path, PathBuf}; + +use clap::Parser; +use datatypes::arrow::record_batch::RecordBatchReader; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use parquet::basic::{BrotliLevel, Compression, Encoding, GzipLevel, ZstdLevel}; +use parquet::file::metadata::ParquetMetaData; +use parquet::file::properties::WriterProperties; +use parquet::file::reader::FileReader; +use parquet::file::serialized_reader::SerializedFileReader; +use parquet::schema::types::ColumnPath; +use serde::{Deserialize, Serialize}; +use snafu::ResultExt; + +use crate::datanode::tool_util::load_local_parquet_metadata; +use crate::error; + +/// Read and rewrite a parquet file with different writer properties. +#[derive(Debug, Parser)] +pub struct ParquetRewriteCommand { + /// Path to input parquet file. + #[clap(long, value_name = "FILE")] + input: PathBuf, + + /// Path to output parquet file in rewrite mode. + #[clap(long, value_name = "FILE")] + output: Option, + + /// Path to writer properties TOML in rewrite mode. + #[clap(long, value_name = "FILE")] + properties: Option, + + /// Dump writer properties TOML inferred from the input parquet file. + #[clap(long, value_name = "FILE")] + dump_properties: Option, + + /// Number of rows per record batch. + #[clap(long)] + batch_size: Option, + + /// Overwrite output files. + #[clap(long, default_value_t = false)] + overwrite: bool, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +#[serde(default, deny_unknown_fields)] +struct RewriteProperties { + writer: WriterConfig, + columns: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)] +#[serde(default, deny_unknown_fields)] +struct WriterConfig { + dictionary_enabled: Option, + compression: Option, + compression_level: Option, + encoding: Option, + max_row_group_row_count: Option, + data_page_size_limit: Option, + data_page_row_count_limit: Option, + dictionary_page_size_limit: Option, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +#[serde(default, deny_unknown_fields)] +struct ColumnConfig { + path: Vec, + dictionary_enabled: Option, + compression: Option, + compression_level: Option, + encoding: Option, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "kebab-case")] +enum CompressionConfig { + Uncompressed, + Snappy, + Gzip, + Lzo, + Brotli, + Lz4, + Zstd, + Lz4Raw, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "kebab-case")] +enum EncodingConfig { + Plain, + DeltaBinaryPacked, + DeltaLengthByteArray, + DeltaByteArray, + ByteStreamSplit, +} + +impl ParquetRewriteCommand { + pub async fn run(&self) -> error::Result<()> { + match (&self.dump_properties, &self.output, &self.properties) { + (Some(path), None, None) => self.dump_properties(path), + (None, Some(output), Some(properties)) => self.rewrite(output, properties), + (Some(_), Some(_), _) | (Some(_), _, Some(_)) => illegal_config( + "use either --dump-properties or rewrite mode, not both".to_string(), + ), + (None, _, _) => illegal_config( + "rewrite mode requires --output and --properties; config dump mode requires --dump-properties".to_string(), + ), + } + } + + fn dump_properties(&self, path: &Path) -> error::Result<()> { + ensure_distinct_paths(&self.input, path)?; + ensure_can_write(path, self.overwrite)?; + let metadata = load_local_parquet_metadata(&self.input)?; + let properties = infer_rewrite_properties(&metadata); + let content = toml::to_string_pretty(&properties).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("failed to serialize writer properties: {e}"), + } + .build() + })?; + std::fs::write(path, content).context(error::FileIoSnafu)?; + println!("Wrote writer properties to {}", path.display()); + Ok(()) + } + + fn rewrite(&self, output: &Path, properties: &Path) -> error::Result<()> { + ensure_distinct_paths(&self.input, output)?; + ensure_can_write(output, self.overwrite)?; + + let props = load_rewrite_properties(properties)?; + let input_reader = SerializedFileReader::new(open_file(&self.input)?) + .map_err(|e| parquet_error("read source parquet metadata", &self.input, e))?; + let key_value_metadata = input_reader + .metadata() + .file_metadata() + .key_value_metadata() + .cloned(); + + let mut reader_builder = ParquetRecordBatchReaderBuilder::try_new(open_file(&self.input)?) + .map_err(|e| parquet_error("open source parquet", &self.input, e))?; + if let Some(batch_size) = self.batch_size { + if batch_size == 0 { + return illegal_config("--batch-size must be greater than 0".to_string()); + } + reader_builder = reader_builder.with_batch_size(batch_size); + } + let reader = reader_builder + .build() + .map_err(|e| parquet_error("build parquet reader", &self.input, e))?; + let schema = reader.schema(); + + let writer_props = build_writer_properties(props, key_value_metadata)?; + let mut writer = ArrowWriter::try_new(create_file(output)?, schema, Some(writer_props)) + .map_err(|e| parquet_error("create parquet writer", output, e))?; + + for batch in reader { + let batch = + batch.map_err(|e| parquet_error("read parquet batch", &self.input, e.into()))?; + writer + .write(&batch) + .map_err(|e| parquet_error("write parquet batch", output, e))?; + } + writer + .close() + .map_err(|e| parquet_error("close parquet writer", output, e))?; + + println!("Wrote parquet file to {}", output.display()); + Ok(()) + } +} + +fn infer_rewrite_properties(metadata: &ParquetMetaData) -> RewriteProperties { + let first_column = metadata + .row_groups() + .first() + .and_then(|row_group| row_group.columns().first()); + let compression = first_column.map(|column| compression_to_config(column.compression())); + let max_row_group_row_count = metadata + .row_groups() + .first() + .and_then(|row_group| usize::try_from(row_group.num_rows()).ok()); + let dictionary_enabled = first_column + .map(|column| column.dictionary_page_offset().is_some()) + .or(Some(true)); + + let columns = metadata + .file_metadata() + .schema_descr() + .columns() + .iter() + .enumerate() + .map(|(idx, column)| { + let first_chunk = metadata + .row_groups() + .first() + .and_then(|row_group| row_group.columns().get(idx)); + ColumnConfig { + path: column.path().parts().to_vec(), + dictionary_enabled: first_chunk + .map(|chunk| chunk.dictionary_page_offset().is_some()), + compression: first_chunk.map(|chunk| compression_to_config(chunk.compression())), + compression_level: None, + encoding: first_chunk.and_then(infer_data_encoding), + } + }) + .collect(); + + RewriteProperties { + writer: WriterConfig { + dictionary_enabled, + compression, + compression_level: None, + encoding: first_column.and_then(infer_data_encoding), + max_row_group_row_count, + data_page_size_limit: None, + data_page_row_count_limit: None, + dictionary_page_size_limit: None, + }, + columns, + } +} + +fn load_rewrite_properties(path: &Path) -> error::Result { + let content = std::fs::read_to_string(path).context(error::FileIoSnafu)?; + toml::from_str(&content).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("failed to parse writer properties {}: {e}", path.display()), + } + .build() + }) +} + +fn build_writer_properties( + config: RewriteProperties, + key_value_metadata: Option>, +) -> error::Result { + let mut builder = WriterProperties::builder().set_key_value_metadata(key_value_metadata); + let writer_compression = config.writer.compression; + if let Some(dictionary_enabled) = config.writer.dictionary_enabled { + builder = builder.set_dictionary_enabled(dictionary_enabled); + } + match (writer_compression, config.writer.compression_level) { + (Some(compression), level) => { + builder = builder.set_compression(to_parquet_compression(compression, level)?); + } + (None, Some(_)) => { + return illegal_config( + "writer compression_level requires writer compression".to_string(), + ); + } + (None, None) => {} + } + if let Some(encoding) = config.writer.encoding { + builder = builder.set_encoding(to_parquet_encoding(encoding)); + } + if let Some(max_row_group_row_count) = config.writer.max_row_group_row_count { + if max_row_group_row_count == 0 { + return illegal_config("max_row_group_row_count must be greater than 0".to_string()); + } + builder = builder.set_max_row_group_row_count(Some(max_row_group_row_count)); + } + if let Some(data_page_size_limit) = config.writer.data_page_size_limit { + builder = builder.set_data_page_size_limit(data_page_size_limit); + } + if let Some(data_page_row_count_limit) = config.writer.data_page_row_count_limit { + builder = builder.set_data_page_row_count_limit(data_page_row_count_limit); + } + if let Some(dictionary_page_size_limit) = config.writer.dictionary_page_size_limit { + builder = builder.set_dictionary_page_size_limit(dictionary_page_size_limit); + } + + for column in config.columns { + if column.path.is_empty() { + return illegal_config("column path must not be empty".to_string()); + } + let path = ColumnPath::new(column.path); + if let Some(dictionary_enabled) = column.dictionary_enabled { + builder = builder.set_column_dictionary_enabled(path.clone(), dictionary_enabled); + } + if column.compression.is_some() || column.compression_level.is_some() { + let Some(compression) = column.compression.or(writer_compression) else { + return illegal_config(format!( + "compression_level for column {} requires column or writer compression", + path.string() + )); + }; + builder = builder.set_column_compression( + path.clone(), + to_parquet_compression(compression, column.compression_level)?, + ); + } + if let Some(encoding) = column.encoding { + builder = builder.set_column_encoding(path, to_parquet_encoding(encoding)); + } + } + + Ok(builder.build()) +} + +fn ensure_distinct_paths(input: &Path, output: &Path) -> error::Result<()> { + if input == output { + return illegal_config(format!( + "input and output paths must be different: {}", + input.display() + )); + } + Ok(()) +} + +fn ensure_can_write(path: &Path, overwrite: bool) -> error::Result<()> { + if !overwrite && path.exists() { + return illegal_config(format!( + "{} already exists; pass --overwrite to replace it", + path.display() + )); + } + Ok(()) +} + +fn open_file(path: &Path) -> error::Result { + File::open(path).context(error::FileIoSnafu) +} + +fn create_file(path: &Path) -> error::Result { + File::create(path).context(error::FileIoSnafu) +} + +fn illegal_config(msg: String) -> error::Result { + error::IllegalConfigSnafu { msg }.fail() +} + +fn parquet_error( + action: &'static str, + path: &Path, + error: parquet::errors::ParquetError, +) -> error::Error { + error::IllegalConfigSnafu { + msg: format!("{action} failed for {}: {error}", path.display()), + } + .build() +} + +fn to_parquet_compression( + compression: CompressionConfig, + level: Option, +) -> error::Result { + if level.is_some() + && !matches!( + compression, + CompressionConfig::Gzip | CompressionConfig::Brotli | CompressionConfig::Zstd + ) + { + return illegal_config(format!( + "compression level is not supported for {}", + compression.name() + )); + } + + Ok(match compression { + CompressionConfig::Uncompressed => Compression::UNCOMPRESSED, + CompressionConfig::Snappy => Compression::SNAPPY, + CompressionConfig::Gzip => Compression::GZIP(match level { + Some(level) => GzipLevel::try_new(level).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid gzip compression level {level}: {e}"), + } + .build() + })?, + None => GzipLevel::default(), + }), + CompressionConfig::Lzo => Compression::LZO, + CompressionConfig::Brotli => Compression::BROTLI(match level { + Some(level) => BrotliLevel::try_new(level).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid brotli compression level {level}: {e}"), + } + .build() + })?, + None => BrotliLevel::default(), + }), + CompressionConfig::Lz4 => Compression::LZ4, + CompressionConfig::Zstd => Compression::ZSTD(match level { + Some(level) => { + let converted_level = i32::try_from(level).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid zstd compression level {level}: {e}"), + } + .build() + })?; + ZstdLevel::try_new(converted_level).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid zstd compression level {level}: {e}"), + } + .build() + })? + } + None => ZstdLevel::default(), + }), + CompressionConfig::Lz4Raw => Compression::LZ4_RAW, + }) +} + +impl CompressionConfig { + fn name(self) -> &'static str { + match self { + CompressionConfig::Uncompressed => "uncompressed", + CompressionConfig::Snappy => "snappy", + CompressionConfig::Gzip => "gzip", + CompressionConfig::Lzo => "lzo", + CompressionConfig::Brotli => "brotli", + CompressionConfig::Lz4 => "lz4", + CompressionConfig::Zstd => "zstd", + CompressionConfig::Lz4Raw => "lz4-raw", + } + } +} + +fn to_parquet_encoding(encoding: EncodingConfig) -> Encoding { + match encoding { + EncodingConfig::Plain => Encoding::PLAIN, + EncodingConfig::DeltaBinaryPacked => Encoding::DELTA_BINARY_PACKED, + EncodingConfig::DeltaLengthByteArray => Encoding::DELTA_LENGTH_BYTE_ARRAY, + EncodingConfig::DeltaByteArray => Encoding::DELTA_BYTE_ARRAY, + EncodingConfig::ByteStreamSplit => Encoding::BYTE_STREAM_SPLIT, + } +} + +fn infer_data_encoding( + column: &parquet::file::metadata::ColumnChunkMetaData, +) -> Option { + let encodings: Vec<_> = column.encodings().collect(); + if encodings.contains(&Encoding::DELTA_BINARY_PACKED) { + Some(EncodingConfig::DeltaBinaryPacked) + } else if encodings.contains(&Encoding::DELTA_LENGTH_BYTE_ARRAY) { + Some(EncodingConfig::DeltaLengthByteArray) + } else if encodings.contains(&Encoding::DELTA_BYTE_ARRAY) { + Some(EncodingConfig::DeltaByteArray) + } else if encodings.contains(&Encoding::BYTE_STREAM_SPLIT) { + Some(EncodingConfig::ByteStreamSplit) + } else if encodings.contains(&Encoding::PLAIN) { + Some(EncodingConfig::Plain) + } else { + None + } +} + +fn compression_to_config(compression: Compression) -> CompressionConfig { + match compression { + Compression::UNCOMPRESSED => CompressionConfig::Uncompressed, + Compression::SNAPPY => CompressionConfig::Snappy, + Compression::GZIP(_) => CompressionConfig::Gzip, + Compression::LZO => CompressionConfig::Lzo, + Compression::BROTLI(_) => CompressionConfig::Brotli, + Compression::LZ4 => CompressionConfig::Lz4, + Compression::ZSTD(_) => CompressionConfig::Zstd, + Compression::LZ4_RAW => CompressionConfig::Lz4Raw, + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datatypes::arrow::array::{Int32Array, RecordBatch, StringArray}; + use datatypes::arrow::datatypes::{DataType, Field, Schema}; + use parquet::file::metadata::KeyValue; + use tempfile::tempdir; + + use super::*; + + #[test] + fn test_dump_and_rewrite_preserves_key_value_metadata_and_disables_dictionary() { + let dir = tempdir().unwrap(); + let input = dir.path().join("input.parquet"); + let output = dir.path().join("output.parquet"); + write_test_parquet(&input, true); + + let metadata = load_local_parquet_metadata(&input).unwrap(); + let mut properties = infer_rewrite_properties(&metadata); + properties.columns[1].dictionary_enabled = Some(false); + let props_path = dir.path().join("props.toml"); + std::fs::write(&props_path, toml::to_string(&properties).unwrap()).unwrap(); + + let command = ParquetRewriteCommand { + input: input.clone(), + output: Some(output.clone()), + properties: Some(props_path), + dump_properties: None, + batch_size: Some(2), + overwrite: false, + }; + command + .rewrite(&output, command.properties.as_ref().unwrap()) + .unwrap(); + + let rewritten = load_local_parquet_metadata(&output).unwrap(); + assert_eq!(rewritten.file_metadata().num_rows(), 4); + let key_values = rewritten + .file_metadata() + .key_value_metadata() + .cloned() + .unwrap_or_default(); + assert!(key_values.iter().any(|kv| kv.key == "greptime:test")); + assert!( + rewritten + .row_groups() + .iter() + .all(|row_group| row_group.column(1).dictionary_page_offset().is_none()) + ); + } + + #[test] + fn test_column_compression_level_inherits_writer_compression() { + let path = ColumnPath::new(vec!["host".to_string()]); + let config = RewriteProperties { + writer: WriterConfig { + compression: Some(CompressionConfig::Zstd), + compression_level: Some(1), + ..Default::default() + }, + columns: vec![ColumnConfig { + path: vec!["host".to_string()], + compression_level: Some(3), + ..Default::default() + }], + }; + + let properties = build_writer_properties(config, None).unwrap(); + assert_eq!( + properties.compression(&path), + Compression::ZSTD(ZstdLevel::try_new(3).unwrap()) + ); + } + + fn write_test_parquet(path: &Path, dictionary_enabled: bool) { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("host", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4])), + Arc::new(StringArray::from(vec!["a", "a", "b", "b"])), + ], + ) + .unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(dictionary_enabled) + .set_key_value_metadata(Some(vec![KeyValue::new( + "greptime:test".to_string(), + "value".to_string(), + )])) + .build(); + let mut writer = + ArrowWriter::try_new(File::create(path).unwrap(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } +} diff --git a/src/cmd/src/datanode/parquetbench.rs b/src/cmd/src/datanode/parquetbench.rs index 2b46a775f76..09ec50f48ef 100644 --- a/src/cmd/src/datanode/parquetbench.rs +++ b/src/cmd/src/datanode/parquetbench.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -42,7 +42,10 @@ use store_api::region_request::PathType; use store_api::storage::consts::{PRIMARY_KEY_COLUMN_NAME, is_internal_column}; use store_api::storage::{ColumnId, FileId, RegionId}; -use crate::datanode::objbench::{build_object_store, extract_region_metadata, parse_config}; +use crate::datanode::tool_util::{ + build_object_store, extract_region_metadata, format_bytes, parse_config, parse_file_id, + parse_path_type, parse_region_id, +}; use crate::error; const DEFAULT_READ_BATCH_SIZE: usize = 8 * 1024; @@ -52,19 +55,23 @@ const DEFAULT_READ_BATCH_SIZE: usize = 8 * 1024; pub struct ParquetbenchCommand { /// Path to config TOML file (same format as standalone/datanode config) #[clap(long, value_name = "FILE")] - config: PathBuf, + config: Option, /// Region ID: either numeric u64 (e.g. "4398046511104") or "table_id:region_num" (e.g. "1024:0") #[clap(long)] - region_id: String, + region_id: Option, /// Table directory relative to data home (e.g. "data/greptime/public/1024/") #[clap(long)] - table_dir: String, + table_dir: Option, /// SST file id to benchmark. #[clap(long)] - file_id: String, + file_id: Option, + + /// Local parquet SST file to benchmark with the direct reader. + #[clap(long, value_name = "FILE")] + file_path: Option, /// Path to scan request JSON config file (supports projection_names only) #[clap(long, value_name = "FILE")] @@ -129,6 +136,32 @@ struct IterationStats { elapsed: Duration, } +#[derive(Debug)] +enum ParquetbenchInput { + LocalFile { + file_path: PathBuf, + }, + Region { + config: PathBuf, + region_id: String, + table_dir: String, + file_id: String, + path_type: PathType, + }, +} + +struct ParquetbenchSource { + object_store: object_store::ObjectStore, + file_path: String, + display_path: String, + region_id_label: String, + region_id: RegionId, + file_id: FileId, + region_file_id: RegionFileId, + path_type: Option, + table_dir: Option, +} + impl ParquetbenchCommand { pub async fn run(&self) -> error::Result<()> { if self.verbose { @@ -145,41 +178,40 @@ impl ParquetbenchCommand { .fail(); } - let region_id = parse_region_id(&self.region_id)?; - let path_type = parse_path_type(&self.path_type)?; - let file_id = FileId::parse_str(&self.file_id).map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid file_id '{}': {}", self.file_id, e), - } - .build() - })?; - let region_file_id = RegionFileId::new(region_id, file_id); - let file_path = sst_file_path(&self.table_dir, region_file_id, path_type); + let input = self.resolve_input()?; + let mut source = build_source(input).await?; - let (store_cfg, _mito_config, _wal_config) = parse_config(&self.config)?; - let object_store = build_object_store(&store_cfg).await?; - - let file_size = object_store - .stat(&file_path) + let file_size = source + .object_store + .stat(&source.file_path) .await .map_err(|e| { error::IllegalConfigSnafu { - msg: format!("stat failed for {}: {}", file_path, e), + msg: format!("stat failed for {}: {}", source.display_path, e), } .build() })? .content_length(); let mut metadata_metrics = MetadataCacheMetrics::default(); - let parquet_meta = MetadataLoader::new(object_store.clone(), &file_path, file_size) - .load(&mut metadata_metrics) - .await - .map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("read parquet metadata failed for {}: {:?}", file_path, e), - } - .build() - })?; - let region_meta = extract_region_metadata(&file_path, &parquet_meta)?; + let parquet_meta = + MetadataLoader::new(source.object_store.clone(), &source.file_path, file_size) + .load(&mut metadata_metrics) + .await + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!( + "read parquet metadata failed for {}: {:?}", + source.display_path, e + ), + } + .build() + })?; + let region_meta = extract_region_metadata(&source.display_path, &parquet_meta)?; + if source.table_dir.is_none() { + source.region_id = region_meta.region_id; + source.region_id_label = source.region_id.as_u64().to_string(); + source.region_file_id = RegionFileId::new(source.region_id, source.file_id); + } let scan_config = self.load_scan_config().await?; let projection = if self.reader == ReaderMode::Direct { resolve_projection_names(&scan_config, ®ion_meta)? @@ -216,10 +248,10 @@ impl ParquetbenchCommand { println!( "{} Region ID: {} (u64: {})", "✓".green(), - self.region_id, - region_id.as_u64() + source.region_id_label, + source.region_id.as_u64() ); - println!("{} File path: {}", "✓".green(), file_path.cyan()); + println!("{} File path: {}", "✓".green(), source.display_path.cyan()); println!( "{} Columns: {}", "✓".green(), @@ -307,8 +339,8 @@ impl ParquetbenchCommand { let mut schema_printed = false; let file_handle = FileHandle::new( FileMeta { - region_id, - file_id, + region_id: source.region_id, + file_id: source.file_id, time_range: Default::default(), level: 0, file_size, @@ -332,9 +364,9 @@ impl ParquetbenchCommand { let stats = match self.reader { ReaderMode::Direct => { run_direct_iteration( - object_store.clone(), - file_path.clone(), - region_file_id, + source.object_store.clone(), + source.file_path.clone(), + source.region_file_id, parquet_meta.clone(), projection.clone(), row_groups.clone(), @@ -345,9 +377,19 @@ impl ParquetbenchCommand { } ReaderMode::FlatPrune => { run_flat_prune_iteration( - object_store.clone(), - self.table_dir.clone(), - path_type, + source.object_store.clone(), + source.table_dir.clone().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: "flat-prune reader requires --table-dir".to_string(), + } + .build() + })?, + source.path_type.ok_or_else(|| { + error::IllegalConfigSnafu { + msg: "flat-prune reader requires --path-type".to_string(), + } + .build() + })?, file_handle.clone(), region_meta.clone(), projection_column_ids.clone(), @@ -450,6 +492,65 @@ impl ParquetbenchCommand { Ok(()) } + fn resolve_input(&self) -> error::Result { + let has_region_args = self.config.is_some() + || self.region_id.is_some() + || self.table_dir.is_some() + || self.file_id.is_some(); + + if let Some(file_path) = &self.file_path { + if self.reader == ReaderMode::FlatPrune { + return Err(error::IllegalConfigSnafu { + msg: "--file-path currently supports only --reader direct".to_string(), + } + .build()); + } + if has_region_args { + return Err(error::IllegalConfigSnafu { + msg: "--file-path cannot be used with --config, --region-id, --table-dir, or --file-id".to_string(), + } + .build()); + } + return Ok(ParquetbenchInput::LocalFile { + file_path: file_path.clone(), + }); + } + + let config = self.config.clone().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: "missing --config unless --file-path is specified".to_string(), + } + .build() + })?; + let region_id = self.region_id.clone().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: "missing --region-id unless --file-path is specified".to_string(), + } + .build() + })?; + let table_dir = self.table_dir.clone().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: "missing --table-dir unless --file-path is specified".to_string(), + } + .build() + })?; + let file_id = self.file_id.clone().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: "missing --file-id unless --file-path is specified".to_string(), + } + .build() + })?; + let path_type = parse_path_type(&self.path_type)?; + + Ok(ParquetbenchInput::Region { + config, + region_id, + table_dir, + file_id, + path_type, + }) + } + async fn load_scan_config(&self) -> error::Result { if let Some(path) = &self.scan_config { let content = tokio::fs::read_to_string(path) @@ -462,6 +563,117 @@ impl ParquetbenchCommand { } } +async fn build_source(input: ParquetbenchInput) -> error::Result { + match input { + ParquetbenchInput::LocalFile { file_path } => build_local_file_source(&file_path), + ParquetbenchInput::Region { + config, + region_id, + table_dir, + file_id, + path_type, + } => build_region_source(&config, ®ion_id, table_dir, &file_id, path_type).await, + } +} + +fn build_local_file_source(file_path: &Path) -> error::Result { + let file_path = std::fs::canonicalize(file_path).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid --file-path {}: {e}", file_path.display()), + } + .build() + })?; + if !file_path.is_file() { + return Err(error::IllegalConfigSnafu { + msg: format!("--file-path {} is not a file", file_path.display()), + } + .build()); + } + let parent = file_path.parent().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: format!( + "--file-path {} has no parent directory", + file_path.display() + ), + } + .build() + })?; + let file_name = file_path + .file_name() + .and_then(|name| name.to_str()) + .ok_or_else(|| { + error::IllegalConfigSnafu { + msg: format!("invalid UTF-8 file name in {}", file_path.display()), + } + .build() + })? + .to_string(); + let object_store = object_store::ObjectStore::new(object_store::services::Fs::default().root( + parent.to_str().ok_or_else(|| { + error::IllegalConfigSnafu { + msg: format!("invalid UTF-8 parent directory in {}", file_path.display()), + } + .build() + })?, + )) + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("failed to build local file object store: {e:?}"), + } + .build() + })? + .finish(); + + let file_id = file_path + .file_stem() + .and_then(|stem| stem.to_str()) + .and_then(|stem| FileId::parse_str(stem).ok()) + .unwrap_or_else(FileId::random); + let region_id = RegionId::new(0, 0); + let region_file_id = RegionFileId::new(region_id, file_id); + let display_path = file_path.display().to_string(); + + Ok(ParquetbenchSource { + object_store, + file_path: file_name, + display_path, + region_id_label: region_id.as_u64().to_string(), + region_id, + file_id, + region_file_id, + path_type: None, + table_dir: None, + }) +} + +async fn build_region_source( + config: &Path, + region_id: &str, + table_dir: String, + file_id: &str, + path_type: PathType, +) -> error::Result { + let region = parse_region_id(region_id)?; + let file_id = parse_file_id(file_id)?; + let region_file_id = RegionFileId::new(region, file_id); + let file_path = sst_file_path(&table_dir, region_file_id, path_type); + + let (store_cfg, _mito_config, _wal_config) = parse_config(config)?; + let object_store = build_object_store(&store_cfg).await?; + + Ok(ParquetbenchSource { + object_store, + display_path: file_path.clone(), + file_path, + region_id_label: region_id.to_string(), + region_id: region, + file_id, + region_file_id, + path_type: Some(path_type), + table_dir: Some(table_dir), + }) +} + #[allow(clippy::too_many_arguments)] async fn run_direct_iteration( object_store: object_store::ObjectStore, @@ -773,60 +985,6 @@ fn parse_batch_size(s: &str) -> Result { Ok(batch_size) } -fn parse_region_id(s: &str) -> error::Result { - if s.contains(':') { - let parts: Vec<&str> = s.splitn(2, ':').collect(); - let table_id: u32 = parts[0].parse().map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid table_id in region_id '{}': {}", s, e), - } - .build() - })?; - let region_num: u32 = parts[1].parse().map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid region_num in region_id '{}': {}", s, e), - } - .build() - })?; - Ok(RegionId::new(table_id, region_num)) - } else { - let id: u64 = s.parse().map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid region_id '{}': {}", s, e), - } - .build() - })?; - Ok(RegionId::from_u64(id)) - } -} - -fn parse_path_type(s: &str) -> error::Result { - match s.to_lowercase().as_str() { - "bare" => Ok(PathType::Bare), - "data" => Ok(PathType::Data), - "metadata" => Ok(PathType::Metadata), - _ => Err(error::IllegalConfigSnafu { - msg: format!("invalid path_type '{}', expected: bare, data, metadata", s), - } - .build()), - } -} - -fn format_bytes(bytes: u64) -> String { - const KIB: u64 = 1024; - const MIB: u64 = 1024 * KIB; - const GIB: u64 = 1024 * MIB; - if bytes >= GIB { - format!("{:.2} GiB", bytes as f64 / GIB as f64) - } else if bytes >= MIB { - format!("{:.2} MiB", bytes as f64 / MIB as f64) - } else if bytes >= KIB { - format!("{:.2} KiB", bytes as f64 / KIB as f64) - } else { - format!("{} B", bytes) - } -} - fn format_rate(rate: f64) -> String { if !rate.is_finite() { return "inf rows".to_string(); @@ -851,6 +1009,25 @@ mod tests { use super::*; + fn test_command() -> ParquetbenchCommand { + ParquetbenchCommand { + config: None, + region_id: None, + table_dir: None, + file_id: None, + file_path: None, + scan_config: None, + iterations: 1, + batch_size: DEFAULT_READ_BATCH_SIZE, + path_type: "bare".to_string(), + verbose: false, + pprof_file: None, + pprof_after_warmup: false, + pk_as_binary: false, + reader: ReaderMode::Direct, + } + } + fn new_test_metadata() -> RegionMetadata { let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 0)); builder @@ -882,19 +1059,27 @@ mod tests { } #[test] - fn test_parse_region_id() { - assert_eq!(parse_region_id("1024:7").unwrap(), RegionId::new(1024, 7)); - assert_eq!( - parse_region_id(&RegionId::new(1, 2).as_u64().to_string()).unwrap(), - RegionId::new(1, 2) - ); + fn test_resolve_input_accepts_direct_file_for_direct_reader() { + let mut command = test_command(); + command.file_path = Some(PathBuf::from("/tmp/source.parquet")); + + let input = command.resolve_input().unwrap(); + match input { + ParquetbenchInput::LocalFile { file_path } => { + assert_eq!(file_path, PathBuf::from("/tmp/source.parquet")); + } + ParquetbenchInput::Region { .. } => panic!("expected local file input"), + } } #[test] - fn test_parse_path_type() { - assert_eq!(parse_path_type("bare").unwrap(), PathType::Bare); - assert_eq!(parse_path_type("data").unwrap(), PathType::Data); - assert_eq!(parse_path_type("metadata").unwrap(), PathType::Metadata); + fn test_resolve_input_rejects_mixed_input_modes() { + let mut command = test_command(); + command.file_path = Some(PathBuf::from("/tmp/source.parquet")); + command.config = Some(PathBuf::from("config.toml")); + + let err = command.resolve_input().unwrap_err(); + assert!(err.to_string().contains("--file-path cannot be used with")); } #[test] diff --git a/src/cmd/src/datanode/scanbench.rs b/src/cmd/src/datanode/scanbench.rs index 4909467f481..1dbee025a7f 100644 --- a/src/cmd/src/datanode/scanbench.rs +++ b/src/cmd/src/datanode/scanbench.rs @@ -52,11 +52,13 @@ use sqlparser::parser::Parser as SqlParser; use store_api::metadata::RegionMetadata; use store_api::path_utils::WAL_DIR; use store_api::region_engine::{PrepareRequest, QueryScanContext, RegionEngine}; -use store_api::region_request::{PathType, RegionOpenRequest, RegionRequest}; +use store_api::region_request::{RegionOpenRequest, RegionRequest}; use store_api::storage::{RegionId, ScanRequest, TimeSeriesDistribution, TimeSeriesRowSelector}; use tokio::fs; -use crate::datanode::objbench::{build_object_store, parse_config}; +use crate::datanode::tool_util::{ + build_object_store, format_bytes, parse_config, parse_path_type, parse_region_id, +}; use crate::error; /// Scan benchmark command - benchmarks scanning a region directly from storage. @@ -164,60 +166,6 @@ fn resolve_projection( Ok(None) } -fn format_bytes(bytes: u64) -> String { - const KIB: u64 = 1024; - const MIB: u64 = 1024 * KIB; - const GIB: u64 = 1024 * MIB; - if bytes >= GIB { - format!("{:.2} GiB", bytes as f64 / GIB as f64) - } else if bytes >= MIB { - format!("{:.2} MiB", bytes as f64 / MIB as f64) - } else if bytes >= KIB { - format!("{:.2} KiB", bytes as f64 / KIB as f64) - } else { - format!("{} B", bytes) - } -} - -fn parse_region_id(s: &str) -> error::Result { - if s.contains(':') { - let parts: Vec<&str> = s.splitn(2, ':').collect(); - let table_id: u32 = parts[0].parse().map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid table_id in region_id '{}': {}", s, e), - } - .build() - })?; - let region_num: u32 = parts[1].parse().map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid region_num in region_id '{}': {}", s, e), - } - .build() - })?; - Ok(RegionId::new(table_id, region_num)) - } else { - let id: u64 = s.parse().map_err(|e| { - error::IllegalConfigSnafu { - msg: format!("invalid region_id '{}': {}", s, e), - } - .build() - })?; - Ok(RegionId::from_u64(id)) - } -} - -fn parse_path_type(s: &str) -> error::Result { - match s.to_lowercase().as_str() { - "bare" => Ok(PathType::Bare), - "data" => Ok(PathType::Data), - "metadata" => Ok(PathType::Metadata), - _ => Err(error::IllegalConfigSnafu { - msg: format!("invalid path_type '{}', expected: bare, data, metadata", s), - } - .build()), - } -} - /// Rewrites literal values in comparison expressions to match the column's arrow type. struct LiteralTypeCaster { schema: DFSchemaRef, diff --git a/src/cmd/src/datanode/sst_replace.rs b/src/cmd/src/datanode/sst_replace.rs new file mode 100644 index 00000000000..df13fcd6f25 --- /dev/null +++ b/src/cmd/src/datanode/sst_replace.rs @@ -0,0 +1,361 @@ +// Copyright 2023 Greptime Team +// +// 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. + +use std::path::PathBuf; + +use clap::Parser; +use colored::Colorize; +use mito2::manifest::action::{RegionEdit, RegionMetaAction, RegionMetaActionList}; +use mito2::manifest::manager::{RegionManifestManager, RegionManifestOptions}; +use mito2::region::ManifestStats; +use mito2::sst::file::RegionFileId; +use mito2::sst::location::{region_dir_from_table_dir, sst_file_path}; +use object_store::ObjectStore; +use parquet::file::FOOTER_SIZE; +use parquet::file::metadata::{FooterTail, ParquetMetaDataReader}; +use store_api::region_request::PathType; + +use crate::datanode::tool_util::{ + build_object_store, max_row_group_uncompressed_size, parse_config, parse_file_id, + parse_path_type, parse_region_id, +}; +use crate::error; + +/// Replace a mito region SST and update the corresponding manifest metadata. +#[derive(Debug, Parser)] +pub struct SstReplaceCommand { + /// Path to config TOML file (same format as standalone/datanode config). + #[clap(long, value_name = "FILE")] + config: PathBuf, + + /// Region ID: either numeric u64 (e.g. "4398046511104") or "table_id:region_num" (e.g. "1024:0"). + #[clap(long)] + region_id: String, + + /// Table directory relative to data home (e.g. "data/greptime/public/1024/"). + #[clap(long)] + table_dir: String, + + /// SST file id to replace. + #[clap(long)] + file_id: String, + + /// Local parquet file used as the replacement. + #[clap(long, value_name = "FILE", conflicts_with = "replacement_object")] + replacement_file: Option, + + /// Object-store parquet path used as the replacement. + #[clap(long, value_name = "PATH", conflicts_with = "replacement_file")] + replacement_object: Option, + + /// Path type for the region: auto, bare, data, metadata. + #[clap(long, default_value = "auto")] + path_type: String, + + /// Actually overwrite the SST object and append a manifest delta. + #[clap(long, default_value_t = false)] + confirm: bool, + + /// Verbose output. + #[clap(short, long, default_value_t = false)] + verbose: bool, +} + +impl SstReplaceCommand { + pub async fn run(&self) -> error::Result<()> { + if self.verbose { + common_telemetry::init_default_ut_logging(); + } + + println!("{}", "Starting sst-replace...".cyan().bold()); + + let region_id = parse_region_id(&self.region_id)?; + let file_id = parse_file_id(&self.file_id)?; + let path_type = parse_optional_path_type(&self.path_type)?; + let replacement = self.replacement()?; + + let (store_cfg, mito_config, _wal_config) = parse_config(&self.config)?; + let object_store = build_object_store(&store_cfg).await?; + println!("{} Object store initialized", "[ok]".green()); + + let candidates = path_type + .map(|path_type| vec![path_type]) + .unwrap_or_else(|| vec![PathType::Bare, PathType::Data, PathType::Metadata]); + let mut found = None; + let mut existing_manifests = Vec::new(); + for candidate in candidates { + let stats = ManifestStats::default(); + let region_dir = region_dir_from_table_dir(&self.table_dir, region_id, candidate); + let manifest_opts = + RegionManifestOptions::new(&mito_config, ®ion_dir, &object_store); + let Some(manifest_manager) = RegionManifestManager::open(manifest_opts, &stats) + .await + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("open {} manifest failed: {e:?}", path_type_name(candidate)), + } + .build() + })? + else { + continue; + }; + existing_manifests.push(path_type_name(candidate)); + let manifest = manifest_manager.manifest(); + if let Some(file) = manifest.files.get(&file_id) { + if found.is_some() { + return error::IllegalConfigSnafu { + msg: format!( + "file {} exists in multiple manifests under region {}; specify --path-type", + file_id, region_id + ), + } + .fail(); + } + found = Some((candidate, manifest_manager, file.clone())); + } + } + let Some((path_type, mut manifest_manager, old_file)) = found else { + let suffix = if existing_manifests.is_empty() { + "no manifest found".to_string() + } else { + format!( + "manifests found for [{}], but none contains file {}", + existing_manifests.join(", "), + file_id + ) + }; + return error::IllegalConfigSnafu { + msg: format!("region manifest not found for {}: {}", region_id, suffix), + } + .fail(); + }; + + let region_file_id = RegionFileId::new(old_file.region_id, old_file.file_id); + let target_path = sst_file_path(&self.table_dir, region_file_id, path_type); + let target_stat = object_store.stat(&target_path).await.map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("stat target SST {} failed: {e}", target_path), + } + .build() + })?; + let replacement_bytes = load_replacement_bytes(&object_store, replacement).await?; + let new_file_size = replacement_bytes.len() as u64; + let parquet_meta = decode_parquet_metadata(&replacement_bytes).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("read replacement parquet metadata failed: {e}"), + } + .build() + })?; + let new_num_rows = parquet_meta.file_metadata().num_rows() as u64; + let new_num_row_groups = parquet_meta.num_row_groups() as u64; + validate_replacement(&old_file, new_num_rows, new_num_row_groups)?; + + let mut new_file = old_file.clone(); + new_file.file_size = new_file_size; + new_file.num_rows = new_num_rows; + new_file.num_row_groups = new_num_row_groups; + new_file.max_row_group_uncompressed_size = max_row_group_uncompressed_size(&parquet_meta); + + println!("{} Region: {}", "[ok]".green(), region_id); + println!("{} Target SST: {}", "[ok]".green(), target_path.cyan()); + println!( + "{} Size: {} -> {} bytes", + "[ok]".green(), + target_stat.content_length(), + new_file_size + ); + println!( + "{} Rows: {}, row groups: {}", + "[ok]".green(), + new_num_rows, + new_num_row_groups + ); + + if !self.confirm { + println!( + "{} Dry run only. Re-run with --confirm to overwrite the SST and update the manifest.", + "[dry-run]".yellow() + ); + return Ok(()); + } + + object_store + .write(&target_path, replacement_bytes) + .await + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("write replacement SST {} failed: {e}", target_path), + } + .build() + })?; + + let edit = RegionEdit { + files_to_add: vec![new_file], + files_to_remove: vec![], + timestamp_ms: None, + compaction_time_window: None, + flushed_entry_id: None, + flushed_sequence: None, + committed_sequence: None, + }; + let version = manifest_manager + .update( + RegionMetaActionList::with_action(RegionMetaAction::Edit(edit)), + false, + ) + .await + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("update manifest failed: {e:?}"), + } + .build() + })?; + + println!( + "{} Replacement committed, manifest version {}", + "[ok]".green(), + version + ); + Ok(()) + } + + fn replacement(&self) -> error::Result> { + match (&self.replacement_file, &self.replacement_object) { + (Some(path), None) => Ok(Replacement::Local(path)), + (None, Some(path)) => Ok(Replacement::Object(path)), + _ => error::IllegalConfigSnafu { + msg: "specify exactly one of --replacement-file or --replacement-object" + .to_string(), + } + .fail(), + } + } +} + +enum Replacement<'a> { + Local(&'a PathBuf), + Object(&'a str), +} + +async fn load_replacement_bytes( + object_store: &ObjectStore, + replacement: Replacement<'_>, +) -> error::Result> { + match replacement { + Replacement::Local(path) => tokio::fs::read(path).await.map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("read replacement file {} failed: {e}", path.display()), + } + .build() + }), + Replacement::Object(path) => { + object_store + .read(path) + .await + .map(|b| b.to_vec()) + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("read replacement object {} failed: {e}", path), + } + .build() + }) + } + } +} + +fn validate_replacement( + old_file: &mito2::sst::file::FileMeta, + new_num_rows: u64, + new_num_row_groups: u64, +) -> error::Result<()> { + if old_file.num_rows != 0 && old_file.num_rows != new_num_rows { + return error::IllegalConfigSnafu { + msg: format!( + "replacement row count mismatch: manifest has {}, replacement has {}", + old_file.num_rows, new_num_rows + ), + } + .fail(); + } + if old_file.num_row_groups != 0 && old_file.num_row_groups != new_num_row_groups { + return error::IllegalConfigSnafu { + msg: format!( + "replacement row group count mismatch: manifest has {}, replacement has {}", + old_file.num_row_groups, new_num_row_groups + ), + } + .fail(); + } + Ok(()) +} + +fn decode_parquet_metadata( + data: &[u8], +) -> Result> { + if data.len() < FOOTER_SIZE { + return Err("file too small".into()); + } + let footer_start = data.len() - FOOTER_SIZE; + let mut footer = [0; FOOTER_SIZE]; + footer.copy_from_slice(&data[footer_start..]); + let footer = FooterTail::try_new(&footer)?; + let metadata_len = footer.metadata_length(); + if footer_start < metadata_len { + return Err("invalid footer/metadata length".into()); + } + let metadata_start = footer_start - metadata_len; + Ok(ParquetMetaDataReader::decode_metadata( + &data[metadata_start..footer_start], + )?) +} + +fn parse_optional_path_type(value: &str) -> error::Result> { + match value.to_lowercase().as_str() { + "auto" => Ok(None), + _ => parse_path_type(value).map(Some).map_err(|_| { + error::IllegalConfigSnafu { + msg: format!("invalid path_type '{value}', expected: auto, bare, data, metadata"), + } + .build() + }), + } +} + +fn path_type_name(path_type: PathType) -> &'static str { + match path_type { + PathType::Bare => "bare", + PathType::Data => "data", + PathType::Metadata => "metadata", + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_validate_replacement_counts() { + let old_file = mito2::sst::file::FileMeta { + num_rows: 4, + num_row_groups: 1, + ..Default::default() + }; + assert!(validate_replacement(&old_file, 4, 1).is_ok()); + assert!(validate_replacement(&old_file, 3, 1).is_err()); + assert!(validate_replacement(&old_file, 4, 2).is_err()); + + let legacy_file = mito2::sst::file::FileMeta::default(); + assert!(validate_replacement(&legacy_file, 4, 1).is_ok()); + } +} diff --git a/src/cmd/src/datanode/tool_util.rs b/src/cmd/src/datanode/tool_util.rs new file mode 100644 index 00000000000..3b16415e89d --- /dev/null +++ b/src/cmd/src/datanode/tool_util.rs @@ -0,0 +1,236 @@ +// Copyright 2023 Greptime Team +// +// 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. + +#[cfg(feature = "dev-tools")] +use std::fs::File; +use std::path::Path; +use std::sync::Arc; + +use common_wal::config::DatanodeWalConfig; +use datanode::config::RegionEngineConfig; +use datanode::store; +use mito2::config::MitoConfig; +use object_store::ObjectStore; +#[cfg(feature = "dev-tools")] +use parquet::basic::Compression; +use parquet::file::metadata::{KeyValue, ParquetMetaData}; +#[cfg(feature = "dev-tools")] +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader}; +use snafu::OptionExt; +#[cfg(feature = "dev-tools")] +use snafu::ResultExt; +use store_api::metadata::{RegionMetadata, RegionMetadataRef}; +use store_api::region_request::PathType; +use store_api::storage::{FileId, RegionId}; + +use crate::datanode::{StorageConfig, StorageConfigWrapper}; +use crate::error; + +pub(crate) fn parse_config( + config_path: &Path, +) -> error::Result<(StorageConfig, MitoConfig, DatanodeWalConfig)> { + let cfg_str = std::fs::read_to_string(config_path).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("failed to read config {}: {e}", config_path.display()), + } + .build() + })?; + + let store_cfg: StorageConfigWrapper = toml::from_str(&cfg_str).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("failed to parse config {}: {e}", config_path.display()), + } + .build() + })?; + + let wal_config = store_cfg.wal; + let storage_config = store_cfg.storage; + let mito_engine_config = store_cfg + .region_engine + .into_iter() + .find_map(|config| match config { + RegionEngineConfig::Mito(mito) => Some(mito), + _ => None, + }) + .with_context(|| error::IllegalConfigSnafu { + msg: format!("Engine config not found in {:?}", config_path), + })?; + + Ok((storage_config, mito_engine_config, wal_config)) +} + +pub(crate) async fn build_object_store(config: &StorageConfig) -> error::Result { + store::new_object_store(config.store.clone(), &config.data_home) + .await + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("Failed to build object store: {e:?}"), + } + .build() + }) +} + +pub(crate) fn extract_region_metadata( + file_path: &str, + metadata: &ParquetMetaData, +) -> error::Result { + let key_values: Option<&Vec> = metadata.file_metadata().key_value_metadata(); + let Some(key_values) = key_values else { + return Err(error::IllegalConfigSnafu { + msg: format!("{file_path}: missing parquet key_value metadata"), + } + .build()); + }; + let json = key_values + .iter() + .find(|key_value| key_value.key == mito2::sst::parquet::PARQUET_METADATA_KEY) + .and_then(|key_value| key_value.value.as_ref()) + .ok_or_else(|| { + error::IllegalConfigSnafu { + msg: format!( + "{file_path}: key {} not found or empty", + mito2::sst::parquet::PARQUET_METADATA_KEY + ), + } + .build() + })?; + let region = RegionMetadata::from_json(json).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid region metadata json: {e}"), + } + .build() + })?; + Ok(Arc::new(region)) +} + +pub(crate) fn parse_region_id(value: &str) -> error::Result { + if let Some((table_id, region_number)) = value.split_once(':') { + let table_id = table_id.parse().map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid table_id in region_id '{value}': {e}"), + } + .build() + })?; + let region_number = region_number.parse().map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid region_num in region_id '{value}': {e}"), + } + .build() + })?; + Ok(RegionId::new(table_id, region_number)) + } else { + value.parse().map(RegionId::from_u64).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid region_id '{value}': {e}"), + } + .build() + }) + } +} + +pub(crate) fn parse_file_id(value: &str) -> error::Result { + FileId::parse_str(value).map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("invalid file_id '{value}': {e}"), + } + .build() + }) +} + +pub(crate) fn parse_path_type(value: &str) -> error::Result { + match value.to_lowercase().as_str() { + "bare" => Ok(PathType::Bare), + "data" => Ok(PathType::Data), + "metadata" => Ok(PathType::Metadata), + _ => Err(error::IllegalConfigSnafu { + msg: format!("invalid path_type '{value}', expected: bare, data, metadata"), + } + .build()), + } +} + +pub(crate) fn format_bytes(bytes: u64) -> String { + const KIB: u64 = 1024; + const MIB: u64 = 1024 * KIB; + const GIB: u64 = 1024 * MIB; + if bytes >= GIB { + format!("{:.2} GiB", bytes as f64 / GIB as f64) + } else if bytes >= MIB { + format!("{:.2} MiB", bytes as f64 / MIB as f64) + } else if bytes >= KIB { + format!("{:.2} KiB", bytes as f64 / KIB as f64) + } else { + format!("{bytes} B") + } +} + +pub(crate) fn max_row_group_uncompressed_size(metadata: &ParquetMetaData) -> u64 { + metadata + .row_groups() + .iter() + .map(|row_group| { + row_group + .columns() + .iter() + .map(|column| column.uncompressed_size() as u64) + .sum::() + }) + .max() + .unwrap_or(0) +} + +#[cfg(feature = "dev-tools")] +pub(crate) fn load_local_parquet_metadata(path: &Path) -> error::Result { + let file = File::open(path).context(error::FileIoSnafu)?; + ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Optional) + .parse_and_finish(&file) + .map_err(|e| { + error::IllegalConfigSnafu { + msg: format!("read parquet metadata failed for {}: {e}", path.display()), + } + .build() + }) +} + +#[cfg(feature = "dev-tools")] +pub(crate) fn compression_name(compression: Compression) -> &'static str { + match compression { + Compression::UNCOMPRESSED => "uncompressed", + Compression::SNAPPY => "snappy", + Compression::GZIP(_) => "gzip", + Compression::LZO => "lzo", + Compression::BROTLI(_) => "brotli", + Compression::LZ4 => "lz4", + Compression::ZSTD(_) => "zstd", + Compression::LZ4_RAW => "lz4-raw", + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_region_and_path_type() { + assert_eq!(parse_region_id("1024:7").unwrap(), RegionId::new(1024, 7)); + assert_eq!( + parse_region_id(&RegionId::new(1, 2).as_u64().to_string()).unwrap(), + RegionId::new(1, 2) + ); + assert_eq!(parse_path_type("bare").unwrap(), PathType::Bare); + assert_eq!(parse_path_type("data").unwrap(), PathType::Data); + assert_eq!(parse_path_type("metadata").unwrap(), PathType::Metadata); + } +} From 3493d2d0fbfd7fe05bf9e3276b98bdf727ee9aaf Mon Sep 17 00:00:00 2001 From: discord9 Date: Tue, 25 Aug 2026 05:16:50 +0000 Subject: [PATCH 75/92] perf(servers)!: speed up Prometheus JSON response building with ryu and per-series entry reuse (#8815) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(perf): align direct-SST CREATE TABLE with baked index metadata The offline fixture generator (query_perf_fixture::direct_sst:: build_region_metadata) bakes greptime:inverted_index / greptime:skipping_index field metadata into the region manifest for tag/field columns, but create_table_sql emitted a bare CREATE TABLE without those declarations. MergeScan's remote-schema validation then failed on any tag/field projection (HTTP 500 'advertised remote stream schema field mismatch'), breaking direct_readable_sst perf cases. CREATE TABLE now declares the matching SKIPPING INDEX WITH (granularity='1') / INVERTED INDEX column options. A round-trip test proves the emitted SQL is parser-valid and yields the exact catalog metadata. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * perf(servers): speed up Prometheus JSON response building with ryu and per-series entry reuse PrometheusJsonResponse::record_batches_to_data spends ~47% of its CPU in f64::to_string() per sample and ~32% in IndexMap::entry() per row (60s profile of concurrent query_range workloads, ~800k series). - Replace f64::to_string() with ryu::Buffer::format_finite for finite values (shortest round-trip, 3-5x faster); NaN/+Inf/-Inf keep the previous std formatting so wire output is unchanged. - Remember the previous row's label vector and entry index; query output is clustered by series, so consecutive rows reuse the same IndexMap entry via get_index_mut instead of rebuilding and hashing the label vector (worst case adds one Vec comparison per series transition). Also adds a query-regression case (prom_json_response) that measures the real Prometheus HTTP range API path (/v1/prometheus/api/v1/query_range), which is the only frontend path that builds the Prometheus JSON response (TQL ANALYZE formats the SQL JSON shape instead), plus a prom_http query kind in the regression runner. Perf (aligned base d90cca4b75, 256 series x 481 points): - prom_range_2h (JSON response path): 29.31ms -> 21.46ms (-26.8%) - tql_range_2h_control (non-JSON path): +1.87% (noise) Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(servers): keep Prometheus wire format for integral floats with ryu ryu::Buffer::format_finite prints integral values as "1.0", but the Prometheus JSON wire format (matching std f64::to_string) expects "1". Strip the trailing ".0" that ryu only emits for integral values; extreme values keep ryu scientific notation, and NaN/Inf keep std output. Adds wire-format tests covering 1.0, 0.0, -0.0, 1.5, 0.1, 1e21, 1e30, 1e-7, f64::MAX, NaN, ±Inf. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(servers): address Prometheus response review feedback Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * perf(servers)!: use ryu for Prometheus sample values Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(cmd): skip Prometheus execution time extraction Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- Cargo.lock | 1 + Cargo.toml | 1 + .../src/bin/query_regression_runner/direct.rs | 106 ++++++- .../bin/query_regression_runner/measure.rs | 64 +++- .../src/bin/query_regression_runner/model.rs | 13 + .../src/bin/query_regression_runner/sql.rs | 72 ++++- src/servers/Cargo.toml | 1 + .../src/http/result/prometheus_resp.rs | 295 +++++++++++++++++- tests-integration/tests/grpc.rs | 8 +- tests-integration/tests/http.rs | 6 +- tests/perf/fixture-format.md | 14 + .../query_cases/prom_json_response/case.toml | 89 ++++++ 12 files changed, 644 insertions(+), 26 deletions(-) create mode 100644 tests/perf/query_cases/prom_json_response/case.toml diff --git a/Cargo.lock b/Cargo.lock index a389391ea90..534f484d4fc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -13470,6 +13470,7 @@ dependencies = [ "rustls", "rustls-pemfile", "rustls-pki-types", + "ryu", "serde", "serde_json", "servers", diff --git a/Cargo.toml b/Cargo.toml index e6b849a6cd0..7c3fccfbbf8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -231,6 +231,7 @@ rstest_reuse = "0.7" rust_decimal = "1.33" rustc-hash = "2.0" rustls = { version = "0.23.25", default-features = false } +ryu = "1.0" sea-query = "0.32" serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["float_roundtrip"] } diff --git a/src/cmd/src/bin/query_regression_runner/direct.rs b/src/cmd/src/bin/query_regression_runner/direct.rs index 04c191c8c55..b003b6ee5d3 100644 --- a/src/cmd/src/bin/query_regression_runner/direct.rs +++ b/src/cmd/src/bin/query_regression_runner/direct.rs @@ -166,7 +166,28 @@ fn create_table_sql(table: &Table) -> Result { let columns = table .columns .iter() - .map(|column| format!("{} {}", sql_ident(&column.name), column.data_type)) + .map(|column| { + // The direct-SST fixture generator (`query_perf_fixture::direct_sst:: + // build_region_metadata`) bakes index metadata into the region + // manifest's column schemas: tag columns get an inverted index and + // a skipping index with granularity 1, field columns get a skipping + // index with granularity 1. The CREATE TABLE must declare the same + // column options, otherwise the table schema (frontend) and the + // region schema (datanode, from the manifest) disagree and the + // MergeScan remote schema validation fails with "advertised remote + // stream schema field mismatch". + let mut sql = format!("{} {}", sql_ident(&column.name), column.data_type); + match column.semantic.as_deref() { + Some("tag") => { + sql.push_str(" SKIPPING INDEX WITH (granularity='1') INVERTED INDEX"); + } + Some("field") => { + sql.push_str(" SKIPPING INDEX WITH (granularity='1')"); + } + _ => {} + } + sql + }) .collect::>() .join(",\n "); let primary_key = table @@ -455,10 +476,12 @@ mod tests { Column { name: "host".to_string(), data_type: "STRING".to_string(), + semantic: Some("tag".to_string()), }, Column { name: "ts".to_string(), data_type: "TIMESTAMP(9)".to_string(), + semantic: Some("timestamp".to_string()), }, ], primary_key: vec!["host".to_string()], @@ -469,7 +492,86 @@ mod tests { }; assert_eq!( create_table_sql(&table).unwrap(), - "CREATE TABLE \"metric\"\"name\" (\n \"host\" STRING,\n \"ts\" TIMESTAMP(9),\n TIME INDEX (\"ts\"),\n PRIMARY KEY (\"host\")\n) ENGINE=mito\nWITH ('append_mode'='true', 'sst_format'='flat');" + "CREATE TABLE \"metric\"\"name\" (\n \"host\" STRING SKIPPING INDEX WITH (granularity='1') INVERTED INDEX,\n \"ts\" TIMESTAMP(9),\n TIME INDEX (\"ts\"),\n PRIMARY KEY (\"host\")\n) ENGINE=mito\nWITH ('append_mode'='true', 'sst_format'='flat');" + ); + } + + #[test] + fn create_table_sql_declares_indexes_matching_fixture_metadata() { + // Mirrors the index metadata baked by + // query_perf_fixture::direct_sst::build_region_metadata: tag columns + // get SKIPPING INDEX WITH (granularity='1') + INVERTED INDEX, field + // columns get SKIPPING INDEX WITH (granularity='1'), timestamp columns + // get no column options. Without these options the frontend table + // schema and the datanode region schema (from the fixture manifest) + // disagree and the MergeScan remote schema validation fails. + let table = Table { + database: "public".to_string(), + name: "metric".to_string(), + engine: "mito".to_string(), + columns: vec![ + Column { + name: "host".to_string(), + data_type: "STRING".to_string(), + semantic: Some("tag".to_string()), + }, + Column { + name: "instance".to_string(), + data_type: "STRING".to_string(), + semantic: Some("tag".to_string()), + }, + Column { + name: "value".to_string(), + data_type: "DOUBLE".to_string(), + semantic: Some("field".to_string()), + }, + Column { + name: "ts".to_string(), + data_type: "TIMESTAMP(9)".to_string(), + semantic: Some("timestamp".to_string()), + }, + ], + primary_key: vec!["host".to_string(), "instance".to_string()], + time_index: Some("ts".to_string()), + append_mode: Some(true), + sst_format: Some("flat".to_string()), + validate_show_create_engine: true, + }; + assert_eq!( + create_table_sql(&table).unwrap(), + "CREATE TABLE \"metric\" (\n \"host\" STRING SKIPPING INDEX WITH (granularity='1') INVERTED INDEX,\n \"instance\" STRING SKIPPING INDEX WITH (granularity='1') INVERTED INDEX,\n \"value\" DOUBLE SKIPPING INDEX WITH (granularity='1'),\n \"ts\" TIMESTAMP(9),\n TIME INDEX (\"ts\"),\n PRIMARY KEY (\"host\", \"instance\")\n) ENGINE=mito\nWITH ('append_mode'='true', 'sst_format'='flat');" + ); + } + + #[test] + fn create_table_sql_without_semantic_keeps_bare_columns() { + // Columns without a semantic type must keep the historical bare form so + // tables that predate the index-metadata fix are unaffected. + let table = Table { + database: "public".to_string(), + name: "metric".to_string(), + engine: "mito".to_string(), + columns: vec![ + Column { + name: "host".to_string(), + data_type: "STRING".to_string(), + semantic: None, + }, + Column { + name: "ts".to_string(), + data_type: "TIMESTAMP(9)".to_string(), + semantic: None, + }, + ], + primary_key: vec!["host".to_string()], + time_index: Some("ts".to_string()), + append_mode: None, + sst_format: None, + validate_show_create_engine: true, + }; + assert_eq!( + create_table_sql(&table).unwrap(), + "CREATE TABLE \"metric\" (\n \"host\" STRING,\n \"ts\" TIMESTAMP(9),\n TIME INDEX (\"ts\"),\n PRIMARY KEY (\"host\")\n) ENGINE=mito;" ); } diff --git a/src/cmd/src/bin/query_regression_runner/measure.rs b/src/cmd/src/bin/query_regression_runner/measure.rs index 2b389f7d782..fd80813f40e 100644 --- a/src/cmd/src/bin/query_regression_runner/measure.rs +++ b/src/cmd/src/bin/query_regression_runner/measure.rs @@ -21,7 +21,7 @@ use serde_json::{Map, Value, json}; use crate::query_regression_runner::model::{Measurement, Query, QueryResult, Scenario, Table}; use crate::query_regression_runner::plan::{load_plan, normalize_scenario}; -use crate::query_regression_runner::sql::{http_post_sql, sql_ident}; +use crate::query_regression_runner::sql::{http_post_prom_range_query, http_post_sql, sql_ident}; use crate::query_regression_runner::{MeasureArgs, Result}; pub(super) async fn run_measure(args: MeasureArgs) -> Result<()> { @@ -106,6 +106,26 @@ fn target_report(name: &str, http_port: u16, result: QueryResult) -> Value { }) } +/// Routes a configured query to the right endpoint: `prom_http` queries hit +/// the Prometheus HTTP range API (which exercises the Prometheus JSON response +/// builder), everything else goes through `/v1/sql` (including `TQL ANALYZE`). +async fn post_query(client: &Client, port: u16, query: &Query, db: &str) -> Value { + if query.kind.as_deref() == Some("prom_http") { + http_post_prom_range_query( + client, + port, + &query.query, + query.start.as_deref(), + query.end.as_deref(), + query.step.as_deref(), + db, + ) + .await + } else { + http_post_sql(client, port, &query.query, db).await + } +} + async fn run_target( port: u16, tables: &[Table], @@ -118,6 +138,9 @@ async fn run_target( name: Some("count_all".to_string()), kind: Some("sql".to_string()), query: format!("SELECT count(*) FROM {}", sql_ident(&tables[0].name)), + start: None, + end: None, + step: None, warmup: 0, iterations: 1, thresholds: Map::new(), @@ -147,7 +170,7 @@ async fn run_target( } validation.push(sample); } - let first = http_post_sql(client, port, &queries[0].query, db).await; + let first = post_query(client, port, &queries[0], db).await; if !first["ok"].as_bool().unwrap_or(false) { validation_errors.push(json!({ "sql": queries[0].query, @@ -160,7 +183,7 @@ async fn run_target( let mut measurements = Vec::with_capacity(queries.len()); for query in &queries { for _ in 0..query.warmup { - let warmup = http_post_sql(client, port, &query.query, db).await; + let warmup = post_query(client, port, query, db).await; if !warmup["ok"].as_bool().unwrap_or(false) { validation_errors.push(json!({ "sql": query.query, @@ -173,10 +196,12 @@ async fn run_target( let mut samples = Vec::with_capacity(query.iterations); let mut good_latencies = Vec::with_capacity(query.iterations); for _ in 0..query.iterations { - let mut sample = http_post_sql(client, port, &query.query, db).await; + let mut sample = post_query(client, port, query, db).await; let execution_time = sample .get("response") - .and_then(extract_execution_time) + .and_then(|response| { + extract_execution_time_for_kind(query.kind.as_deref(), response) + }) .cloned() .unwrap_or(Value::Null); sample @@ -242,6 +267,14 @@ fn response_text(body: &Value) -> String { .unwrap_or_else(|| serde_json::to_string(body).unwrap_or_default()) } +fn extract_execution_time_for_kind<'a>(kind: Option<&str>, body: &'a Value) -> Option<&'a Value> { + if kind == Some("prom_http") { + None + } else { + extract_execution_time(body) + } +} + fn extract_execution_time(body: &Value) -> Option<&Value> { match body { Value::Object(map) => { @@ -379,12 +412,33 @@ mod tests { ); } + #[test] + fn prom_http_does_not_extract_execution_time_from_response() { + let response = json!({ + "data": { + "result": [{"metric": {"job": "api", "elapsed": "label"}}], + "elapsed": 42 + } + }); + assert_eq!( + extract_execution_time_for_kind(Some("prom_http"), &response), + None + ); + assert_eq!( + extract_execution_time_for_kind(Some("sql"), &response), + Some(&json!(42)) + ); + } + #[test] fn threshold_rejects_unknown_keys_and_zero_base() { let query = Query { name: Some("q".to_string()), kind: None, query: "SELECT 1".to_string(), + start: None, + end: None, + step: None, warmup: 0, iterations: 1, thresholds: Map::from_iter([ diff --git a/src/cmd/src/bin/query_regression_runner/model.rs b/src/cmd/src/bin/query_regression_runner/model.rs index 50d41b3f85f..6831a5256a4 100644 --- a/src/cmd/src/bin/query_regression_runner/model.rs +++ b/src/cmd/src/bin/query_regression_runner/model.rs @@ -90,6 +90,12 @@ pub(super) struct Column { pub(super) name: String, #[serde(rename = "type")] pub(super) data_type: String, + /// Semantic type ("tag", "field", "timestamp") from the case TOML. The + /// direct-SST fixture generator bakes index metadata into the region + /// manifest based on this semantic, so the CREATE TABLE must declare the + /// matching column options (see `create_table_sql` in `direct.rs`). + #[serde(default)] + pub(super) semantic: Option, } const fn default_show_create_engine() -> bool { @@ -203,6 +209,13 @@ pub(super) struct Query { #[serde(default)] pub(super) kind: Option, pub(super) query: String, + /// Prometheus HTTP range query parameters (`kind = "prom_http"` only). + #[serde(default)] + pub(super) start: Option, + #[serde(default)] + pub(super) end: Option, + #[serde(default)] + pub(super) step: Option, #[serde(default)] pub(super) warmup: usize, #[serde(default = "one")] diff --git a/src/cmd/src/bin/query_regression_runner/sql.rs b/src/cmd/src/bin/query_regression_runner/sql.rs index 4d8baab2c8c..b33c0a1c22b 100644 --- a/src/cmd/src/bin/query_regression_runner/sql.rs +++ b/src/cmd/src/bin/query_regression_runner/sql.rs @@ -117,10 +117,62 @@ pub(super) fn value_f64(value: Option<&Value>) -> Option { } pub(super) async fn http_post_sql(client: &Client, port: u16, sql: &str, db: &str) -> Value { + let mut sample = post_form( + client, + format!("http://127.0.0.1:{port}/v1/sql"), + &[("sql", sql), ("db", db), ("format", "json")], + ) + .await; + sample + .as_object_mut() + .expect("HTTP samples are objects") + .insert("sql".to_string(), Value::String(sql.to_string())); + sample +} + +/// Posts a Prometheus HTTP API range query (`/v1/prometheus/api/v1/query_range`) +/// and measures the full request-to-body latency. This exercises the Prometheus +/// JSON response building path (`PrometheusJsonResponse::record_batches_to_data`) +/// that `/v1/sql` (including `TQL ANALYZE`) does not go through. +pub(super) async fn http_post_prom_range_query( + client: &Client, + port: u16, + query: &str, + start: Option<&str>, + end: Option<&str>, + step: Option<&str>, + db: &str, +) -> Value { + let mut sample = post_form( + client, + prom_range_query_url(port, db), + &[ + ("query", query), + ("start", start.unwrap_or_default()), + ("end", end.unwrap_or_default()), + ("step", step.unwrap_or_default()), + ], + ) + .await; + sample + .as_object_mut() + .expect("HTTP samples are objects") + .insert("query".to_string(), Value::String(query.to_string())); + sample +} + +fn prom_range_query_url(port: u16, db: &str) -> String { + let mut url = reqwest::Url::parse(&format!( + "http://127.0.0.1:{port}/v1/prometheus/api/v1/query_range" + )) + .expect("fixed Prometheus range query URL must be valid"); + url.query_pairs_mut().append_pair("db", db); + url.into() +} + +async fn post_form(client: &Client, url: String, form: &[(&str, &str)]) -> Value { let started = Instant::now(); - let request = client - .post(format!("http://127.0.0.1:{port}/v1/sql")) - .form(&[("sql", sql), ("db", db), ("format", "json")]); + let request = client.post(url).form(form); match request.send().await { Ok(response) => { let status = response.status().as_u16(); @@ -133,7 +185,6 @@ pub(super) async fn http_post_sql(client: &Client, port: u16, sql: &str, db: &st "status": status, "latency_ms": started.elapsed().as_secs_f64() * 1000.0, "response": body, - "sql": sql, }); if status >= 400 { sample @@ -148,7 +199,6 @@ pub(super) async fn http_post_sql(client: &Client, port: u16, sql: &str, db: &st "status": status, "latency_ms": started.elapsed().as_secs_f64() * 1000.0, "error": error.to_string(), - "sql": sql, }), } } @@ -157,7 +207,6 @@ pub(super) async fn http_post_sql(client: &Client, port: u16, sql: &str, db: &st "status": Value::Null, "latency_ms": started.elapsed().as_secs_f64() * 1000.0, "error": error.to_string(), - "sql": sql, }), } } @@ -208,6 +257,17 @@ pub(super) fn sql_ident(name: &str) -> String { mod tests { use super::*; + #[test] + fn prom_range_query_url_places_database_in_query_parameters() { + let url = reqwest::Url::parse(&prom_range_query_url(4000, "catalog-schema name")) + .expect("generated URL must be valid"); + assert_eq!(url.path(), "/v1/prometheus/api/v1/query_range"); + assert_eq!( + url.query_pairs().collect::>(), + vec![("db".into(), "catalog-schema name".into())] + ); + } + #[test] fn top_level_errors_do_not_inspect_rows() { assert!(response_has_error(&json!({"error_code": 7}))); diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 59d00b90ad7..7936c96fbda 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -115,6 +115,7 @@ rust_decimal = { workspace = true, features = ["db-postgres"] } rustls = { workspace = true, default-features = false, features = ["aws_lc_rs", "logging", "std", "tls12"] } rustls-pemfile = "2.0" rustls-pki-types = "1.0" +ryu.workspace = true serde.workspace = true serde_json.workspace = true session.workspace = true diff --git a/src/servers/src/http/result/prometheus_resp.rs b/src/servers/src/http/result/prometheus_resp.rs index e79628fac0d..bc1d0849037 100644 --- a/src/servers/src/http/result/prometheus_resp.rs +++ b/src/servers/src/http/result/prometheus_resp.rs @@ -38,6 +38,7 @@ use datatypes::prelude::ConcreteDataType; use indexmap::IndexMap; use promql_parser::label::METRIC_NAME; use promql_parser::parser::value::ValueType; +use ryu::Buffer; use serde::{Deserialize, Serialize}; use serde_json::Value; use snafu::{OptionExt, ResultExt}; @@ -70,6 +71,19 @@ fn prometheus_native_histogram(histogram: &NativeHistogram) -> Result String { + if value.is_finite() { + Buffer::new().format_finite(value).to_string() + } else { + value.to_string() + } +} + #[derive(Debug, Default, Serialize, Deserialize, PartialEq)] pub struct PrometheusJsonResponse { pub status: String, @@ -307,6 +321,15 @@ impl PrometheusJsonResponse { // Tag order matters, e.g., after sorc and sort_desc, the output order must be kept. let mut buffer = IndexMap::, PromSeriesSamples>::new(); + // Query output is clustered by series (the range plan sorts by series + // key + timestamp), so consecutive rows usually belong to the same + // series. Remember the index of the previous row's entry in `buffer`, + // and reuse it directly when its tags are unchanged. This avoids + // building and hashing the label vector on every row; the worst case + // adds one `Vec` comparison per series transition before falling back + // to the map lookup. + let mut last_entry_index = None; + let schema = batches.schema(); for batch in batches.iter() { // prepare things... @@ -380,15 +403,32 @@ impl PrometheusJsonResponse { } } - let entry = buffer.entry(tags).or_default(); + let reuse = last_entry_index.filter(|index| { + buffer + .get_index(*index) + .is_some_and(|(key, _)| key == &tags) + }); + let samples = if let Some(index) = reuse { + buffer + .get_index_mut(index) + .map(|(_, samples)| samples) + .with_context(|| UnexpectedResultSnafu { + reason: "reused series entry must exist", + })? + } else { + let entry = buffer.entry(tags); + last_entry_index = Some(entry.index()); + entry.or_default() + }; if let Some((timestamp_millis, histogram)) = histogram { - entry + samples .histograms .push((timestamp_millis as f64 / 1000.0, histogram)); } else if let Some((timestamp_millis, value)) = value { - entry - .values - .push((timestamp_millis as f64 / 1000.0, value.to_string())); + samples.values.push(( + timestamp_millis as f64 / 1000.0, + format_prometheus_sample_value(value), + )); } } } @@ -564,6 +604,67 @@ mod tests { Arc::new(StructVector::try_new(histogram_type, histogram_array).unwrap()) } + #[test] + fn format_prometheus_sample_value_uses_ryu_for_finite_values() { + let values = [ + 1.5, + 0.1, + 1.0, + 0.0, + -0.0, + 100.0, + 1e-6, + 1e-7, + 1e21, + 1e30, + f64::MAX, + f64::MIN_POSITIVE, + ]; + + for value in values { + let output = format_prometheus_sample_value(value); + assert_eq!(output.parse::().unwrap().to_bits(), value.to_bits()); + } + + // Representative integral values use ryu's explicit .0 form. + assert_eq!(format_prometheus_sample_value(1.0), "1.0"); + assert_eq!(format_prometheus_sample_value(-0.0), "-0.0"); + assert_eq!(format_prometheus_sample_value(100.0), "100.0"); + assert_eq!(format_prometheus_sample_value(1e-6), "1e-6"); + assert_eq!(format_prometheus_sample_value(1e-7), "1e-7"); + assert_eq!(format_prometheus_sample_value(1e21), "1e21"); + + // These known shortest-roundtrip tie cases have different text but + // remain numerically equivalent to Rust's representation. + for value in [ + f64::from_bits(0x42374876e8000400), + f64::from_bits(0x3ff0000800000000), + f64::from_bits(0x430a8e5672bc7312), + ] { + let ryu_output = format_prometheus_sample_value(value); + let std_output = value.to_string(); + assert_ne!(ryu_output, std_output); + assert_eq!(ryu_output.parse::().unwrap(), value); + assert_eq!(std_output.parse::().unwrap(), value); + } + } + + #[test] + fn format_prometheus_sample_value_preserves_nonfinite_values() { + assert_eq!(format_prometheus_sample_value(f64::NAN), "NaN"); + assert_eq!(format_prometheus_sample_value(f64::INFINITY), "inf"); + assert_eq!(format_prometheus_sample_value(f64::NEG_INFINITY), "-inf"); + + // Parsing a NaN does not preserve its payload bits, so NaN is checked + // by its required semantic spelling rather than by to_bits(). + assert!( + format_prometheus_sample_value(f64::NAN) + .parse::() + .unwrap() + .is_nan() + ); + } + #[test] fn matrix_response_preserves_ordinary_nan_and_filters_stale_markers() { let schema = Arc::new(Schema::new(vec![ @@ -604,10 +705,190 @@ mod tests { assert_eq!(series.len(), 1); assert_eq!( series[0].values, - vec![(1.0, "1".to_string()), (2.0, "NaN".to_string())] + vec![(1.0, "1.0".to_string()), (2.0, "NaN".to_string())] ); } + #[test] + fn record_batches_to_data_formats_values_with_ryu() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true), + ])); + let batch = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![ + 1_000, 2_000, 3_000, 4_000, 5_000, 6_000, + ])) as _, + Arc::new(Float64Vector::from(vec![ + Some(0.0), + Some(-0.0), + Some(1.25), + Some(1e30), + Some(1e-7), + Some(f64::MAX), + ])) as _, + ], + ) + .unwrap(); + let batches = RecordBatches::try_new(schema, vec![batch]).unwrap(); + + let response = + PrometheusJsonResponse::record_batches_to_data(batches, None, ValueType::Matrix) + .unwrap(); + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Matrix(series), + .. + }) = response + else { + panic!("expected matrix response"); + }; + + assert_eq!(series.len(), 1); + let input_values = [0.0, -0.0, 1.25, 1e30, 1e-7, f64::MAX]; + // Keep this expected-value generation independent from the production + // formatter while still asserting the Arrow/batch-to-Prometheus path. + let expected = input_values + .into_iter() + .enumerate() + .map(|(index, value)| { + let expected_value = if value.is_finite() { + let mut buffer = Buffer::new(); + buffer.format_finite(value).to_string() + } else { + value.to_string() + }; + ((index + 1) as f64, expected_value) + }) + .collect::>(); + assert_eq!(series[0].values, expected); + } + + #[test] + fn record_batches_to_data_preserves_infinity_output() { + // NaN and infinities use Rust's `f64::to_string()` output. + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true), + ])); + let batch = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![ + 1_000, 2_000, 3_000, + ])) as _, + Arc::new(Float64Vector::from(vec![ + Some(f64::INFINITY), + Some(f64::NEG_INFINITY), + Some(f64::NAN), + ])) as _, + ], + ) + .unwrap(); + let batches = RecordBatches::try_new(schema, vec![batch]).unwrap(); + + let response = + PrometheusJsonResponse::record_batches_to_data(batches, None, ValueType::Matrix) + .unwrap(); + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Matrix(series), + .. + }) = response + else { + panic!("expected matrix response"); + }; + + assert_eq!(series.len(), 1); + assert_eq!( + series[0].values, + vec![ + (1.0, "inf".to_string()), + (2.0, "-inf".to_string()), + (3.0, "NaN".to_string()), + ] + ); + } + + #[test] + fn record_batches_to_data_reuses_entries_for_clustered_series() { + // Rows are clustered by series (a, b, a, a, b, c): consecutive rows of + // the same series exercise the entry-reuse fast path, while series + // transitions fall back to the map lookup. The result must keep the + // first-occurrence order and accumulate values per series as before. + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), false), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true), + ])); + let batch = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_vec(vec![ + 1_000, 2_000, 3_000, 4_000, 5_000, 6_000, + ])) as _, + Arc::new(StringVector::from(vec![ + Some("a"), + Some("b"), + Some("a"), + Some("a"), + Some("b"), + Some("c"), + ])) as _, + Arc::new(Float64Vector::from(vec![ + Some(1.0), + Some(2.0), + Some(3.0), + Some(4.0), + Some(5.0), + Some(6.0), + ])) as _, + ], + ) + .unwrap(); + let batches = RecordBatches::try_new(schema, vec![batch]).unwrap(); + + let response = PrometheusJsonResponse::record_batches_to_data( + batches, + Some("metric".to_string()), + ValueType::Vector, + ) + .unwrap(); + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Vector(series), + .. + }) = response + else { + panic!("expected vector response"); + }; + + assert_eq!(series.len(), 3); + // Output order is first-occurrence order: a, b, c. + assert_eq!( + series + .iter() + .map(|series| series.metric["host"].as_str()) + .collect::>(), + vec!["a", "b", "c"] + ); + // Vector results keep the last sample of each series. + assert_eq!(series[0].value, Some((4.0, "4.0".to_string()))); + assert_eq!(series[1].value, Some((5.0, "5.0".to_string()))); + assert_eq!(series[2].value, Some((6.0, "6.0".to_string()))); + } + #[test] fn record_batches_to_data_preserves_mixed_float_and_histogram_rows() { let schema = Arc::new(Schema::new(vec![ @@ -713,7 +994,7 @@ mod tests { assert_eq!(series[0].metric["__name__"], "label_replace_repro"); assert_eq!(series[0].metric["host"], "server-01"); assert_eq!(series[0].metric["host_copy"], "server-01"); - assert_eq!(series[0].value, Some((1.0, "1".to_string()))); + assert_eq!(series[0].value, Some((1.0, "1.0".to_string()))); } #[test] diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index 9d0e13dddc9..737b4fb6c1d 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -1071,7 +1071,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) { ] .into_iter() .collect(), - value: Some((5.0, "1".to_string())), + value: Some((5.0, "1.0".to_string())), ..Default::default() }, PromSeriesVector { @@ -1081,7 +1081,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) { ] .into_iter() .collect(), - value: Some((5.0, "2".to_string())), + value: Some((5.0, "2.0".to_string())), ..Default::default() }, ] @@ -1133,7 +1133,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) { ] .into_iter() .collect(), - values: vec![(5.0, "1".to_string()), (10.0, "1".to_string())], + values: vec![(5.0, "1.0".to_string()), (10.0, "1.0".to_string())], ..Default::default() }, PromSeriesMatrix { @@ -1143,7 +1143,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) { ] .into_iter() .collect(), - values: vec![(5.0, "2".to_string()), (10.0, "2".to_string())], + values: vec![(5.0, "2.0".to_string()), (10.0, "2.0".to_string())], ..Default::default() }, ] diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 2a66c4353e7..7a787833376 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -935,7 +935,7 @@ pub async fn test_prom_http_api(store_type: StorageType) { assert_eq!( body.data, serde_json::from_value::( - json!({"resultType":"scalar","result":[1.0,"2"]}) + json!({"resultType":"scalar","result":[1.0,"2.0"]}) ) .unwrap() ); @@ -2967,6 +2967,8 @@ pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: Storag .await; assert_eq!(res.status(), StatusCode::OK); let body = res.json::().await; + // Finite vector sample values use the HTTP response's ryu contract, so + // integral f64 values retain their explicit `.0`; timestamps remain JSON numbers. assert_eq!( serde_json::to_string_pretty(&body).unwrap(), r#"{ @@ -2982,7 +2984,7 @@ pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: Storag }, "value": [ 4.0, - "6" + "6.0" ] } ] diff --git a/tests/perf/fixture-format.md b/tests/perf/fixture-format.md index ca60ae54d72..361f0ec79fd 100644 --- a/tests/perf/fixture-format.md +++ b/tests/perf/fixture-format.md @@ -58,6 +58,20 @@ warmup = 0 iterations = 1 ``` +### Query kinds + +`kind = "prom_http"` runs a Prometheus range query by POSTing form fields to +`/v1/prometheus/api/v1/query_range`. `query`, `start`, `end`, and `step` are +sent as form fields; `start` and `end` define the range, and `step` defines its +evaluation interval. `start`, `end`, and `step` are optional in the case model +and default to empty form values when omitted. The table database is sent as the +`db` URL query parameter, rather than as a form field. + +The measurement starts before the HTTP request is sent and ends after the full +response body is read and parsed as JSON. It therefore includes request send, +server-side Prometheus JSON response building, and response-body read and JSON +parsing; it is not server-side-only latency. + `series_layout = "round_robin"` advances the timestamp once per generated row and cycles series labels across rows. `series_layout = "timestamp_major"` writes all series for one timestamp before advancing to the next timestamp; use it for diff --git a/tests/perf/query_cases/prom_json_response/case.toml b/tests/perf/query_cases/prom_json_response/case.toml new file mode 100644 index 00000000000..17472cd5f32 --- /dev/null +++ b/tests/perf/query_cases/prom_json_response/case.toml @@ -0,0 +1,89 @@ +# Prometheus JSON response building benchmark (issue #8805). +# +# Exercises the serial per-request +# `PrometheusJsonResponse::record_batches_to_data` +# (src/servers/src/http/result/prometheus_resp.rs) path, ryu float formatting, +# and its per-series IndexMap entry reuse. +# +# 256 series x 481 evaluation timestamps (2h at 15s) = ~123k output samples +# per request. The primary `prom_http` query builds the Prometheus JSON +# response. The TQL ANALYZE control runs the same workload without building the +# Prometheus JSON response, to measure query-engine/scan variance. + +[case] +name = "prom_json_response" +description = "Prometheus HTTP API JSON response building for large range queries" + +[scenario] +kind = "direct_readable_sst" +seed = 8805 + +[[scenario.tables]] +database = "public" +name = "prom_json_response" +engine = "mito" +append_mode = true +sst_format = "flat" +primary_key = ["host", "instance"] +time_index = "ts" + +[[scenario.tables.columns]] +name = "host" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 16, prefix = "host" } + +[[scenario.tables.columns]] +name = "instance" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 256, prefix = "instance" } + +[[scenario.tables.columns]] +name = "value" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 1000.0 } + +[[scenario.tables.columns]] +name = "ts" +type = "TIMESTAMP(9)" +semantic = "timestamp" + +[scenario.layout] +regions = 1 +sst_count = 16 +rows_per_sst = 7680 +row_group_size = 1920 +series_count = 256 +start_unix_nanos = 1704067200000000000 +step_nanos = 15000000000 +time_range_layout = "non_overlapping_per_sst" +series_layout = "timestamp_major" + +# Primary case: Prometheus HTTP range query over the full 2h window. +# 256 series x 481 points -> ~123k samples built and serialized per request. +[[scenario.queries]] +name = "prom_range_2h" +kind = "prom_http" +query = "prom_json_response{host=~\"host.*\"}" +start = "1704067200" +end = "1704074400" +step = "15s" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +# Control: same workload through the SQL/TQL frontend path, which does not +# build the Prometheus JSON response. Isolates query-engine/scan variance. +[[scenario.queries]] +name = "tql_range_2h_control" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704067200, 1704074400, '15s') prom_json_response{host=~'host.*'}" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 From 04614175feff271fdf3608462c84e4ecef36b601 Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Tue, 25 Aug 2026 06:26:09 +0000 Subject: [PATCH 76/92] refactor: centralize native histogram encoding in common-query (#8945) * refactor: centralize native histogram encoding in common-query Signed-off-by: shuiyisong * fix: fmt Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/common/query/src/native_histogram.rs | 3 +++ .../query/src/native_histogram/encoding.rs} | 18 +++++++++++------- src/servers/src/lib.rs | 1 - src/servers/src/otlp/metrics.rs | 5 +++-- src/servers/src/prom_remote_write/v2.rs | 5 +++-- 5 files changed, 20 insertions(+), 12 deletions(-) rename src/{servers/src/native_histogram.rs => common/query/src/native_histogram/encoding.rs} (96%) diff --git a/src/common/query/src/native_histogram.rs b/src/common/query/src/native_histogram.rs index 1c7bb09868e..d8a559a2dc5 100644 --- a/src/common/query/src/native_histogram.rs +++ b/src/common/query/src/native_histogram.rs @@ -19,6 +19,8 @@ //! [`NativeHistogram`] is the query-time representation and therefore normalizes //! integer and floating-point payloads to absolute `f64` counts. +mod encoding; + use std::collections::BTreeMap; use std::sync::Arc; @@ -34,6 +36,7 @@ use datafusion::arrow::datatypes::{ use datafusion_common::{DataFusionError, Result as DfResult}; use datatypes::data_type::{ConcreteDataType, DataType}; use datatypes::types::{StructField, StructType}; +pub use encoding::{NativeHistogramError, encode_native_histogram, native_histogram_column_schema}; use once_cell::sync::Lazy; use crate::prelude::greptime_native_histogram; diff --git a/src/servers/src/native_histogram.rs b/src/common/query/src/native_histogram/encoding.rs similarity index 96% rename from src/servers/src/native_histogram.rs rename to src/common/query/src/native_histogram/encoding.rs index f1dae526993..be7b35ec454 100644 --- a/src/servers/src/native_histogram.rs +++ b/src/common/query/src/native_histogram/encoding.rs @@ -17,23 +17,27 @@ use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram}; use api::helper::ColumnDataTypeWrapper; use api::v1::value::ValueData; use api::v1::{ColumnSchema, ListValue, SemanticType, Value}; -use common_query::native_histogram::*; -use common_query::prelude::greptime_native_histogram; use snafu::{Snafu, ensure}; -const MAX_NATIVE_HISTOGRAM_SCHEMA: i32 = 8; +use crate::native_histogram::{ + CUSTOM_BUCKETS_SCHEMA, MAX_EXPONENTIAL_SCHEMA, NATIVE_HISTOGRAM_FIELD_NAMES, + exponential_overflow_bucket_index, native_histogram_value_type, +}; +use crate::prelude::greptime_native_histogram; + const MAX_REDUCIBLE_NATIVE_HISTOGRAM_SCHEMA: i32 = 52; +/// Error returned while validating or encoding a native histogram. #[derive(Debug, Snafu)] #[snafu(display("{message}"))] -pub(crate) struct NativeHistogramError { +pub struct NativeHistogramError { message: String, } type Result = std::result::Result; /// Returns the canonical column schema for a native histogram value. -pub(crate) fn native_histogram_column_schema() -> Result { +pub fn native_histogram_column_schema() -> Result { let (datatype, datatype_extension) = ColumnDataTypeWrapper::try_from(native_histogram_value_type().clone()) .map_err(|error| NativeHistogramError { @@ -51,7 +55,7 @@ pub(crate) fn native_histogram_column_schema() -> Result { } /// Validates and encodes a Prometheus histogram into the canonical Struct value. -pub(crate) fn encode_native_histogram(histogram: &Histogram) -> Result { +pub fn encode_native_histogram(histogram: &Histogram) -> Result { let uses_float_counts = native_histogram_uses_float_counts(histogram)?; validate_native_histogram(histogram, uses_float_counts)?; @@ -193,7 +197,7 @@ fn validate_native_histogram_schema(schema: i32) -> Result> { return Ok(Some(overflow_index)); } - if (MAX_NATIVE_HISTOGRAM_SCHEMA + 1..=MAX_REDUCIBLE_NATIVE_HISTOGRAM_SCHEMA).contains(&schema) { + if (MAX_EXPONENTIAL_SCHEMA + 1..=MAX_REDUCIBLE_NATIVE_HISTOGRAM_SCHEMA).contains(&schema) { Err(NativeHistogramError { message: format!("native histogram schema {schema} must be reduced before ingestion"), }) diff --git a/src/servers/src/lib.rs b/src/servers/src/lib.rs index 79fd392a1a9..2a3796aab21 100644 --- a/src/servers/src/lib.rs +++ b/src/servers/src/lib.rs @@ -34,7 +34,6 @@ pub mod interceptor; pub mod metrics; pub mod metrics_handler; pub mod mysql; -pub(crate) mod native_histogram; pub mod opentsdb; pub mod otel_arrow; pub mod otlp; diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index 332fa46f1d1..bd0bfb2ce77 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -20,7 +20,9 @@ use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram as Pr use api::v1::value::ValueData; use api::v1::{RowInsertRequests, SemanticType, Value}; use common_grpc::precision::Precision; -use common_query::native_histogram::native_histogram_value_type; +use common_query::native_histogram::{ + encode_native_histogram, native_histogram_column_schema, native_histogram_value_type, +}; use common_query::prelude::{GREPTIME_COUNT, greptime_timestamp, greptime_value}; use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; use lazy_static::lazy_static; @@ -34,7 +36,6 @@ use table::requests::{ }; use crate::error::{self, Result}; -use crate::native_histogram::{encode_native_histogram, native_histogram_column_schema}; use crate::otlp::trace::{KEY_SERVICE_INSTANCE_ID, KEY_SERVICE_NAME}; use crate::query_handler::MetricsIngestOutcome; use crate::row_writer::{self, MultiTableData, TableData}; diff --git a/src/servers/src/prom_remote_write/v2.rs b/src/servers/src/prom_remote_write/v2.rs index 846ef3ee6be..b388071106f 100644 --- a/src/servers/src/prom_remote_write/v2.rs +++ b/src/servers/src/prom_remote_write/v2.rs @@ -30,9 +30,11 @@ use api::v1::value::ValueData; use api::v1::{ColumnDataType, RowInsertRequest, Rows, SemanticType, Value}; use bytes::{Buf, Bytes}; use common_grpc::precision::Precision; -use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD; #[cfg(test)] use common_query::native_histogram::*; +use common_query::native_histogram::{ + NATIVE_HISTOGRAM_FIELD, encode_native_histogram, native_histogram_column_schema, +}; use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; use pipeline::{ContextOpt, ContextReq}; use prost::encoding::{ @@ -46,7 +48,6 @@ use table::requests::{ }; use crate::error::{self, Result}; -use crate::native_histogram::{encode_native_histogram, native_histogram_column_schema}; use crate::prom_remote_write::row_builder::PromCtx; use crate::prom_remote_write::validation::validate_label_name; use crate::prom_remote_write::{REMOTE_WRITE_V2_VERSION, try_decompress}; From a3a0db63b80a5ef00f52610c070f4b8d9707cd23 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Tue, 25 Aug 2026 07:01:10 +0000 Subject: [PATCH 77/92] feat(mito2): add series index searcher (#8926) * feat(mito2): add series index searcher Signed-off-by: evenyag * refactor(mito2): use parquet push decoder for series index Signed-off-by: evenyag * fix(mito2): handle evolved series index schemas Signed-off-by: evenyag --------- Signed-off-by: evenyag --- src/mito2/AGENTS.md | 2 +- src/mito2/src/read/series_candidate.rs | 19 +- src/mito2/src/read/series_reader.rs | 3 +- src/mito2/src/series_index.rs | 26 +- src/mito2/src/series_index/searcher.rs | 734 +++++++++++++++++++++++++ src/mito2/src/series_index/writer.rs | 8 +- src/mito2/src/sst/parquet/format.rs | 50 +- src/mito2/src/sst/parquet/prefilter.rs | 38 +- 8 files changed, 832 insertions(+), 48 deletions(-) create mode 100644 src/mito2/src/series_index/searcher.rs diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index e68ca586096..20ca9b53ca0 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -26,7 +26,7 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`. | `compaction` | `src/mito2/src/compaction/` | Compaction scheduler (`scheduler.rs` + `scheduler/`), TWCS picker, strict-window manual picker, compactor, memory control | | `access_layer` | `src/mito2/src/access_layer.rs` | SST read/write over the object store | | `sst` | `src/mito2/src/sst/` | Parquet format, file metadata, index layout | -| `series_index` | `src/mito2/src/series_index/` | Incremental writer for aggregate series index files | +| `series_index` | `src/mito2/src/series_index/` | Incremental writer and predicate searcher for aggregate series index files | | `read` | `src/mito2/src/read/` | `ScanRegion`, merge, dedup, projection, streaming | | `manifest` | `src/mito2/src/manifest/` | `RegionManifestManager`, manifest actions/edits | | `cache` | `src/mito2/src/cache.rs` | Write/file/page caches | diff --git a/src/mito2/src/read/series_candidate.rs b/src/mito2/src/read/series_candidate.rs index 2933327a38e..73591d4ffa3 100644 --- a/src/mito2/src/read/series_candidate.rs +++ b/src/mito2/src/read/series_candidate.rs @@ -30,7 +30,6 @@ use datatypes::arrow::compute::SortOptions; use datatypes::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::ConcreteDataType; -use futures::stream::BoxStream; use futures::{StreamExt, TryStreamExt}; use mito_codec::row_converter::{PrimaryKeyFilter, SparsePrimaryKeyCodec}; use snafu::{OptionExt, ResultExt, ensure}; @@ -51,6 +50,7 @@ use crate::read::range_cache::{ }; use crate::read::scan_region::StreamContext; use crate::read::scan_util::{PartitionMetrics, new_filter_metrics, scan_flat_mem_ranges}; +use crate::series_index::{METRIC_SERIES_ID_BATCH_SIZE, MetricSeriesId, MetricSeriesIdStream}; use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; use crate::sst::parquet::format::PrimaryKeyArray; use crate::sst::parquet::prefilter::{ @@ -59,17 +59,6 @@ use crate::sst::parquet::prefilter::{ use crate::sst::parquet::reader::ReaderMetrics; use crate::sst::parquet::row_group::ParquetFetchMetrics; -const CANDIDATE_SERIES_BATCH_SIZE: usize = 500; - -/// Identifies one series in a physical metric region. -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub(crate) struct MetricSeriesId { - pub(crate) table_id: u32, - pub(crate) tsid: u64, -} - -pub(crate) type MetricSeriesIdStream = BoxStream<'static, Result>>; - /// Builds candidate metric series from the ranges assigned to a [`SeriesScan`](super::series_scan::SeriesScan). pub(crate) struct SeriesCandidateScanner { stream_ctx: Arc, @@ -537,7 +526,7 @@ fn decode_metric_series( let codec = SparsePrimaryKeyCodec::new(&metadata); Ok(Box::pin(try_stream! { let mut last_series = None; - let mut output = Vec::with_capacity(CANDIDATE_SERIES_BATCH_SIZE); + let mut output = Vec::with_capacity(METRIC_SERIES_ID_BATCH_SIZE); while let Some(batch) = input.try_next().await? { let array = batch .column(0) @@ -556,10 +545,10 @@ fn decode_metric_series( } last_series = Some(series); output.push(series); - if output.len() == CANDIDATE_SERIES_BATCH_SIZE { + if output.len() == METRIC_SERIES_ID_BATCH_SIZE { yield std::mem::replace( &mut output, - Vec::with_capacity(CANDIDATE_SERIES_BATCH_SIZE), + Vec::with_capacity(METRIC_SERIES_ID_BATCH_SIZE), ); } } diff --git a/src/mito2/src/read/series_reader.rs b/src/mito2/src/read/series_reader.rs index 0a3e008f27a..15e47924323 100644 --- a/src/mito2/src/read/series_reader.rs +++ b/src/mito2/src/read/series_reader.rs @@ -40,7 +40,8 @@ use crate::read::scan_util::{ should_split_flat_batches_for_merge, }; use crate::read::seq_scan::SeqScan; -use crate::read::series_candidate::{MetricSeriesId, validate_metric_metadata}; +use crate::read::series_candidate::validate_metric_metadata; +use crate::series_index::MetricSeriesId; use crate::sst::parquet::DEFAULT_READ_BATCH_SIZE; use crate::sst::parquet::flat_format::primary_key_column_index; use crate::sst::parquet::prefilter::prefilter_flat_batch_by_primary_key; diff --git a/src/mito2/src/series_index.rs b/src/mito2/src/series_index.rs index 21ced00cf64..e35769bec4c 100644 --- a/src/mito2/src/series_index.rs +++ b/src/mito2/src/series_index.rs @@ -12,10 +12,34 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Series index writer. +//! Series index writer and searcher. +mod searcher; mod writer; +use futures::stream::BoxStream; +pub use searcher::SeriesIndexSearcher; pub use writer::{ SeriesIndexWriter, SeriesIndexWriterMetrics, SeriesIndexWriterOptions, series_index_schema, }; + +use crate::error::Result; + +pub(crate) const MIN_TS_COLUMN: &str = "__series_min_ts"; +pub(crate) const MAX_TS_COLUMN: &str = "__series_max_ts"; +pub(crate) const ROW_COUNT_COLUMN: &str = "__series_row_count"; +pub(crate) const TABLE_ID_COLUMN: &str = "__table_id"; +pub(crate) const TSID_COLUMN: &str = "__tsid"; +pub(crate) const METRIC_SERIES_ID_BATCH_SIZE: usize = 500; + +/// Identifies one series in a physical metric region. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct MetricSeriesId { + /// Logical table ID inside the physical metric region. + pub table_id: u32, + /// Time-series ID inside the logical table. + pub tsid: u64, +} + +/// Stream of bounded batches of matching metric-series IDs. +pub type MetricSeriesIdStream = BoxStream<'static, Result>>; diff --git a/src/mito2/src/series_index/searcher.rs b/src/mito2/src/series_index/searcher.rs new file mode 100644 index 00000000000..a303bcaebda --- /dev/null +++ b/src/mito2/src/series_index/searcher.rs @@ -0,0 +1,734 @@ +// Copyright 2023 Greptime Team +// +// 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. + +use std::collections::HashSet; +use std::ops::Range; +use std::sync::Arc; + +use api::v1::SemanticType; +use async_stream::try_stream; +use bytes::Bytes; +use common_recordbatch::filter::SimpleFilterEvaluator; +use common_time::range::TimestampRange; +use datafusion_common::pruning::PruningStatistics; +use datafusion_common::{Column, ScalarValue}; +use datafusion_expr::{Expr, col, lit}; +use datatypes::arrow::array::{ArrayRef, BooleanArray, UInt32Array, UInt64Array}; +use datatypes::arrow::buffer::BooleanBuffer; +use datatypes::arrow::datatypes::{DataType, SchemaRef}; +use object_store::ObjectStore; +use parquet::DecodeResult; +use parquet::arrow::ProjectionMask; +use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; +use parquet::arrow::push_decoder::ParquetPushDecoderBuilder; +use parquet::file::metadata::{ParquetMetaData, RowGroupMetaData}; +use snafu::{OptionExt, ResultExt, ensure}; +use store_api::metadata::RegionMetadataRef; +use table::predicate::Predicate; + +use crate::error::{ + InvalidMetaSnafu, InvalidRecordBatchSnafu, OpenDalSnafu, ReadParquetSnafu, RecordBatchSnafu, + Result, UnexpectedSnafu, +}; +use crate::series_index::{ + MAX_TS_COLUMN, METRIC_SERIES_ID_BATCH_SIZE, MIN_TS_COLUMN, MetricSeriesId, + MetricSeriesIdStream, ROW_COUNT_COLUMN, TABLE_ID_COLUMN, TSID_COLUMN, series_index_schema, +}; +use crate::sst::parquet::format::{column_null_counts, column_values_by_type}; +use crate::sst::parquet::helper::fetch_byte_ranges; +use crate::sst::parquet::metadata::MetadataLoader; +use crate::sst::parquet::prefilter::simple_tag_filters; +use crate::sst::parquet::reader::MetadataCacheMetrics; + +#[derive(Clone)] +struct SeriesIndexRangeFetcher { + object_store: ObjectStore, +} + +impl SeriesIndexRangeFetcher { + async fn fetch(&self, path: &str, ranges: &[Range]) -> Result> { + fetch_byte_ranges(path, self.object_store.clone(), ranges) + .await + .context(OpenDalSnafu) + } +} + +#[derive(Clone)] +struct SeriesIndexMetadataProvider { + object_store: ObjectStore, +} + +impl SeriesIndexMetadataProvider { + async fn load(&self, path: &str) -> Result> { + let mut metrics = MetadataCacheMetrics::default(); + MetadataLoader::new(self.object_store.clone(), path, 0) + .load(&mut metrics) + .await + .map(Arc::new) + } +} + +/// Searches a series-index file for metric series matching query predicates. +pub struct SeriesIndexSearcher { + range_fetcher: SeriesIndexRangeFetcher, + metadata_provider: SeriesIndexMetadataProvider, + filters: Vec<(Expr, SimpleFilterEvaluator)>, + empty_time_range: bool, +} + +impl SeriesIndexSearcher { + /// Creates a searcher reusable across series-index files of `metadata`. + pub fn try_new( + metadata: RegionMetadataRef, + object_store: ObjectStore, + predicate: Option<&Predicate>, + time_range: Option, + ) -> Result { + // Keep search-time metadata validation identical to the writer. + series_index_schema(&metadata)?; + + let mut filters = simple_tag_filters(&metadata, None, predicate); + let (empty_time_range, time_exprs) = time_range_filters(&metadata, time_range)?; + for expr in time_exprs { + let filter = SimpleFilterEvaluator::try_new(&expr).context(UnexpectedSnafu { + reason: "failed to build an internal series-index time filter", + })?; + filters.push((expr, filter)); + } + + Ok(Self { + range_fetcher: SeriesIndexRangeFetcher { + object_store: object_store.clone(), + }, + metadata_provider: SeriesIndexMetadataProvider { object_store }, + filters, + empty_time_range, + }) + } + + /// Searches `path` and returns sorted batches of matching metric-series IDs. + pub async fn search(&self, path: &str) -> Result { + if self.empty_time_range { + return Ok(Box::pin(futures::stream::empty())); + } + + let parquet_metadata = self.metadata_provider.load(path).await?; + let arrow_metadata = + ArrowReaderMetadata::try_new(parquet_metadata, ArrowReaderOptions::new()) + .with_context(|_| ReadParquetSnafu { + path: path.to_string(), + })?; + validate_index_schema(arrow_metadata.schema())?; + + // An older index file may not contain tags added by schema evolution. + // Ignore filters on those tags to preserve a conservative candidate set. + let (pruning_predicate, filters) = self.filters_for_schema(arrow_metadata.schema()); + let row_groups = row_groups_to_read( + arrow_metadata.metadata().row_groups(), + arrow_metadata.schema().clone(), + &pruning_predicate, + ); + let projection = projection_mask( + arrow_metadata.parquet_schema(), + arrow_metadata.schema(), + &filters, + )?; + let mut decoder = ParquetPushDecoderBuilder::new_with_metadata(arrow_metadata) + .with_row_groups(row_groups) + .with_projection(projection) + .build() + .with_context(|_| ReadParquetSnafu { + path: path.to_string(), + })?; + let path = path.to_string(); + let range_fetcher = self.range_fetcher.clone(); + + Ok(Box::pin(try_stream! { + let mut last_series = None; + let mut output = Vec::with_capacity(METRIC_SERIES_ID_BATCH_SIZE); + loop { + let batch = match decoder + .try_decode() + .with_context(|_| ReadParquetSnafu { path: path.clone() })? + { + DecodeResult::NeedsData(ranges) => { + let data = range_fetcher.fetch(&path, &ranges).await?; + decoder + .push_ranges(ranges, data) + .with_context(|_| ReadParquetSnafu { path: path.clone() })?; + continue; + } + DecodeResult::Data(batch) => batch, + DecodeResult::Finished => break, + }; + + let mut mask = BooleanBuffer::new_set(batch.num_rows()); + for filter in &filters { + let column = column(&batch, filter.column_name())?; + let evaluated = filter.evaluate_array(column).context(RecordBatchSnafu)?; + mask = &mask & &evaluated; + } + + let table_ids = column(&batch, TABLE_ID_COLUMN)? + .as_any() + .downcast_ref::() + .context(InvalidRecordBatchSnafu { + reason: "series index __table_id is not UInt32", + })?; + let tsids = column(&batch, TSID_COLUMN)? + .as_any() + .downcast_ref::() + .context(InvalidRecordBatchSnafu { + reason: "series index __tsid is not UInt64", + })?; + + for (row, matched) in mask.iter().enumerate() { + if !matched { + continue; + } + let series = MetricSeriesId { + table_id: table_ids.value(row), + tsid: tsids.value(row), + }; + if last_series == Some(series) { + continue; + } + last_series = Some(series); + output.push(series); + if output.len() == METRIC_SERIES_ID_BATCH_SIZE { + yield std::mem::replace( + &mut output, + Vec::with_capacity(METRIC_SERIES_ID_BATCH_SIZE), + ); + } + } + } + if !output.is_empty() { + yield output; + } + })) + } + + fn filters_for_schema(&self, schema: &SchemaRef) -> (Predicate, Vec) { + let (exprs, filters): (Vec<_>, Vec<_>) = self + .filters + .iter() + .filter(|(_, filter)| schema.field_with_name(filter.column_name()).is_ok()) + .cloned() + .unzip(); + (Predicate::new(exprs), filters) + } +} + +fn time_range_filters( + metadata: &RegionMetadataRef, + time_range: Option, +) -> Result<(bool, Vec)> { + let Some(time_range) = time_range else { + return Ok((false, Vec::new())); + }; + if time_range.is_empty() { + return Ok((true, Vec::new())); + } + + let time_index = metadata.time_index_column(); + ensure!( + time_index.semantic_type == SemanticType::Timestamp, + InvalidMetaSnafu { + reason: "series index metadata has no timestamp time index", + } + ); + let timestamp_type = + time_index + .column_schema + .data_type + .as_timestamp() + .context(InvalidMetaSnafu { + reason: "series index time index is not a timestamp", + })?; + let unit = timestamp_type.unit(); + let mut exprs = Vec::with_capacity(2); + // A series overlaps [start, end) only if its maximum is at least start. + // Round start up so a series ending before an unaligned start is pruned. + if let Some(start) = time_range + .start() + .and_then(|start| start.convert_to_ceil(unit)) + { + exprs.push(col(MAX_TS_COLUMN).gt_eq(lit(start.value()))); + } + // A series overlaps [start, end) only if its minimum is less than end. + // Round the exclusive end up to avoid pruning the containing unit interval. + if let Some(end) = time_range.end().and_then(|end| end.convert_to_ceil(unit)) { + exprs.push(col(MIN_TS_COLUMN).lt(lit(end.value()))); + } + Ok((false, exprs)) +} + +fn validate_index_schema(schema: &SchemaRef) -> Result<()> { + for (name, data_type) in [ + (MIN_TS_COLUMN, DataType::Int64), + (MAX_TS_COLUMN, DataType::Int64), + (ROW_COUNT_COLUMN, DataType::UInt64), + (TABLE_ID_COLUMN, DataType::UInt32), + (TSID_COLUMN, DataType::UInt64), + ] { + let field = schema + .field_with_name(name) + .ok() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!("series index is missing internal column {name}"), + })?; + ensure!( + field.data_type() == &data_type && !field.is_nullable(), + InvalidRecordBatchSnafu { + reason: format!( + "series index internal column {name} must be non-nullable {data_type:?}, got {:?}", + field.data_type() + ), + } + ); + } + Ok(()) +} + +fn projection_mask( + parquet_schema: &parquet::schema::types::SchemaDescriptor, + arrow_schema: &SchemaRef, + filters: &[SimpleFilterEvaluator], +) -> Result { + let mut indices = HashSet::new(); + for name in [TABLE_ID_COLUMN, TSID_COLUMN] { + let index = arrow_schema + .index_of(name) + .ok() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!("series index is missing internal column {name}"), + })?; + indices.insert(index); + } + for filter in filters { + let index = arrow_schema + .index_of(filter.column_name()) + .ok() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!( + "series index is missing predicate column {}", + filter.column_name() + ), + })?; + indices.insert(index); + } + Ok(ProjectionMask::roots(parquet_schema, indices)) +} + +fn column<'a>( + batch: &'a datatypes::arrow::record_batch::RecordBatch, + name: &str, +) -> Result<&'a ArrayRef> { + let index = batch + .schema() + .index_of(name) + .ok() + .with_context(|| InvalidRecordBatchSnafu { + reason: format!("series index batch is missing column {name}"), + })?; + Ok(batch.column(index)) +} + +struct SeriesIndexPruningStats<'a> { + row_groups: &'a [RowGroupMetaData], + schema: SchemaRef, +} + +impl PruningStatistics for SeriesIndexPruningStats<'_> { + fn min_values(&self, column: &Column) -> Option { + self.column_values(column, true) + } + + fn max_values(&self, column: &Column) -> Option { + self.column_values(column, false) + } + + fn num_containers(&self) -> usize { + self.row_groups.len() + } + + fn null_counts(&self, column: &Column) -> Option { + let column_index = self.schema.index_of(&column.name).ok()?; + column_null_counts(self.row_groups, column_index) + } + + fn row_counts(&self, _column: &Column) -> Option { + None + } + + fn contained(&self, _column: &Column, _values: &HashSet) -> Option { + None + } +} + +impl SeriesIndexPruningStats<'_> { + fn column_values(&self, column: &Column, is_min: bool) -> Option { + let column_index = self.schema.index_of(&column.name).ok()?; + let data_type = self.schema.field(column_index).data_type(); + column_values_by_type(self.row_groups, data_type, column_index, is_min) + } +} + +fn row_groups_to_read( + row_groups: &[RowGroupMetaData], + schema: SchemaRef, + predicate: &Predicate, +) -> Vec { + let stats = SeriesIndexPruningStats { row_groups, schema }; + predicate + .prune_with_stats(&stats, &stats.schema) + .into_iter() + .enumerate() + .filter_map(|(row_group, keep)| keep.then_some(row_group)) + .collect() +} + +#[cfg(test)] +mod tests { + use datafusion_expr::{col, lit}; + use datatypes::arrow::array::{BinaryArray, TimestampMillisecondArray, UInt8Array}; + use datatypes::arrow::datatypes::{Field, Schema}; + use datatypes::arrow::record_batch::RecordBatch; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::ColumnSchema; + use futures::TryStreamExt; + use object_store::services::Memory; + use store_api::codec::PrimaryKeyEncoding; + use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder}; + + use super::*; + use crate::series_index::{SeriesIndexWriter, SeriesIndexWriterOptions}; + use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding}; + + fn object_store() -> ObjectStore { + ObjectStore::new(Memory::default()).unwrap().finish() + } + + fn flat_batch(primary_keys: &[Vec], timestamps: &[i64]) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(datatypes::arrow::datatypes::TimeUnit::Millisecond, None), + false, + ), + Field::new("__primary_key", DataType::Binary, false), + Field::new("__sequence", DataType::UInt64, false), + Field::new("__op_type", DataType::UInt8, false), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(TimestampMillisecondArray::from(timestamps.to_vec())), + Arc::new(BinaryArray::from_iter_values( + primary_keys.iter().map(Vec::as_slice), + )), + Arc::new(UInt64Array::from(vec![1; timestamps.len()])), + Arc::new(UInt8Array::from(vec![0; timestamps.len()])), + ], + ) + .unwrap() + } + + async fn write_index( + metadata: RegionMetadataRef, + object_store: ObjectStore, + path: &str, + rows: &[(u32, u64, &str, &str, i64)], + row_group_size: usize, + ) { + let primary_keys = rows + .iter() + .map(|(table_id, tsid, tag_0, tag_1, _)| { + new_sparse_primary_key(&[*tag_0, *tag_1], &metadata, *table_id, *tsid) + }) + .collect::>(); + let timestamps = rows.iter().map(|row| row.4).collect::>(); + let mut writer = SeriesIndexWriter::try_new( + metadata, + object_store, + path, + SeriesIndexWriterOptions { row_group_size }, + ) + .await + .unwrap(); + writer + .write(&flat_batch(&primary_keys, ×tamps)) + .await + .unwrap(); + writer.finish().await.unwrap(); + } + + async fn collect_ids(stream: MetricSeriesIdStream) -> Vec { + stream + .try_collect::>() + .await + .unwrap() + .into_iter() + .flatten() + .collect() + } + + #[tokio::test] + async fn search_applies_candidate_tag_filters_and_time_overlap() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let object_store = object_store(); + let path = "search.parquet"; + write_index( + metadata.clone(), + object_store.clone(), + path, + &[ + (1, 10, "a", "x", 10), + (1, 20, "b", "x", 20), + (1, 30, "a", "y", 30), + ], + 2, + ) + .await; + + // The field filter is not available in the series index and is ignored, + // matching candidate-primary-key filter behavior. + let predicate = Predicate::new(vec![ + col("tag_0").eq(lit("a")), + col("field_0").gt(lit(0_u64)), + ]); + let time_range = TimestampRange::new( + common_time::Timestamp::new_millisecond(20), + common_time::Timestamp::new_millisecond(31), + ) + .unwrap(); + let searcher = SeriesIndexSearcher::try_new( + metadata.clone(), + object_store.clone(), + Some(&predicate), + Some(time_range), + ) + .unwrap(); + let ids = collect_ids(searcher.search(path).await.unwrap()).await; + assert_eq!( + ids, + vec![MetricSeriesId { + table_id: 1, + tsid: 30 + }] + ); + + // Both bounds fall between millisecond ticks. Rounding the inclusive + // start and exclusive end upward leaves only the 30 ms series. + let time_range = TimestampRange::new( + common_time::Timestamp::new_microsecond(20_001), + common_time::Timestamp::new_microsecond(30_001), + ) + .unwrap(); + let searcher = SeriesIndexSearcher::try_new( + metadata.clone(), + object_store.clone(), + None, + Some(time_range), + ) + .unwrap(); + let ids = collect_ids(searcher.search(path).await.unwrap()).await; + assert_eq!( + ids, + vec![MetricSeriesId { + table_id: 1, + tsid: 30 + }] + ); + + // The stored maximum equal to the query start intersects, while the + // stored minimum equal to the exclusive query end does not. + let time_range = TimestampRange::new( + common_time::Timestamp::new_millisecond(20), + common_time::Timestamp::new_millisecond(30), + ) + .unwrap(); + let searcher = + SeriesIndexSearcher::try_new(metadata, object_store, None, Some(time_range)).unwrap(); + let ids = collect_ids(searcher.search(path).await.unwrap()).await; + assert_eq!( + ids, + vec![MetricSeriesId { + table_id: 1, + tsid: 20 + }] + ); + } + + #[tokio::test] + async fn search_skips_filters_for_columns_missing_from_older_index() { + let old_metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let object_store = object_store(); + let path = "schema-evolution.parquet"; + write_index( + old_metadata.clone(), + object_store.clone(), + path, + &[ + (1, 10, "a", "x", 10), + (1, 20, "b", "x", 20), + (1, 30, "a", "y", 30), + ], + 2, + ) + .await; + + let mut builder = RegionMetadataBuilder::from_existing(old_metadata.as_ref().clone()); + builder.push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new("tag_2", ConcreteDataType::string_datatype(), true), + semantic_type: SemanticType::Tag, + column_id: 4, + }); + let mut primary_key = old_metadata.primary_key.clone(); + primary_key.push(4); + builder.primary_key(primary_key); + let current_metadata = Arc::new(builder.build().unwrap()); + + let predicate = + Predicate::new(vec![col("tag_0").eq(lit("a")), col("tag_2").eq(lit("new"))]); + let searcher = + SeriesIndexSearcher::try_new(current_metadata, object_store, Some(&predicate), None) + .unwrap(); + let ids = collect_ids(searcher.search(path).await.unwrap()).await; + assert_eq!( + ids, + vec![ + MetricSeriesId { + table_id: 1, + tsid: 10, + }, + MetricSeriesId { + table_id: 1, + tsid: 30, + }, + ] + ); + } + + #[tokio::test] + async fn search_streams_fixed_size_batches() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let object_store = object_store(); + let rows = (0..501_u64) + .map(|tsid| (1, tsid, "a", "x", tsid as i64)) + .collect::>(); + write_index( + metadata.clone(), + object_store.clone(), + "batching.parquet", + &rows, + 100, + ) + .await; + + let searcher = SeriesIndexSearcher::try_new(metadata, object_store, None, None).unwrap(); + let batches = searcher + .search("batching.parquet") + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(batches.iter().map(Vec::len).collect::>(), [500, 1]); + assert_eq!( + batches[0][0], + MetricSeriesId { + table_id: 1, + tsid: 0 + } + ); + assert_eq!( + batches[1][0], + MetricSeriesId { + table_id: 1, + tsid: 500 + } + ); + } + + #[tokio::test] + async fn search_prunes_row_groups_and_empty_ranges() { + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let object_store = object_store(); + let path = "pruning.parquet"; + write_index( + metadata.clone(), + object_store.clone(), + path, + &[ + (1, 0, "a", "x", 0), + (1, 1, "b", "x", 1), + (1, 2, "m", "x", 2), + (1, 3, "m", "x", 3), + (1, 4, "y", "x", 4), + (1, 5, "z", "x", 5), + ], + 2, + ) + .await; + + let predicate = Predicate::new(vec![col("tag_0").eq(lit("m"))]); + let searcher = SeriesIndexSearcher::try_new( + metadata.clone(), + object_store.clone(), + Some(&predicate), + None, + ) + .unwrap(); + let parquet_metadata = searcher.metadata_provider.load(path).await.unwrap(); + let arrow_metadata = + ArrowReaderMetadata::try_new(parquet_metadata, ArrowReaderOptions::new()).unwrap(); + let (pruning_predicate, _) = searcher.filters_for_schema(arrow_metadata.schema()); + assert_eq!( + row_groups_to_read( + arrow_metadata.metadata().row_groups(), + arrow_metadata.schema().clone(), + &pruning_predicate, + ), + vec![1] + ); + + let empty = SeriesIndexSearcher::try_new( + metadata, + object_store, + None, + Some(TimestampRange::empty()), + ) + .unwrap(); + assert!( + empty + .search("does-not-need-to-exist.parquet") + .await + .unwrap() + .try_collect::>() + .await + .unwrap() + .is_empty() + ); + } +} diff --git a/src/mito2/src/series_index/writer.rs b/src/mito2/src/series_index/writer.rs index ce5b7405f21..9ebba289fb4 100644 --- a/src/mito2/src/series_index/writer.rs +++ b/src/mito2/src/series_index/writer.rs @@ -45,15 +45,13 @@ use crate::error::{ DecodeSnafu, InvalidMetaSnafu, InvalidRecordBatchSnafu, NewRecordBatchSnafu, OpenDalSnafu, Result, UnexpectedSnafu, WriteParquetSnafu, }; +use crate::series_index::{ + MAX_TS_COLUMN, MIN_TS_COLUMN, ROW_COUNT_COLUMN, TABLE_ID_COLUMN, TSID_COLUMN, +}; use crate::sst::parquet::DEFAULT_ROW_GROUP_SIZE; use crate::sst::parquet::flat_format::{primary_key_column_index, time_index_column_index}; use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY}; -const MIN_TS_COLUMN: &str = "__series_min_ts"; -const MAX_TS_COLUMN: &str = "__series_max_ts"; -const ROW_COUNT_COLUMN: &str = "__series_row_count"; -const TABLE_ID_COLUMN: &str = "__table_id"; -const TSID_COLUMN: &str = "__tsid"; const WRITE_BATCH_SIZE: usize = 1024; type ParquetWriter = AsyncArrowWriter; diff --git a/src/mito2/src/sst/parquet/format.rs b/src/mito2/src/sst/parquet/format.rs index f3f0a34a682..4ccdcb41089 100644 --- a/src/mito2/src/sst/parquet/format.rs +++ b/src/mito2/src/sst/parquet/format.rs @@ -36,7 +36,7 @@ use datafusion_common::ScalarValue; use datatypes::arrow::array::{ ArrayRef, BinaryArray, BinaryDictionaryBuilder, DictionaryArray, UInt64Array, }; -use datatypes::arrow::datatypes::{SchemaRef, UInt32Type}; +use datatypes::arrow::datatypes::{DataType as ArrowDataType, SchemaRef, UInt32Type}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::DataType; use datatypes::types::json_type::JsonNativeType; @@ -136,12 +136,22 @@ pub(crate) fn column_values( column_index: usize, is_min: bool, ) -> Option { - let null_scalar: ScalarValue = column - .column_schema - .data_type - .as_arrow_type() - .try_into() - .ok()?; + column_values_by_type( + row_groups, + &column.column_schema.data_type.as_arrow_type(), + column_index, + is_min, + ) +} + +/// Returns min/max values of a parquet column with the given Arrow data type. +pub(crate) fn column_values_by_type( + row_groups: &[impl Borrow], + data_type: &ArrowDataType, + column_index: usize, + is_min: bool, +) -> Option { + let null_scalar: ScalarValue = data_type.try_into().ok()?; let scalar_values = row_groups .iter() .map(|meta| { @@ -152,16 +162,22 @@ pub(crate) fn column_values( } else { *s.max_opt()? }))), - Statistics::Int32(s) => Some(ScalarValue::Int32(Some(if is_min { - *s.min_opt()? - } else { - *s.max_opt()? - }))), - Statistics::Int64(s) => Some(ScalarValue::Int64(Some(if is_min { - *s.min_opt()? - } else { - *s.max_opt()? - }))), + Statistics::Int32(s) => { + let value = if is_min { *s.min_opt()? } else { *s.max_opt()? }; + if data_type == &ArrowDataType::UInt32 { + Some(ScalarValue::UInt32(Some(value as u32))) + } else { + Some(ScalarValue::Int32(Some(value))) + } + } + Statistics::Int64(s) => { + let value = if is_min { *s.min_opt()? } else { *s.max_opt()? }; + if data_type == &ArrowDataType::UInt64 { + Some(ScalarValue::UInt64(Some(value as u64))) + } else { + Some(ScalarValue::Int64(Some(value))) + } + } Statistics::Int96(_) => None, Statistics::Float(s) => Some(ScalarValue::Float32(Some(if is_min { *s.min_opt()? diff --git a/src/mito2/src/sst/parquet/prefilter.rs b/src/mito2/src/sst/parquet/prefilter.rs index fcaaae4a54d..3af3b20d9df 100644 --- a/src/mito2/src/sst/parquet/prefilter.rs +++ b/src/mito2/src/sst/parquet/prefilter.rs @@ -24,6 +24,7 @@ use std::sync::Arc; use api::v1::SemanticType; use common_recordbatch::filter::SimpleFilterEvaluator; +use datafusion_expr::Expr; use datatypes::arrow::array::{Array, BinaryArray, BooleanArray, BooleanBufferBuilder}; use datatypes::arrow::buffer::BooleanBuffer; use datatypes::arrow::datatypes::SchemaRef; @@ -272,15 +273,9 @@ pub(crate) fn build_primary_key_filter( expected_metadata: Option<&RegionMetadata>, predicate: Option<&Predicate>, ) -> Option { - let filters = predicate + let filters = simple_tag_filters(sst_metadata, expected_metadata, predicate) .into_iter() - .flat_map(|predicate| predicate.exprs()) - .filter_map(|expr| SimpleFilterContext::new_opt(sst_metadata, expected_metadata, expr)) - .filter_map(|filter_ctx| { - (filter_ctx.semantic_type() == SemanticType::Tag) - .then(|| filter_ctx.filter().as_filter().cloned()) - .flatten() - }) + .map(|(_, filter)| filter) .collect::>(); if filters.is_empty() { return None; @@ -291,6 +286,33 @@ pub(crate) fn build_primary_key_filter( Some(CachedPrimaryKeyFilter::new(filter)) } +/// Extracts simple tag filters that can be applied to encoded primary keys or series indexes. +pub(crate) fn simple_tag_filters( + sst_metadata: &RegionMetadataRef, + expected_metadata: Option<&RegionMetadata>, + predicate: Option<&Predicate>, +) -> Vec<(Expr, SimpleFilterEvaluator)> { + predicate + .into_iter() + .flat_map(|predicate| predicate.exprs()) + .filter_map(|expr| { + SimpleFilterContext::new_opt(sst_metadata, expected_metadata, expr) + .map(|filter_ctx| (expr, filter_ctx)) + }) + .filter_map(|(expr, filter_ctx)| { + (filter_ctx.semantic_type() == SemanticType::Tag) + .then(|| { + filter_ctx + .filter() + .as_filter() + .cloned() + .map(|filter| (expr.clone(), filter)) + }) + .flatten() + }) + .collect() +} + /// How the parquet reader should apply each predicate. /// /// The reader runs in two phases. Predicates routed into `prefilter_builder` From 9bfabca72b56b15cdb9872a0fe9defe31c896262 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 25 Aug 2026 09:02:12 +0000 Subject: [PATCH 78/92] chore(deps): bump postgres-protocol from 0.6.8 to 0.6.12 (#8948) Bumps [postgres-protocol](https://github.com/rust-postgres/rust-postgres) from 0.6.8 to 0.6.12. - [Release notes](https://github.com/rust-postgres/rust-postgres/releases) - [Commits](https://github.com/rust-postgres/rust-postgres/compare/postgres-protocol-v0.6.8...postgres-protocol-v0.6.12) --- updated-dependencies: - dependency-name: postgres-protocol dependency-version: 0.6.12 dependency-type: indirect ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- Cargo.lock | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 534f484d4fc..ec477e5c9a3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10906,19 +10906,19 @@ dependencies = [ [[package]] name = "postgres-protocol" -version = "0.6.8" +version = "0.6.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76ff0abab4a9b844b93ef7b81f1efc0a366062aaef2cd702c76256b5dc075c54" +checksum = "08808e3c483c46e999108051c78334f473d5adb59d78bb80a1268c7e6aa6c514" dependencies = [ "base64 0.22.1", "byteorder", "bytes", "fallible-iterator", - "hmac 0.12.1", - "md-5 0.10.6", + "hmac 0.13.0", + "md-5 0.11.0", "memchr", - "rand 0.9.4", - "sha2 0.10.9", + "rand 0.10.1", + "sha2 0.11.0", "stringprep", ] @@ -11299,7 +11299,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4" dependencies = [ "bytes", - "heck 0.5.0", + "heck 0.4.1", "itertools 0.12.1", "log", "multimap", @@ -11319,7 +11319,7 @@ version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ - "heck 0.5.0", + "heck 0.4.1", "itertools 0.14.0", "log", "multimap", @@ -13785,7 +13785,7 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1961e2ef424c1424204d3a5d6975f934f56b6d50ff5732382d84ebf460e147f7" dependencies = [ - "heck 0.5.0", + "heck 0.4.1", "proc-macro2", "quote", "syn 2.0.117", From 046862f067847c2f37d40f4c3601354c79b7f339 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Tue, 25 Aug 2026 09:22:53 +0000 Subject: [PATCH 79/92] feat(cmd): improve parquet rewrite fidelity and scanbench output (#8947) * feat(cmd): preserve parquet rewrite metadata settings Signed-off-by: evenyag * feat(cmd): print scanner metrics in verbose scanbench Signed-off-by: evenyag * fix(cmd): preserve SST truncation settings in property dumps Signed-off-by: evenyag --------- Signed-off-by: evenyag --- docs/parquet-tools.md | 4 +- src/cmd/src/datanode/parquet_meta.rs | 72 +++++++++++++++++++++---- src/cmd/src/datanode/parquet_rewrite.rs | 65 +++++++++++++++++++--- src/cmd/src/datanode/scanbench.rs | 19 +++++++ 4 files changed, 144 insertions(+), 16 deletions(-) diff --git a/docs/parquet-tools.md b/docs/parquet-tools.md index 8528f969431..9dfe404626e 100644 --- a/docs/parquet-tools.md +++ b/docs/parquet-tools.md @@ -89,7 +89,9 @@ Supported compression names are `uncompressed`, `snappy`, `gzip`, `lzo`, The `[writer]` table also accepts `compression_level`, `data_page_size_limit`, `data_page_row_count_limit`, and -`dictionary_page_size_limit`. A `[[columns]]` entry overrides dictionary, +`dictionary_page_size_limit`. It also accepts +`column_index_truncate_length` and `statistics_truncate_length`; set either to +zero to disable that truncation. A `[[columns]]` entry overrides dictionary, compression, compression level, or encoding for its column path. Unknown TOML fields are rejected. Compression levels are supported for `gzip`, `brotli`, and `zstd`. A column-level `compression_level` without `compression` inherits the diff --git a/src/cmd/src/datanode/parquet_meta.rs b/src/cmd/src/datanode/parquet_meta.rs index 425da45e3c0..bbd89a5bb20 100644 --- a/src/cmd/src/datanode/parquet_meta.rs +++ b/src/cmd/src/datanode/parquet_meta.rs @@ -17,6 +17,7 @@ use std::path::{Path, PathBuf}; use clap::{Parser, ValueEnum}; use parquet::file::metadata::ParquetMetaData; +use parquet::file::page_index::offset_index::PageLocation; use serde::Serialize; use snafu::ResultExt; @@ -78,6 +79,7 @@ struct ColumnChunkMetaView { dictionary_page_offset: Option, dictionary_page_bytes: Option, data_pages: Option, + max_data_page_row_count: Option, has_statistics: bool, column_index_offset: Option, column_index_length: Option, @@ -123,10 +125,14 @@ fn build_file_meta_view(path: &Path, metadata: &ParquetMetaData) -> FileMetaView .iter() .enumerate() .map(|(column_idx, column)| { - let data_pages = offset_index + let page_locations = offset_index .and_then(|index| index.get(row_group_idx)) .and_then(|columns| columns.get(column_idx)) - .map(|index| index.page_locations().len()); + .map(|index| index.page_locations()); + let data_pages = page_locations.map(|locations| locations.len()); + let max_data_page_row_count = page_locations.and_then(|locations| { + max_data_page_row_count(locations, row_group.num_rows()) + }); let dictionary_page_bytes = dictionary_page_bytes( column.dictionary_page_offset(), column.data_page_offset(), @@ -149,6 +155,7 @@ fn build_file_meta_view(path: &Path, metadata: &ParquetMetaData) -> FileMetaView dictionary_page_offset: column.dictionary_page_offset(), dictionary_page_bytes, data_pages, + max_data_page_row_count, has_statistics: column.statistics().is_some(), column_index_offset: column.column_index_offset(), column_index_length: column.column_index_length(), @@ -219,7 +226,7 @@ fn print_meta_text(view: &FileMetaView) { ); for column in &row_group.columns { println!( - " column[{}] path={}: type={}, compression={}, encodings=[{}], values={}, uncompressed_size={}, compressed_size={}, compression_ratio={}, data_pages={}, dictionary_page_offset={}, dictionary_page_bytes={}, data_page_offset={}, statistics={}, column_index={}/{}, offset_index={}/{}, bloom_filter={}/{}", + " column[{}] path={}: type={}, compression={}, encodings=[{}], values={}, uncompressed_size={}, compressed_size={}, compression_ratio={}, data_pages={}, max_data_page_row_count={}, dictionary_page_offset={}, dictionary_page_bytes={}, data_page_offset={}, statistics={}, column_index={}/{}, offset_index={}/{}, bloom_filter={}/{}", column.index, column.path, column.physical_type, @@ -230,6 +237,7 @@ fn print_meta_text(view: &FileMetaView) { column.compressed_size, format_ratio(column.compression_ratio), format_optional_usize(column.data_pages), + format_optional_i64(column.max_data_page_row_count), format_optional_i64(column.dictionary_page_offset), column .dictionary_page_bytes @@ -248,6 +256,33 @@ fn print_meta_text(view: &FileMetaView) { } } +fn max_data_page_row_count( + page_locations: &[PageLocation], + row_group_row_count: i64, +) -> Option { + if page_locations.is_empty() { + return Some(0); + } + if page_locations.first()?.first_row_index != 0 { + return None; + } + + let mut max_row_count = 0; + for locations in page_locations.windows(2) { + let row_count = locations[1].first_row_index - locations[0].first_row_index; + if row_count < 0 { + return None; + } + max_row_count = max_row_count.max(row_count); + } + + let last_page_row_count = row_group_row_count - page_locations.last()?.first_row_index; + if last_page_row_count < 0 { + return None; + } + Some(max_row_count.max(last_page_row_count)) +} + fn dictionary_page_bytes( dictionary_page_offset: Option, data_page_offset: i64, @@ -301,7 +336,7 @@ mod tests { fn test_meta_view_unknown_page_count() { let dir = tempdir().unwrap(); let path = dir.path().join("input.parquet"); - write_test_parquet(&path); + write_test_parquet(&path, None); let metadata = load_local_parquet_metadata(&path).unwrap(); let view = build_file_meta_view(&path, &metadata); @@ -313,7 +348,21 @@ mod tests { assert_eq!(view.row_groups[0].columns.len(), 2); } - fn write_test_parquet(path: &Path) { + #[test] + fn test_meta_view_reports_max_data_page_row_count() { + let dir = tempdir().unwrap(); + let path = dir.path().join("input.parquet"); + write_test_parquet(&path, Some(2)); + + let metadata = load_local_parquet_metadata(&path).unwrap(); + let view = build_file_meta_view(&path, &metadata); + + for column in &view.row_groups[0].columns { + assert_eq!(Some(2), column.max_data_page_row_count); + } + } + + fn write_test_parquet(path: &Path, data_page_row_count_limit: Option) { let schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::Int32, false), Field::new("host", DataType::Utf8, false), @@ -326,12 +375,17 @@ mod tests { ], ) .unwrap(); - let props = WriterProperties::builder() - .set_key_value_metadata(Some(vec![KeyValue::new( + let mut props = + WriterProperties::builder().set_key_value_metadata(Some(vec![KeyValue::new( "greptime:test".to_string(), "value".to_string(), - )])) - .build(); + )])); + if let Some(limit) = data_page_row_count_limit { + props = props + .set_write_batch_size(limit) + .set_data_page_row_count_limit(limit); + } + let props = props.build(); let mut writer = ArrowWriter::try_new(File::create(path).unwrap(), schema, Some(props)).unwrap(); writer.write(&batch).unwrap(); diff --git a/src/cmd/src/datanode/parquet_rewrite.rs b/src/cmd/src/datanode/parquet_rewrite.rs index 47346aac889..b04db93edcb 100644 --- a/src/cmd/src/datanode/parquet_rewrite.rs +++ b/src/cmd/src/datanode/parquet_rewrite.rs @@ -77,6 +77,8 @@ struct WriterConfig { data_page_size_limit: Option, data_page_row_count_limit: Option, dictionary_page_size_limit: Option, + column_index_truncate_length: Option, + statistics_truncate_length: Option, } #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] @@ -189,6 +191,15 @@ impl ParquetRewriteCommand { } fn infer_rewrite_properties(metadata: &ParquetMetaData) -> RewriteProperties { + let is_greptime_sst = metadata + .file_metadata() + .key_value_metadata() + .is_some_and(|key_values| { + key_values + .iter() + .any(|key_value| key_value.key == mito2::sst::parquet::PARQUET_METADATA_KEY) + }); + let truncate_length = is_greptime_sst.then_some(0); let first_column = metadata .row_groups() .first() @@ -234,6 +245,8 @@ fn infer_rewrite_properties(metadata: &ParquetMetaData) -> RewriteProperties { data_page_size_limit: None, data_page_row_count_limit: None, dictionary_page_size_limit: None, + column_index_truncate_length: truncate_length, + statistics_truncate_length: truncate_length, }, columns, } @@ -287,6 +300,16 @@ fn build_writer_properties( if let Some(dictionary_page_size_limit) = config.writer.dictionary_page_size_limit { builder = builder.set_dictionary_page_size_limit(dictionary_page_size_limit); } + if let Some(column_index_truncate_length) = config.writer.column_index_truncate_length { + builder = builder.set_column_index_truncate_length( + (column_index_truncate_length > 0).then_some(column_index_truncate_length), + ); + } + if let Some(statistics_truncate_length) = config.writer.statistics_truncate_length { + builder = builder.set_statistics_truncate_length( + (statistics_truncate_length > 0).then_some(statistics_truncate_length), + ); + } for column in config.columns { if column.path.is_empty() { @@ -492,7 +515,7 @@ mod tests { let dir = tempdir().unwrap(); let input = dir.path().join("input.parquet"); let output = dir.path().join("output.parquet"); - write_test_parquet(&input, true); + write_test_parquet(&input, true, Some("greptime:test")); let metadata = load_local_parquet_metadata(&input).unwrap(); let mut properties = infer_rewrite_properties(&metadata); @@ -528,6 +551,37 @@ mod tests { ); } + #[test] + fn test_infer_truncation_settings_for_greptime_sst() { + let dir = tempdir().unwrap(); + let greptime_sst = dir.path().join("greptime.parquet"); + let generic_parquet = dir.path().join("generic.parquet"); + write_test_parquet( + &greptime_sst, + true, + Some(mito2::sst::parquet::PARQUET_METADATA_KEY), + ); + write_test_parquet(&generic_parquet, true, None); + + let metadata = load_local_parquet_metadata(&greptime_sst).unwrap(); + let config = infer_rewrite_properties(&metadata); + assert_eq!(Some(0), config.writer.column_index_truncate_length); + assert_eq!(Some(0), config.writer.statistics_truncate_length); + + let serialized = toml::to_string(&config).unwrap(); + assert!(serialized.contains("column_index_truncate_length = 0")); + assert!(serialized.contains("statistics_truncate_length = 0")); + + let properties = build_writer_properties(config, None).unwrap(); + assert_eq!(None, properties.column_index_truncate_length()); + assert_eq!(None, properties.statistics_truncate_length()); + + let metadata = load_local_parquet_metadata(&generic_parquet).unwrap(); + let config = infer_rewrite_properties(&metadata); + assert_eq!(None, config.writer.column_index_truncate_length); + assert_eq!(None, config.writer.statistics_truncate_length); + } + #[test] fn test_column_compression_level_inherits_writer_compression() { let path = ColumnPath::new(vec!["host".to_string()]); @@ -551,7 +605,7 @@ mod tests { ); } - fn write_test_parquet(path: &Path, dictionary_enabled: bool) { + fn write_test_parquet(path: &Path, dictionary_enabled: bool, metadata_key: Option<&str>) { let schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::Int32, false), Field::new("host", DataType::Utf8, false), @@ -564,12 +618,11 @@ mod tests { ], ) .unwrap(); + let key_value_metadata = + metadata_key.map(|key| vec![KeyValue::new(key.to_string(), "value".to_string())]); let props = WriterProperties::builder() .set_dictionary_enabled(dictionary_enabled) - .set_key_value_metadata(Some(vec![KeyValue::new( - "greptime:test".to_string(), - "value".to_string(), - )])) + .set_key_value_metadata(key_value_metadata) .build(); let mut writer = ArrowWriter::try_new(File::create(path).unwrap(), schema, Some(props)).unwrap(); diff --git a/src/cmd/src/datanode/scanbench.rs b/src/cmd/src/datanode/scanbench.rs index 1dbee025a7f..fd6c08fe6db 100644 --- a/src/cmd/src/datanode/scanbench.rs +++ b/src/cmd/src/datanode/scanbench.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::HashMap; +use std::fmt; use std::path::PathBuf; use std::sync::Arc; use std::time::Instant; @@ -31,6 +32,7 @@ use datafusion::logical_expr::{BinaryExpr, Expr as DfExpr, ExprSchemable, Operat use datafusion_common::tree_node::{Transformed, TreeNodeRewriter}; use datafusion_common::{DFSchemaRef, ScalarValue, ToDFSchema}; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion_physical_plan::{DisplayAs, DisplayFormatType}; use datatypes::arrow::compute; use futures::StreamExt; use futures::stream::FuturesUnordered; @@ -61,6 +63,15 @@ use crate::datanode::tool_util::{ }; use crate::error; +/// Displays a scanner using DataFusion's verbose explain format. +struct VerboseScannerDisplay<'a, T: ?Sized>(&'a T); + +impl fmt::Display for VerboseScannerDisplay<'_, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt_as(DisplayFormatType::Verbose, f) + } +} + /// Scan benchmark command - benchmarks scanning a region directly from storage. #[derive(Debug, Parser)] pub struct ScanbenchCommand { @@ -688,6 +699,14 @@ impl ScanbenchCommand { format_bytes(total_estimated_size), ); + if self.verbose { + println!( + " {} Scanner explain: {}", + "ℹ".blue(), + VerboseScannerDisplay(scanner.as_ref()) + ); + } + // Start profiling after the first iteration (warmup) if pprof_after_warmup is set #[cfg(unix)] if iteration == 0 From 8a473c5bf0d3cb88d0bd328c14f33d5f2fea2224 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 25 Aug 2026 10:30:59 +0000 Subject: [PATCH 80/92] fix(meta): allow manual migration from offline datanodes (#8934) * fix(meta): allow migration from offline datanodes Signed-off-by: WenyXu * test: fix offline migration event actor Signed-off-by: WenyXu * test: read migration routes from metadata Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/meta-srv/src/service/procedure.rs | 3 +- tests-integration/tests/region_migration.rs | 135 +++++++++++++++++++- 2 files changed, 134 insertions(+), 4 deletions(-) diff --git a/src/meta-srv/src/service/procedure.rs b/src/meta-srv/src/service/procedure.rs index df27e2cb32b..be4eedc2ab7 100644 --- a/src/meta-srv/src/service/procedure.rs +++ b/src/meta-srv/src/service/procedure.rs @@ -26,6 +26,7 @@ use api::v1::meta::{ use common_event_recorder::{PersistentEventContext, ProcedureEventInput}; use common_meta::key::TableMetadataManagerRef; use common_meta::key::table_name::TableNameKey; +use common_meta::peer::Peer; use common_meta::procedure_executor::ExecutorContext; use common_meta::rpc::ddl::{ CREATE_DATABASE_CREATOR_EXTENSION_KEY, CREATE_DATABASE_CREATOR_METADATA_KEY, @@ -181,7 +182,7 @@ impl procedure_service_server::ProcedureService for Metasrv { let from_peer = self .lookup_datanode_peer(from_peer) .await? - .context(error::PeerUnavailableSnafu { peer_id: from_peer })?; + .unwrap_or_else(|| Peer::empty(from_peer)); let to_peer = self .lookup_datanode_peer(to_peer) .await? diff --git a/tests-integration/tests/region_migration.rs b/tests-integration/tests/region_migration.rs index 17a58de2eed..0df42ee1e88 100644 --- a/tests-integration/tests/region_migration.rs +++ b/tests-integration/tests/region_migration.rs @@ -25,8 +25,10 @@ use common_event_recorder::{ DEFAULT_EVENTS_TABLE_NAME, DEFAULT_FLUSH_INTERVAL_SECONDS, EVENTS_TABLE_TIMESTAMP_COLUMN_NAME, EVENTS_TABLE_TYPE_COLUMN_NAME, PersistentEventContext, TriggerReason, }; +use common_meta::distributed_time_constants::default_distributed_time_constants; use common_meta::key::{RegionDistribution, RegionRoleSet, TableMetadataManagerRef}; use common_meta::peer::Peer; +use common_meta::rpc::store::BatchDeleteRequest; use common_procedure::ProcedureContext; use common_procedure::event::{ EVENTS_TABLE_PROCEDURE_ID_COLUMN_NAME, EVENTS_TABLE_PROCEDURE_STATE_COLUMN_NAME, @@ -46,6 +48,7 @@ use futures::future::BoxFuture; use meta_srv::error; use meta_srv::error::Result as MetaResult; use meta_srv::event::region_migration::REGION_MIGRATION_EVENT_TYPE; +use meta_srv::key::DatanodeLeaseKey; use meta_srv::metasrv::SelectorContext; use meta_srv::procedure::region_migration::{ RegionMigrationProcedureTask, RegionMigrationTriggerReason, @@ -102,6 +105,7 @@ macro_rules! region_migration_tests { test_region_migration, test_region_migration_by_sql, + test_region_migration_with_offline_source_by_sql, test_region_migration_multiple_regions, test_region_migration_all_regions, test_region_migration_incorrect_from_peer, @@ -419,8 +423,24 @@ pub async fn test_metric_table_region_migration_by_sql( .await; } -/// A naive region migration test by SQL function +/// A naive region migration test by SQL function. pub async fn test_region_migration_by_sql(store_type: StorageType, endpoints: Vec) { + test_region_migration_by_sql_inner(store_type, endpoints, false).await; +} + +/// A region migration test by SQL function with an offline source datanode. +pub async fn test_region_migration_with_offline_source_by_sql( + store_type: StorageType, + endpoints: Vec, +) { + test_region_migration_by_sql_inner(store_type, endpoints, true).await; +} + +async fn test_region_migration_by_sql_inner( + store_type: StorageType, + endpoints: Vec, + simulate_offline_source: bool, +) { let cluster_name = "test_region_migration"; let peer_factory = |id| Peer { id, @@ -437,7 +457,7 @@ pub async fn test_region_migration_by_sql(store_type: StorageType, endpoints: Ve peer_factory(2), peer_factory(3), ])); - let cluster = builder + let mut cluster = builder .with_datanodes(datanodes as u32) .with_store_config(store_config) .with_datanode_wal_config(DatanodeWalConfig::Kafka(DatanodeKafkaConfig { @@ -463,6 +483,7 @@ pub async fn test_region_migration_by_sql(store_type: StorageType, endpoints: Ve .with_meta_selector(const_selector.clone()) .build(true) .await; + let table_metadata_manager = cluster.metasrv.table_metadata_manager().clone(); let (actor_db, _actor_grpc_server) = setup_authenticated_grpc_database( cluster.fe_instance().clone(), PROCEDURE_ACTOR, @@ -487,7 +508,7 @@ pub async fn test_region_migration_by_sql(store_type: StorageType, endpoints: Ve let old_distribution = distribution.clone(); // Selecting target of region migration. - let region_migration_manager = cluster.metasrv.region_migration_manager(); + let region_migration_manager = cluster.metasrv.region_migration_manager().clone(); let (from_peer_id, from_regions) = distribution.pop_first().unwrap(); info!( "Selecting from peer: {from_peer_id}, and regions: {:?}", @@ -545,6 +566,114 @@ pub async fn test_region_migration_by_sql(store_type: StorageType, endpoints: Ve // Asserts the writes. assert_values(cluster.fe_instance()).await; + if simulate_offline_source { + let mut expected_distribution = + find_region_distribution(&table_metadata_manager, table_id).await; + let (offline_from_peer_id, offline_region_number) = expected_distribution + .iter() + .find(|(peer_id, regions)| { + **peer_id != from_peer_id + && **peer_id != to_peer_id + && !regions.leader_regions.is_empty() + }) + .map(|(peer_id, regions)| (*peer_id, regions.leader_regions[0])) + .unwrap(); + let offline_region_id = RegionId::new(table_id, offline_region_number); + + // Simulates scale-in: the source datanode stops and its lease disappears. + cluster + .datanode_instances + .get_mut(&offline_from_peer_id) + .unwrap() + .shutdown() + .await + .unwrap(); + let source_lease_key: Vec = DatanodeLeaseKey { + node_id: offline_from_peer_id, + } + .try_into() + .unwrap(); + cluster + .metasrv + .in_memory() + .batch_delete(BatchDeleteRequest { + keys: vec![source_lease_key], + prev_kv: false, + }) + .await + .unwrap(); + + let offline_procedure_id = trigger_migration_by_sql( + &cluster, + offline_region_id.as_u64(), + offline_from_peer_id, + from_peer_id, + ) + .await; + let frontend = cluster.fe_instance().clone(); + let procedure_id_for_closure = offline_procedure_id.clone(); + wait_condition( + default_distributed_time_constants().region_lease + Duration::from_secs(10), + Box::pin(async move { + loop { + let state = query_procedure_by_sql(&frontend, &procedure_id_for_closure).await; + if state == "{\"status\":\"Done\"}" { + info!("Offline-source migration done: {state}"); + break; + } + info!("Offline-source migration not finished: {state}"); + tokio::time::sleep(Duration::from_millis(200)).await; + } + }), + ) + .await; + + check_region_migration_events_system_table( + cluster.fe_instance(), + &offline_procedure_id, + offline_region_id.as_u64(), + offline_from_peer_id, + from_peer_id, + Some("greptime"), + ) + .await; + + let remove_offline_source = { + let source_regions = expected_distribution + .get_mut(&offline_from_peer_id) + .unwrap(); + source_regions + .leader_regions + .retain(|region_number| *region_number != offline_region_number); + source_regions.leader_regions.is_empty() && source_regions.follower_regions.is_empty() + }; + if remove_offline_source { + expected_distribution.remove(&offline_from_peer_id); + } + let target_regions = expected_distribution.entry(from_peer_id).or_default(); + target_regions.add_leader_region(offline_region_number); + target_regions.sort(); + + let table_metadata_manager = table_metadata_manager.clone(); + wait_condition( + Duration::from_secs(10), + Box::pin(async move { + loop { + let distribution = + find_region_distribution(&table_metadata_manager, table_id).await; + if distribution == expected_distribution { + break; + } + info!("Offline-source migration has unexpected distribution: {distribution:?}"); + tokio::time::sleep(Duration::from_millis(200)).await; + } + }), + ) + .await; + + assert_values(cluster.fe_instance()).await; + } + // Triggers again. let err = region_migration_manager .submit_procedure( From 28398138ec03d154524a8217672d463f28e7f466 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Tue, 25 Aug 2026 10:31:06 +0000 Subject: [PATCH 81/92] fix(flight): bound DoGet response wait (#8943) * fix(flight): defer datanode query initialization Signed-off-by: WenyXu * fix(client): retain Flight stream peer context Signed-off-by: WenyXu * fix(client): improve Flight stream diagnostics Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/client/src/region.rs | 116 +++++++++++++++++++++----- src/datanode/src/region_server.rs | 20 +++-- src/servers/src/grpc/flight.rs | 6 +- src/servers/src/grpc/flight/stream.rs | 104 +++++++++++++++++++---- tests-integration/src/grpc/flight.rs | 103 ++++++++++++++++++++++- 5 files changed, 299 insertions(+), 50 deletions(-) diff --git a/src/client/src/region.rs b/src/client/src/region.rs index f1730a93c8e..2bb5f7dd0a6 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::sync::Arc; +use std::time::Duration; use api::region::RegionResponse; use api::v1::ResponseHeader; @@ -24,7 +25,7 @@ use arc_swap::ArcSwapOption; use arrow_flight::Ticket; use async_stream::stream; use async_trait::async_trait; -use common_error::ext::BoxedError; +use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; use common_grpc::flight::{FlightDecoder, FlightMessage}; use common_meta::error::{self as meta_error, Result as MetaResult}; @@ -49,6 +50,8 @@ use crate::error::{ use crate::flight::decode_flight_data; use crate::{Client, metrics}; +const FLIGHT_DO_GET_TIMEOUT: Duration = Duration::from_secs(10); + #[derive(Debug)] pub struct RegionRequester { client: Client, @@ -109,20 +112,24 @@ impl RegionRequester { let mut flight_client = self .client .make_flight_client(self.send_compression, self.accept_compression)?; + // Limit Flight DoGet response time without limiting query stream execution. + let addr = flight_client.addr().to_string(); + let mut request = tonic::Request::new(ticket); + request.set_timeout(FLIGHT_DO_GET_TIMEOUT); let response = flight_client .mut_inner() - .do_get(ticket) + .do_get(request) .await .or_else(|e| { let tonic_code = e.code(); let e: error::Error = e.into(); error!( e; "Failed to do Flight get, addr: {}, code: {}", - flight_client.addr(), + addr, tonic_code ); Err(BoxedError::new(e)).with_context(|_| FlightGetSnafu { - addr: flight_client.addr().to_string(), + addr: addr.clone(), tonic_code, }) })?; @@ -133,7 +140,7 @@ impl RegionRequester { let flight_message_stream = flight_data_stream .filter_map(move |flight_data| decode_flight_data(&mut decoder, flight_data)); - recordbatches_from_flight_message_stream(flight_message_stream).await + recordbatches_from_flight_message_stream(addr, flight_message_stream).await } async fn handle_inner(&self, request: RegionRequest) -> Result { @@ -195,6 +202,7 @@ impl RegionRequester { } async fn recordbatches_from_flight_message_stream( + addr: String, mut flight_message_stream: S, ) -> Result where @@ -204,13 +212,17 @@ where return IllegalFlightMessagesSnafu { reason: "Expect the response not to be empty", } - .fail(); + .fail() + .map_err(|error| flight_stream_error(&addr, error)); }; - let FlightMessage::Schema(schema) = first_flight_message? else { + let FlightMessage::Schema(schema) = + first_flight_message.map_err(|e| flight_stream_error(&addr, e))? + else { return IllegalFlightMessagesSnafu { reason: "Expect schema to be the first flight message", } - .fail(); + .fail() + .map_err(|error| flight_stream_error(&addr, error)); }; let metrics = Arc::new(ArcSwapOption::from(None)); @@ -221,6 +233,7 @@ where let schema = Arc::new(datatypes::schema::Schema::try_from(schema).context(error::ConvertSchemaSnafu)?); let schema_cloned = schema.clone(); + let stream_addr = addr.clone(); let stream = Box::pin(stream!({ let _span = tracing_context.attach(common_telemetry::tracing::info_span!( "poll_flight_data_stream" @@ -240,7 +253,8 @@ where let flight_message = match flight_message_item { Some(Ok(message)) => message, Some(Err(e)) => { - yield Err(BoxedError::new(e)).context(ExternalSnafu); + yield Err(BoxedError::new(flight_stream_error(&stream_addr, e))) + .context(ExternalSnafu); break; } None => break, @@ -273,7 +287,8 @@ where break; } Err(e) => { - yield Err(BoxedError::new(e)).context(ExternalSnafu); + yield Err(BoxedError::new(flight_stream_error(&stream_addr, e))) + .context(ExternalSnafu); break; } } @@ -312,6 +327,23 @@ where Ok(Box::pin(record_batch_stream)) } +fn flight_stream_error(addr: &str, error: error::Error) -> error::Error { + let tonic_code = error.tonic_code().unwrap_or(tonic::Code::Unknown); + if error.status_code().should_log_error() { + error!( + error; "Failed to receive Flight data, addr: {}, code: {}", + addr, + tonic_code + ); + } + + error::Error::FlightGet { + addr: addr.to_string(), + tonic_code, + source: BoxedError::new(error), + } +} + pub fn build_remote_dyn_filter_update_request( query_id: impl Into, update: RemoteDynFilterUpdate, @@ -389,6 +421,44 @@ mod test { use super::*; use crate::Error::{self, IllegalDatabaseResponse, Server}; + #[test] + fn test_flight_stream_error_preserves_peer_address() { + let error = flight_stream_error( + "127.0.0.1:4001", + tonic::Status::unavailable("datanode unavailable").into(), + ); + + assert!(matches!( + error, + error::Error::FlightGet { + addr, + tonic_code: tonic::Code::Unavailable, + .. + } if addr == "127.0.0.1:4001" + )); + } + + #[tokio::test] + async fn test_empty_flight_stream_preserves_peer_address() { + let Err(error) = recordbatches_from_flight_message_stream( + "127.0.0.1:4001".to_string(), + stream::empty::>(), + ) + .await + else { + panic!("expected empty Flight stream to fail"); + }; + + assert!(matches!( + error, + error::Error::FlightGet { + addr, + tonic_code: tonic::Code::Unknown, + .. + } if addr == "127.0.0.1:4001" + )); + } + fn test_schema() -> Arc { Arc::new(Schema::new(vec![ColumnSchema::new( "v", @@ -509,11 +579,14 @@ mod test { ) .unwrap(); - let mut recordbatches = recordbatches_from_flight_message_stream(stream::iter(vec![ - Ok(FlightMessage::Schema(schema.arrow_schema().clone())), - Ok(FlightMessage::Metrics(test_metrics_json())), - Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), - ])) + let mut recordbatches = recordbatches_from_flight_message_stream( + "test-peer".to_string(), + stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::Metrics(test_metrics_json())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + ]), + ) .await .unwrap(); @@ -528,11 +601,14 @@ mod test { #[tokio::test] async fn test_record_batch_stream_exposes_error_after_pre_batch_metrics() { let schema = test_schema(); - let mut recordbatches = recordbatches_from_flight_message_stream(stream::iter(vec![ - Ok(FlightMessage::Schema(schema.arrow_schema().clone())), - Ok(FlightMessage::Metrics(test_metrics_json())), - Err(Error::from(Status::internal("boom after metrics"))), - ])) + let mut recordbatches = recordbatches_from_flight_message_stream( + "test-peer".to_string(), + stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::Metrics(test_metrics_json())), + Err(Error::from(Status::internal("boom after metrics"))), + ]), + ) .await .unwrap(); diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 57facfd53ba..0d9398b0889 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -66,7 +66,9 @@ use servers::error::{ self as servers_error, ExecuteGrpcRequestSnafu, Result as ServerResult, SuspendedSnafu, }; use servers::grpc::FlightCompression; -use servers::grpc::flight::{FlightCraft, FlightRecordBatchStream, TonicStream}; +use servers::grpc::flight::{ + FlightCraft, FlightRecordBatchSource, FlightRecordBatchStream, TonicStream, +}; use servers::grpc::region_server::RegionServerHandler; use session::context::{ FLIGHT_METRICS_HEARTBEAT_INTERVAL, QueryContext, QueryContextBuilder, QueryContextRef, @@ -976,13 +978,19 @@ impl FlightCraft for RegionServer { .map(|h| Arc::new(QueryContext::from(h))) .unwrap_or(QueryContext::arc()); - let result = self - .handle_remote_read(request, query_ctx.clone()) - .trace(tracing_context.attach(info_span!("RegionServer::handle_read"))) - .await?; + let region_server = self.clone(); + let initializer_query_ctx = query_ctx.clone(); + let initializer_tracing_context = tracing_context.clone(); + let initializer = async move { + region_server + .handle_remote_read(request, initializer_query_ctx) + .trace(initializer_tracing_context.attach(info_span!("RegionServer::handle_read"))) + .await + .map_err(Into::into) + }; let stream = Box::pin(FlightRecordBatchStream::new( - result, + FlightRecordBatchSource::initializer(initializer), tracing_context, self.flight_compression, query_ctx, diff --git a/src/servers/src/grpc/flight.rs b/src/servers/src/grpc/flight.rs index 97a579d952c..e6f35e17b29 100644 --- a/src/servers/src/grpc/flight.rs +++ b/src/servers/src/grpc/flight.rs @@ -53,7 +53,7 @@ use tokio_stream::wrappers::ReceiverStream; use tonic::{Request, Response, Status, Streaming}; use crate::error::{InvalidParameterSnafu, Result, ToJsonSnafu}; -pub use crate::grpc::flight::stream::FlightRecordBatchStream; +pub use crate::grpc::flight::stream::{FlightRecordBatchSource, FlightRecordBatchStream}; use crate::grpc::greptime_handler::{ GreptimeRequestHandler, create_query_context, get_request_type, }; @@ -583,7 +583,7 @@ fn to_flight_data_stream( match output.data { OutputData::Stream(stream) => { let stream = FlightRecordBatchStream::new( - stream, + FlightRecordBatchSource::RecordBatches(stream), tracing_context, flight_compression, query_ctx, @@ -592,7 +592,7 @@ fn to_flight_data_stream( } OutputData::RecordBatches(x) => { let stream = FlightRecordBatchStream::new( - x.as_stream(), + FlightRecordBatchSource::RecordBatches(x.as_stream()), tracing_context, flight_compression, query_ctx, diff --git a/src/servers/src/grpc/flight/stream.rs b/src/servers/src/grpc/flight/stream.rs index c222c72e1cd..319a797d9ca 100644 --- a/src/servers/src/grpc/flight/stream.rs +++ b/src/servers/src/grpc/flight/stream.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::VecDeque; +use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; use std::time::{Duration, Instant}; @@ -40,6 +41,30 @@ use crate::error; use crate::grpc::FlightCompression; use crate::grpc::flight::TonicResult; +pub type FlightRecordBatchStreamInitializer = Pin< + Box< + dyn Future> + + Send + + 'static, + >, +>; + +pub enum FlightRecordBatchSource { + RecordBatches(SendableRecordBatchStream), + Initializer(FlightRecordBatchStreamInitializer), +} + +impl FlightRecordBatchSource { + pub fn initializer(initializer: F) -> Self + where + F: Future> + + Send + + 'static, + { + Self::Initializer(Box::pin(initializer)) + } +} + /// Metrics collector for Flight stream with RAII logging pattern struct StreamMetrics { send_schema_duration: Duration, @@ -136,7 +161,7 @@ impl FlightRecordBatchStream { } pub fn new( - recordbatches: SendableRecordBatchStream, + source: FlightRecordBatchSource, tracing_context: TracingContext, compression: FlightCompression, query_ctx: QueryContextRef, @@ -148,17 +173,43 @@ impl FlightRecordBatchStream { .remote_query_id() .zip(query_ctx.extension(SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY)) .is_some_and(|(remote_query_id, capability)| capability == remote_query_id); - let (tx, rx) = mpsc::channel::>(1); - let join_handle = common_runtime::spawn_global(async move { - Self::flight_data_stream( - recordbatches, - tx, - should_send_partial_metrics, - can_send_metrics_before_batch, - ) - .trace(tracing_context.attach(info_span!("flight_data_stream"))) - .await - }); + let (mut tx, rx) = mpsc::channel::>(1); + let source_type = match &source { + FlightRecordBatchSource::RecordBatches(_) => "record_batches", + FlightRecordBatchSource::Initializer(_) => "initializer", + }; + let initializer_tracing_context = tracing_context.clone(); + let join_handle = common_runtime::spawn_global( + async move { + let recordbatches = async move { + match source { + FlightRecordBatchSource::RecordBatches(recordbatches) => Ok(recordbatches), + FlightRecordBatchSource::Initializer(initializer) => initializer.await, + } + } + .trace( + initializer_tracing_context + .attach(info_span!("flight_data_stream_init", source_type)), + ) + .await; + + match recordbatches { + Ok(recordbatches) => { + Self::flight_data_stream( + recordbatches, + tx, + should_send_partial_metrics, + can_send_metrics_before_batch, + ) + .await; + } + Err(status) => { + let _ = tx.send(Err(status)).await; + } + } + } + .trace(tracing_context.attach(info_span!("flight_data_stream"))), + ); let encoder = if compression.arrow_compression() { FlightEncoder::default() } else { @@ -434,7 +485,7 @@ mod test { .unwrap() .as_stream(); let mut stream = FlightRecordBatchStream::new( - recordbatches, + FlightRecordBatchSource::RecordBatches(recordbatches), TracingContext::default(), FlightCompression::default(), QueryContext::arc(), @@ -468,6 +519,23 @@ mod test { } } + #[tokio::test] + async fn test_flight_record_batch_stream_forwards_initializer_error() { + let mut stream = FlightRecordBatchStream::new( + FlightRecordBatchSource::initializer(async { + Err(tonic::Status::unavailable( + "remote read initialization failed", + )) + }), + TracingContext::default(), + FlightCompression::default(), + QueryContext::arc(), + ); + + let error = stream.next().await.unwrap().unwrap_err(); + assert_eq!(tonic::Code::Unavailable, error.code()); + assert!(stream.next().await.is_none()); + } #[tokio::test] async fn test_flight_record_batch_stream_emits_metrics_while_pending() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( @@ -486,7 +554,7 @@ mod test { let query_ctx = query_context_with_live_metrics_and_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - recordbatches, + FlightRecordBatchSource::RecordBatches(recordbatches), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -541,7 +609,7 @@ mod test { let query_ctx = query_context_with_live_metrics_and_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - recordbatches, + FlightRecordBatchSource::RecordBatches(recordbatches), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -595,7 +663,7 @@ mod test { let query_ctx = query_context_with_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - recordbatches, + FlightRecordBatchSource::RecordBatches(recordbatches), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -638,7 +706,7 @@ mod test { let query_ctx = Arc::new(query_ctx); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - recordbatches, + FlightRecordBatchSource::RecordBatches(recordbatches), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -677,7 +745,7 @@ mod test { }); let query_ctx = query_context_with_matching_capability(); let mut stream = FlightRecordBatchStream::new( - recordbatches, + FlightRecordBatchSource::RecordBatches(recordbatches), TracingContext::default(), FlightCompression::default(), query_ctx, diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index 520635518a4..cf75ae7fff1 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -17,35 +17,132 @@ mod test { use std::collections::HashMap; use std::net::SocketAddr; use std::sync::Arc; - use std::time::Duration; + use std::time::{Duration, Instant}; use api::v1::auth_header::AuthScheme; use api::v1::query_request::Query; use api::v1::{Basic, ColumnDataType, ColumnDef, CreateTableExpr, QueryRequest, SemanticType}; - use arrow_flight::FlightDescriptor; + use arrow_flight::flight_service_server::FlightServiceServer; + use arrow_flight::{FlightData, FlightDescriptor, Ticket}; use auth::user_provider_from_option; use client::{Client, Database}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_grpc::flight::do_put::DoPutMetadata; use common_grpc::flight::{FlightEncoder, FlightMessage}; use common_query::OutputData; - use common_recordbatch::RecordBatch; use common_recordbatch::adapter::RegionWatermarkEntry; + use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; + use common_telemetry::tracing_context::TracingContext; use datatypes::prelude::{ConcreteDataType, ScalarVector, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{Int32Vector, StringVector, TimestampMillisecondVector}; use futures_util::StreamExt; + use hyper_util::rt::TokioIo; use itertools::Itertools; use servers::grpc::builder::GrpcServerBuilder; + use servers::grpc::flight::{ + FlightCraft, FlightCraftWrapper, FlightRecordBatchSource, FlightRecordBatchStream, + TonicStream, + }; use servers::grpc::greptime_handler::GreptimeRequestHandler; use servers::grpc::{FlightCompression, GrpcServerConfig}; use servers::server::Server; + use tonic::Response; + use tonic::transport::Server as TonicServer; + use tower::service_fn; use crate::cluster::GreptimeDbClusterBuilder; use crate::grpc::query_and_expect; use crate::test_util::{StorageType, setup_grpc_server}; use crate::tests::test_util::MockInstance; + struct SlowFlightCraft; + + fn slow_recordbatch_stream() -> SendableRecordBatchStream { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "value", + ConcreteDataType::int32_datatype(), + false, + )])); + let recordbatch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_vec(vec![1])) as VectorRef], + ) + .unwrap(); + + RecordBatches::try_new(schema, vec![recordbatch]) + .unwrap() + .as_stream() + } + + #[async_trait::async_trait] + impl FlightCraft for SlowFlightCraft { + async fn do_get( + &self, + _: tonic::Request, + ) -> std::result::Result>, tonic::Status> { + let stream = FlightRecordBatchStream::new( + FlightRecordBatchSource::initializer(async { + tokio::time::sleep(Duration::from_secs(2)).await; + Ok(slow_recordbatch_stream()) + }), + TracingContext::default(), + FlightCompression::default(), + session::context::QueryContext::arc(), + ); + + Ok(Response::new(Box::pin(stream))) + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_do_get_timeout_does_not_cancel_slow_flight_stream() { + let (client_io, server_io) = tokio::io::duplex(1024); + tokio::spawn(async move { + TonicServer::builder() + .add_service(FlightServiceServer::new(FlightCraftWrapper( + SlowFlightCraft, + ))) + .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>( + server_io, + )])) + .await + .unwrap(); + }); + + let channel_manager = ChannelManager::with_config(ChannelConfig::new().timeout(None), None); + let mut client_io = Some(client_io); + channel_manager + .reset_with_connector( + "slow-flight", + service_fn(move |_| { + let client_io = client_io.take(); + + async move { + client_io + .map(TokioIo::new) + .ok_or_else(|| std::io::Error::other("Client already taken")) + } + }), + ) + .unwrap(); + let client = Client::with_manager_and_urls(channel_manager, ["slow-flight"]); + let mut flight_client = client.make_flight_client(false, false).unwrap(); + + let start = Instant::now(); + let mut request = tonic::Request::new(Ticket::default()); + request.set_timeout(Duration::from_secs(1)); + let response = flight_client.mut_inner().do_get(request).await.unwrap(); + assert!(start.elapsed() < Duration::from_secs(1)); + let mut stream = response.into_inner(); + + let start = Instant::now(); + assert!(stream.message().await.unwrap().is_some()); + assert!(start.elapsed() >= Duration::from_secs(1)); + + assert!(stream.message().await.unwrap().is_some()); + } #[tokio::test(flavor = "multi_thread")] async fn test_standalone_flight_do_put() { common_telemetry::init_default_ut_logging(); From 932f87f7a82d6d7af270b93e3bc2420f0f5cc7ab Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Tue, 25 Aug 2026 11:11:06 +0000 Subject: [PATCH 82/92] refactor(json2): support querying v2 storage layout (#8940) * feat(json2): support querying v2 storage layout - route missing JSON2 paths to the v2 remainder field - reconstruct complete values from explicit fields and remainder data - preserve root JSON2 columns across projections and filters - support nested JSON values in json_get string results - add and reorganize JSON2 sqlness coverage Signed-off-by: luofucong * resolve PR comments Signed-off-by: luofucong --------- Signed-off-by: luofucong --- .../function/src/scalars/json/json_get.rs | 70 ++++++-- src/mito2/src/read/flat_projection.rs | 4 +- src/mito2/src/sst/parquet/read_columns.rs | 167 +++++++++++++++++- .../src/optimizer/json_type_concretize.rs | 59 +++++-- .../standalone/common/types/json/json2.result | 25 +++ .../standalone/common/types/json/json2.sql | 4 + .../common/types/json/json2_join.result | 68 +++++++ .../common/types/json/json2_join.sql | 37 ++++ .../common/types/json/json2_limit.result | 77 -------- .../common/types/json/json2_limit.sql | 43 ----- 10 files changed, 392 insertions(+), 162 deletions(-) create mode 100644 tests/cases/standalone/common/types/json/json2_join.result create mode 100644 tests/cases/standalone/common/types/json/json2_join.sql diff --git a/src/common/function/src/scalars/json/json_get.rs b/src/common/function/src/scalars/json/json_get.rs index ce49f9d59fc..2357bc88f91 100644 --- a/src/common/function/src/scalars/json/json_get.rs +++ b/src/common/function/src/scalars/json/json_get.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::borrow::Cow; use std::sync::Arc; use arrow::array::{ArrayRef, BinaryViewArray, new_null_array}; @@ -54,12 +55,17 @@ trait JsonGetResultBuilder { fn build(&mut self) -> ArrayRef; } -fn result_builder(len: usize, with_type: &DataType) -> Result> { +fn result_builder( + len: usize, + with_type: &DataType, + is_json2: bool, +) -> Result> { let builder = match with_type { - DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { - Box::new(StringResultBuilder(StringViewBuilder::with_capacity(len))) - as Box - } + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Box::new(StringResultBuilder { + inner: StringViewBuilder::with_capacity(len), + is_json2, + }) + as Box, DataType::Int64 => Box::new(IntResultBuilder(Int64Builder::with_capacity(len))), DataType::Float64 => Box::new(FloatResultBuilder(Float64Builder::with_capacity(len))), DataType::Boolean => Box::new(BoolResultBuilder(BooleanBuilder::with_capacity(len))), @@ -71,20 +77,30 @@ fn result_builder(len: usize, with_type: &DataType) -> Result Result<()> { - self.0.append_option(jsonb::to_str(value).ok()); + // Scalar casts stay unquoted and map JSON null to SQL NULL; only containers + // use `to_string` to preserve their JSON representation. + let value = if self.is_json2 && (jsonb::is_array(value) || jsonb::is_object(value)) { + Some(jsonb::to_string(value)) + } else { + jsonb::to_str(value).ok() + }; + self.inner.append_option(value); Ok(()) } fn append_null(&mut self) { - self.0.append_null(); + self.inner.append_null(); } fn build(&mut self) -> ArrayRef { - Arc::new(self.0.finish()) + Arc::new(self.inner.finish()) } } @@ -404,18 +420,21 @@ impl Function for JsonGetWithType { let result = match arg0.data_type() { DataType::Binary | DataType::LargeBinary | DataType::BinaryView => { let arg0 = compute::cast(&arg0, &DataType::BinaryView)?; + let is_json2 = args.arg_fields.first().is_some_and(is_json2_extension_type); - if args.arg_fields.first().is_some_and(is_json2_extension_type) { - // Query concretization projects nested JSON2 paths as Struct arrays. A binary - // JSON2 argument is therefore an already-selected scalar or root value that - // only needs conversion from its JSONB representation to the requested type. + if is_json2 && path.trim_start_matches('$').split('.').all(str::is_empty) { JsonArray::from(&arg0) .project_to(&with_type) .map_err(|e| exec_datafusion_err!("{e:?}"))? } else { let jsons = arg0.as_binary_view(); - let mut builder = result_builder(len, &with_type)?; - jsonb_get(jsons, path, builder.as_mut())?; + let path = if is_json2 && !path.starts_with('$') { + Cow::Owned(format!("$.{path}")) + } else { + Cow::Borrowed(path) + }; + let mut builder = result_builder(len, &with_type, is_json2)?; + jsonb_get(jsons, &path, builder.as_mut())?; builder.build() } } @@ -511,6 +530,7 @@ mod tests { use datafusion_common::ScalarValue; use datafusion_common::arrow::array::{BinaryArray, BinaryViewArray, StringArray}; use datafusion_common::arrow::datatypes::{Float64Type, Int64Type}; + use datatypes::extension::json::Json2ExtensionType; use datatypes::types::parse_string_to_jsonb; use serde_json::json; @@ -566,6 +586,15 @@ mod tests { )) } + fn test_json_field(json: &ArrayRef, is_json2: bool) -> Arc { + let field = Field::new("json", json.data_type().clone(), true); + Arc::new(if is_json2 { + field.with_extension_type(Json2ExtensionType::default()) + } else { + field + }) + } + #[test] fn test_json_get_int() { let json_get_int = JsonGetInt::default(); @@ -850,7 +879,10 @@ mod tests { ColumnarValue::Array(json.clone()), ColumnarValue::Scalar(path.into()), ], - arg_fields: vec![], + arg_fields: vec![ + test_json_field(json, i >= json_strings.len()), + Arc::new(Field::new("path", DataType::Utf8, false)), + ], number_rows: 1, return_field: Arc::new(Field::new("x", DataType::Utf8View, false)), config_options: Arc::new(Default::default()), @@ -984,7 +1016,11 @@ mod tests { ColumnarValue::Scalar(path.into()), ColumnarValue::Scalar(ScalarValue::Utf8View(None)), ], - arg_fields: vec![], + arg_fields: vec![ + test_json_field(json, i >= json_strings.len()), + Arc::new(Field::new("path", DataType::Utf8, false)), + Arc::new(Field::new("with_type", DataType::Utf8View, true)), + ], number_rows: 1, return_field: Arc::new(Field::new("x", DataType::Utf8View, false)), config_options: Arc::new(Default::default()), diff --git a/src/mito2/src/read/flat_projection.rs b/src/mito2/src/read/flat_projection.rs index e631cf45db4..5178778ba9c 100644 --- a/src/mito2/src/read/flat_projection.rs +++ b/src/mito2/src/read/flat_projection.rs @@ -348,9 +348,9 @@ impl FlatProjectionMapper { } let field = &self.output_schema.arrow_schema().fields()[output_idx]; - if is_json2_extension_type(field) { + if is_json2_extension_type(field) && array.data_type() != field.data_type() { array = JsonArray::from(&array) - .project_to(field.data_type()) + .project_to_v2(batch.schema_ref().field(*index), field.data_type()) .context(DataTypesSnafu)?; } diff --git a/src/mito2/src/sst/parquet/read_columns.rs b/src/mito2/src/sst/parquet/read_columns.rs index 6f632aee300..68d7cae619d 100644 --- a/src/mito2/src/sst/parquet/read_columns.rs +++ b/src/mito2/src/sst/parquet/read_columns.rs @@ -14,6 +14,7 @@ use std::collections::{HashMap, HashSet}; +use datatypes::extension::json::JSON2_REMAINDER_FIELD_NAME; use parquet::arrow::ProjectionMask; use parquet::basic::{ConvertedType, Type as PhysicalType}; use parquet::schema::types::{ColumnDescriptor, SchemaDescriptor}; @@ -179,7 +180,7 @@ pub struct ProjectionMaskPlan { /// returned plan keeps `k` in the projection mask and marks `j` as /// not present in the output, so it can be synthesized during /// post-processing. -pub fn build_projection_plan( +pub(crate) fn build_projection_plan( parquet_read_cols: &ParquetReadColumns, parquet_schema_desc: &SchemaDescriptor, ) -> ProjectionMaskPlan { @@ -261,13 +262,30 @@ fn build_parquet_leaves_indices( } } - // Then fallback prefix misses to their nearest variant parent. + // Then include v2 remainder leaves or fallback prefix misses to their nearest variant parent. // TODO(fys): Gate fallback planning on the root being JSON2. A raw Binary // leaf is a JSONB variant only under a JSON2 root; plain struct Binary // children should not enter this fallback path. for col in &projection.cols { - for (path_idx, nested_path) in col.nested_paths.iter().enumerate() { - if prefix_matched[&col.root_index][path_idx] { + let path_matches = &prefix_matched[&col.root_index]; + let needs_remainder = col + .nested_paths + .iter() + .zip(path_matches) + .any(|(path, matched)| { + !*matched || path_points_to_struct(parquet_schema_desc, col.root_index, path) + }); + if needs_remainder { + let remainder_leaves = find_remainder_leaves(parquet_schema_desc, col.root_index); + if !remainder_leaves.is_empty() { + matched_leaves.extend(remainder_leaves); + matched_roots.insert(col.root_index); + continue; + } + } + + for (matched, nested_path) in path_matches.iter().zip(&col.nested_paths) { + if *matched { continue; } @@ -287,6 +305,49 @@ fn build_parquet_leaves_indices( (matched_leaves, matched_roots) } +/// Returns whether a nested path points to an explicitly materialized object. +/// +/// JSON2 v2 can split an object's children between its Struct field and the remainder, +/// so reading the Struct leaves alone may produce an incomplete object. +fn path_points_to_struct( + parquet_schema_desc: &SchemaDescriptor, + root_idx: usize, + path: &[String], +) -> bool { + let Some(mut field) = parquet_schema_desc.root_schema().get_fields().get(root_idx) else { + return false; + }; + for name in path.iter().skip(1) { + if !field.is_group() { + return false; + } + let Some(child) = field.get_fields().iter().find(|field| field.name() == name) else { + return false; + }; + field = child; + } + field.is_group() +} + +/// Finds the Parquet leaves backing a JSON2 v2 remainder field. +/// +/// The remainder is a sibling of explicitly materialized fields, so prefix matching a +/// requested path cannot find it. These leaves are needed when an explicit path is absent +/// or an explicitly materialized object may have additional children in the remainder. +fn find_remainder_leaves(parquet_schema_desc: &SchemaDescriptor, root_idx: usize) -> Vec { + parquet_schema_desc + .columns() + .iter() + .enumerate() + .filter_map(|(i, column)| { + let path = column.path().parts(); + (parquet_schema_desc.get_column_root_idx(i) == root_idx + && path.get(1).is_some_and(|x| x == JSON2_REMAINDER_FIELD_NAME)) + .then_some(i) + }) + .collect::>() +} + fn find_nearest_variant_parent( parquet_schema_desc: &SchemaDescriptor, root_idx: usize, @@ -329,6 +390,7 @@ mod tests { use std::sync::Arc; use parquet::basic::{ConvertedType, LogicalType, Repetition}; + use parquet::errors::ParquetError; use parquet::schema::types::Type; use super::*; @@ -441,6 +503,54 @@ mod tests { ); } + #[test] + fn test_v2_routes_missing_path_to_remainder() -> Result<(), ParquetError> { + let parquet = build_test_v2_schema()?; + let projection = + ParquetReadColumns::from_deduped(vec![ParquetReadColumn::new(0).with_nested_paths( + vec![ + vec!["j".to_string(), "cold".to_string()], + vec!["j".to_string(), "another".to_string()], + ], + )]); + + let plan = build_projection_plan(&projection, &parquet); + + assert_eq!(vec![true], plan.projected_root_presence); + assert_eq!(ProjectionMask::leaves(&parquet, [0, 1]), plan.mask); + Ok(()) + } + + #[test] + fn test_v2_explicit_path_does_not_read_remainder() -> Result<(), ParquetError> { + let parquet = build_test_v2_schema()?; + let projection = ParquetReadColumns::from_deduped(vec![ + ParquetReadColumn::new(0) + .with_nested_paths(vec![vec!["j".to_string(), "hot".to_string()]]), + ]); + + let plan = build_projection_plan(&projection, &parquet); + + assert_eq!(vec![true], plan.projected_root_presence); + assert_eq!(ProjectionMask::leaves(&parquet, [3]), plan.mask); + Ok(()) + } + + #[test] + fn test_v2_container_path_reads_remainder() -> Result<(), ParquetError> { + let parquet = build_test_v2_schema()?; + let projection = ParquetReadColumns::from_deduped(vec![ + ParquetReadColumn::new(0) + .with_nested_paths(vec![vec!["j".to_string(), "commit".to_string()]]), + ]); + + let plan = build_projection_plan(&projection, &parquet); + + assert_eq!(vec![true], plan.projected_root_presence); + assert_eq!(ProjectionMask::leaves(&parquet, [0, 1, 2]), plan.mask); + Ok(()) + } + #[test] fn test_merges_mixed_paths() { let parquet_schema_desc = build_test_nested_parquet_schema(); @@ -673,6 +783,55 @@ mod tests { SchemaDescriptor::new(schema) } + fn build_test_v2_schema() -> Result { + let metadata = Arc::new( + Type::primitive_type_builder("metadata", parquet::basic::Type::BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .build()?, + ); + let value = Arc::new( + Type::primitive_type_builder("value", parquet::basic::Type::BYTE_ARRAY) + .with_repetition(Repetition::REQUIRED) + .build()?, + ); + let remainder = Arc::new( + Type::group_type_builder(JSON2_REMAINDER_FIELD_NAME) + .with_repetition(Repetition::OPTIONAL) + .with_logical_type(Some(LogicalType::Variant { + specification_version: None, + })) + .with_fields(vec![metadata, value]) + .build()?, + ); + let operation = Arc::new( + Type::primitive_type_builder("operation", parquet::basic::Type::INT64) + .with_repetition(Repetition::OPTIONAL) + .build()?, + ); + let commit = Arc::new( + Type::group_type_builder("commit") + .with_repetition(Repetition::OPTIONAL) + .with_fields(vec![operation]) + .build()?, + ); + let hot = Arc::new( + Type::primitive_type_builder("hot", parquet::basic::Type::INT64) + .with_repetition(Repetition::OPTIONAL) + .build()?, + ); + let root = Arc::new( + Type::group_type_builder("j") + .with_repetition(Repetition::OPTIONAL) + .with_fields(vec![remainder, commit, hot]) + .build()?, + ); + Ok(SchemaDescriptor::new(Arc::new( + Type::group_type_builder("schema") + .with_fields(vec![root]) + .build()?, + ))) + } + // Test schema: // schema // `- j diff --git a/src/query/src/optimizer/json_type_concretize.rs b/src/query/src/optimizer/json_type_concretize.rs index 4625b619784..c9964e119d0 100644 --- a/src/query/src/optimizer/json_type_concretize.rs +++ b/src/query/src/optimizer/json_type_concretize.rs @@ -111,13 +111,26 @@ fn apply_json_type_hint( fn deduce_json_types(plan: &LogicalPlan) -> Result> { let mut json_types = HashMap::::new(); + // JSON2 columns in the final output must retain their complete values even when + // predicates or other expressions access only specific paths. + // For example, `SELECT j FROM t WHERE json_get(j, 'a') = 1`. + plan.schema() + .fields() + .iter() + .filter(|field| is_json2_extension_type(field)) + .for_each(|field| { + json_types.insert(field.name().clone(), JsonNativeType::Variant); + }); + plan.apply(|plan| { for expr in plan.expressions() { expr.apply(|expr| { if let Some((column, json_type)) = deduce_json_type(expr)? { json_types.entry(column).or_default().merge(&json_type); + Ok(TreeNodeRecursion::Jump) + } else { + Ok(TreeNodeRecursion::Continue) } - Ok(TreeNodeRecursion::Continue) })?; } Ok(TreeNodeRecursion::Continue) @@ -128,6 +141,7 @@ fn deduce_json_types(plan: &LogicalPlan) -> Result Result> { let f = match expr { Expr::ScalarFunction(f) if f.name().eq_ignore_ascii_case(JsonGetWithType::NAME) => f, + Expr::Column(c) => return Ok(Some((c.name.clone(), JsonNativeType::Variant))), _ => return Ok(None), }; @@ -360,25 +374,13 @@ mod tests { .rewrite(plan, &OptimizerContext::default())? .transformed ); - assert!(provider.scan_request().json_type_hint.contains_key("j")); - Ok(()) - } - - #[test] - fn test_allow_json2_passthrough_for_later_projection() -> Result<()> { - let json_get = json_get_expr(col("j"), path_expr("a"), Some(DataType::Int64))?; - let (provider, plan) = build_json2_scan()?; - let plan = plan - .project(vec![json_get.alias("__common_expr"), col("j")])? - .aggregate(Vec::::new(), vec![count(lit(1))])? - .build()?; - - assert!( - JsonTypeConcretizeRule - .rewrite(plan, &OptimizerContext::default())? - .transformed + assert_eq!( + Some(&JsonNativeType::Object(JsonObjectType::from([( + "a".to_string(), + JsonNativeType::i64(), + )]))), + provider.scan_request().json_type_hint.get("j") ); - assert!(provider.scan_request().json_type_hint.contains_key("j")); Ok(()) } @@ -402,6 +404,25 @@ mod tests { Ok(()) } + #[test] + fn test_allow_json2_filter_with_root_projection() -> Result<()> { + let predicate = + json_get_expr(col("j"), path_expr("a"), Some(DataType::Int64))?.eq(lit(1_i64)); + let (provider, plan) = build_json2_scan()?; + let plan = plan.filter(predicate)?.build()?; + + assert!( + JsonTypeConcretizeRule + .rewrite(plan, &OptimizerContext::default())? + .transformed + ); + assert_eq!( + Some(&JsonNativeType::Variant), + provider.scan_request().json_type_hint.get("j") + ); + Ok(()) + } + #[test] fn test_deduce_json_type_with_non_column_base() -> Result<()> { let expr = json_get_expr( diff --git a/tests/cases/standalone/common/types/json/json2.result b/tests/cases/standalone/common/types/json/json2.result index 2e9c9c8ee83..6afa0a759ca 100644 --- a/tests/cases/standalone/common/types/json/json2.result +++ b/tests/cases/standalone/common/types/json/json2.result @@ -134,6 +134,31 @@ select j.a, j.a.x from json2_table order by ts; | {"b":10,"x":null} | | +-----------------------------------+-------------------------------------+ +select j, j.a from json2_table order by ts; + ++--------------------------------------------------------------------+-----------------------------------+ +| j | json_get(json2_table.j,Utf8("a")) | ++--------------------------------------------------------------------+-----------------------------------+ +| {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1,"g":null}}]} | {"b":1} | +| {"a":{"b":-2},"c":"s2","d":[{"e":{"f":0.2,"g":null}}]} | {"b":-2} | +| {"a":{"b":3},"c":"s3","d":null} | {"b":3} | +| {"a":{"b":-4},"c":null,"d":[{"e":{"f":null,"g":-0.4}}]} | {"b":-4} | +| {"a":null,"c":"s5","d":null} | | +| {"a":null,"c":"s6","d":null} | | +| {"a":{"b":"s7"},"c":[1],"d":[{"e":{"g":-0.7}}]} | {"b":"s7"} | +| {"a":{"b":8},"c":"s8","d":null} | {"b":8} | +| {"a":{"b":null,"x":true},"c":"s9","d":[{"e":{"g":-0.9}}],"y":null} | {"b":null,"x":true} | +| {"a":{"b":10,"x":null},"c":null,"d":null,"y":false} | {"b":10,"x":null} | ++--------------------------------------------------------------------+-----------------------------------+ + +select j from json2_table where j.a.b = 1; + ++-------------------------------------------------------+ +| j | ++-------------------------------------------------------+ +| {"a":{"b":1},"c":"s1","d":[{"e":{"f":0.1,"g":null}}]} | ++-------------------------------------------------------+ + select j.c, j.y from json2_table order by ts; +-----------------------------------+-----------------------------------+ diff --git a/tests/cases/standalone/common/types/json/json2.sql b/tests/cases/standalone/common/types/json/json2.sql index f0bb5286b4c..406cc4af54d 100644 --- a/tests/cases/standalone/common/types/json/json2.sql +++ b/tests/cases/standalone/common/types/json/json2.sql @@ -44,6 +44,10 @@ select j.a.b from json2_table order by ts; select j.a, j.a.x from json2_table order by ts; +select j, j.a from json2_table order by ts; + +select j from json2_table where j.a.b = 1; + select j.c, j.y from json2_table order by ts; select j from json2_table order by ts; diff --git a/tests/cases/standalone/common/types/json/json2_join.result b/tests/cases/standalone/common/types/json/json2_join.result new file mode 100644 index 00000000000..cd5b215a980 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_join.result @@ -0,0 +1,68 @@ +create table json2_join_same_name_left ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +Affected Rows: 0 + +create table json2_join_same_name_right ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +Affected Rows: 0 + +insert into json2_join_same_name_left values + (1, 'a', '{"a": 1, "left_only": "kept"}'); + +Affected Rows: 1 + +insert into json2_join_same_name_right values + (1, 'a', '{"a": "right", "right_only": "should be kept"}'); + +Affected Rows: 1 + +admin flush_table('json2_join_same_name_left'); + ++------------------------------------------------+ +| ADMIN flush_table('json2_join_same_name_left') | ++------------------------------------------------+ +| 0 | ++------------------------------------------------+ + +admin flush_table('json2_join_same_name_right'); + ++-------------------------------------------------+ +| ADMIN flush_table('json2_join_same_name_right') | ++-------------------------------------------------+ +| 0 | ++-------------------------------------------------+ + +-- Conflicting hints for same-named JSON2 columns preserve both values as Variant. +select json_get(r.j, 'a')::string, json_get(l.j, 'a')::int64 +from json2_join_same_name_left l +join json2_join_same_name_right r +on l.k = r.k; + ++-------------------------+---------------------------------------------------+ +| json_get(r.j,Utf8("a")) | arrow_cast(json_get(l.j,Utf8("a")),Utf8("Int64")) | ++-------------------------+---------------------------------------------------+ +| right | 1 | ++-------------------------+---------------------------------------------------+ + +drop table json2_join_same_name_left; + +Affected Rows: 0 + +drop table json2_join_same_name_right; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_join.sql b/tests/cases/standalone/common/types/json/json2_join.sql new file mode 100644 index 00000000000..3d49d740f62 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_join.sql @@ -0,0 +1,37 @@ +create table json2_join_same_name_left ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +create table json2_join_same_name_right ( + ts timestamp time index, + k string, + j json2 +) +with ( + 'append_mode' = 'true' +); + +insert into json2_join_same_name_left values + (1, 'a', '{"a": 1, "left_only": "kept"}'); + +insert into json2_join_same_name_right values + (1, 'a', '{"a": "right", "right_only": "should be kept"}'); + +admin flush_table('json2_join_same_name_left'); + +admin flush_table('json2_join_same_name_right'); + +-- Conflicting hints for same-named JSON2 columns preserve both values as Variant. +select json_get(r.j, 'a')::string, json_get(l.j, 'a')::int64 +from json2_join_same_name_left l +join json2_join_same_name_right r +on l.k = r.k; + +drop table json2_join_same_name_left; + +drop table json2_join_same_name_right; diff --git a/tests/cases/standalone/common/types/json/json2_limit.result b/tests/cases/standalone/common/types/json/json2_limit.result index 035c4ed07bf..0394bfca66e 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.result +++ b/tests/cases/standalone/common/types/json/json2_limit.result @@ -65,87 +65,10 @@ order by json_get(j, 'a.b'); | 2 | 1 | +---------------------------------------------------+----------+ -select j, j.a from json2_whole_and_path_read; - -Error: 3001(EngineExecuteQuery), Invalid argument error: column types must match schema types, expected Binary but found Struct("a": Utf8View) at column index 0 - -select j from json2_whole_and_path_read where j.a.b = 1; - -Error: 3001(EngineExecuteQuery), Invalid argument error: column types must match schema types, expected Binary but found Struct("a": Struct("b": Int64)) at column index 0 - drop table json2_whole_and_path_read; Affected Rows: 0 -create table json2_join_same_name_left ( - ts timestamp time index, - k string, - j json2 -) -with ( - 'append_mode' = 'true' -); - -Affected Rows: 0 - -create table json2_join_same_name_right ( - ts timestamp time index, - k string, - j json2 -) -with ( - 'append_mode' = 'true' -); - -Affected Rows: 0 - -insert into json2_join_same_name_left values - (1, 'a', '{"a": 1, "left_only": "kept"}'); - -Affected Rows: 1 - -insert into json2_join_same_name_right values - (1, 'a', '{"a": "right", "right_only": "should be kept"}'); - -Affected Rows: 1 - -admin flush_table('json2_join_same_name_left'); - -+------------------------------------------------+ -| ADMIN flush_table('json2_join_same_name_left') | -+------------------------------------------------+ -| 0 | -+------------------------------------------------+ - -admin flush_table('json2_join_same_name_right'); - -+-------------------------------------------------+ -| ADMIN flush_table('json2_join_same_name_right') | -+-------------------------------------------------+ -| 0 | -+-------------------------------------------------+ - --- FIXME: This should return `right` and `1`. The current NULL values are caused --- by JSON type hints losing the table qualifier in joins. -select json_get(r.j, 'a')::string, json_get(l.j, 'a')::int64 -from json2_join_same_name_left l -join json2_join_same_name_right r -on l.k = r.k; - -+---------------------------------------------+---------------------------------------------------+ -| json_get(r.j,Utf8("a")) | arrow_cast(json_get(l.j,Utf8("a")),Utf8("Int64")) | -+---------------------------------------------+---------------------------------------------------+ -| {"a":"right","right_only":"should be kept"} | | -+---------------------------------------------+---------------------------------------------------+ - -drop table json2_join_same_name_left; - -Affected Rows: 0 - -drop table json2_join_same_name_right; - -Affected Rows: 0 - create table json2_without_append_mode ( ts timestamp time index, j json2 diff --git a/tests/cases/standalone/common/types/json/json2_limit.sql b/tests/cases/standalone/common/types/json/json2_limit.sql index 986ed58eb2a..fc80e97cd03 100644 --- a/tests/cases/standalone/common/types/json/json2_limit.sql +++ b/tests/cases/standalone/common/types/json/json2_limit.sql @@ -38,51 +38,8 @@ from json2_whole_and_path_read group by json_get(j, 'a.b') order by json_get(j, 'a.b'); -select j, j.a from json2_whole_and_path_read; - -select j from json2_whole_and_path_read where j.a.b = 1; - drop table json2_whole_and_path_read; -create table json2_join_same_name_left ( - ts timestamp time index, - k string, - j json2 -) -with ( - 'append_mode' = 'true' -); - -create table json2_join_same_name_right ( - ts timestamp time index, - k string, - j json2 -) -with ( - 'append_mode' = 'true' -); - -insert into json2_join_same_name_left values - (1, 'a', '{"a": 1, "left_only": "kept"}'); - -insert into json2_join_same_name_right values - (1, 'a', '{"a": "right", "right_only": "should be kept"}'); - -admin flush_table('json2_join_same_name_left'); - -admin flush_table('json2_join_same_name_right'); - --- FIXME: This should return `right` and `1`. The current NULL values are caused --- by JSON type hints losing the table qualifier in joins. -select json_get(r.j, 'a')::string, json_get(l.j, 'a')::int64 -from json2_join_same_name_left l -join json2_join_same_name_right r -on l.k = r.k; - -drop table json2_join_same_name_left; - -drop table json2_join_same_name_right; - create table json2_without_append_mode ( ts timestamp time index, j json2 From 73c4140938f4035c0782e18f6e766558041b9c63 Mon Sep 17 00:00:00 2001 From: "Lei, HUANG" <6406592+v0y4g3r@users.noreply.github.com> Date: Tue, 25 Aug 2026 11:31:47 +0000 Subject: [PATCH 83/92] feat(function): expose uddsketch rank (#8929) * feat(function): expose uddsketch rank Signed-off-by: Lei, HUANG * test(function): cover uddsketch rank in sqlness Signed-off-by: Lei, HUANG * fix(function): support legacy uddsketch rank Signed-off-by: Lei, HUANG * docs(function): document uddsketch rank behavior Signed-off-by: Lei, HUANG --------- Signed-off-by: Lei, HUANG --- src/common/function/src/function_registry.rs | 7 + src/common/function/src/scalars.rs | 1 + .../function/src/scalars/uddsketch_rank.rs | 224 ++++++++++++ src/common/function/src/uddsketch_compat.rs | 338 ++++++++++++++++-- .../common/aggregate/uddsketch.result | 36 ++ .../standalone/common/aggregate/uddsketch.sql | 18 + .../cases/uddsketch_legacy_state/case.toml | 2 +- .../uddsketch_legacy_state/verify.result | 22 +- .../cases/uddsketch_legacy_state/verify.sql | 14 +- 9 files changed, 633 insertions(+), 29 deletions(-) create mode 100644 src/common/function/src/scalars/uddsketch_rank.rs diff --git a/src/common/function/src/function_registry.rs b/src/common/function/src/function_registry.rs index ef2eb3658d4..15721ccabc6 100644 --- a/src/common/function/src/function_registry.rs +++ b/src/common/function/src/function_registry.rs @@ -41,6 +41,7 @@ use crate::scalars::primary_key::DecodePrimaryKeyFunction; use crate::scalars::string::register_string_functions; use crate::scalars::timestamp::TimestampFunction; use crate::scalars::uddsketch_calc::UddSketchCalcFunction; +use crate::scalars::uddsketch_rank::UddSketchRankFunction; use crate::scalars::vector::VectorFunction as VectorScalarFunction; use crate::system::SystemFunction; @@ -209,6 +210,7 @@ pub static FUNCTION_REGISTRY: LazyLock> = LazyLock::new(|| DateFunction::register(&function_registry); ExpressionFunction::register(&function_registry); UddSketchCalcFunction::register(&function_registry); + UddSketchRankFunction::register(&function_registry); HllCalcFunction::register(&function_registry); DecodePrimaryKeyFunction::register(&function_registry); @@ -363,6 +365,11 @@ mod tests { assert_eq!(1, registry.scalar_functions().len()); } + #[test] + fn test_uddsketch_rank_registered() { + assert!(FUNCTION_REGISTRY.get_function("uddsketch_rank").is_some()); + } + #[test] fn test_register_if_absent_registers_new_function() { let registry = FunctionRegistry::default(); diff --git a/src/common/function/src/scalars.rs b/src/common/function/src/scalars.rs index 0c089dedbf1..75dec84ea62 100644 --- a/src/common/function/src/scalars.rs +++ b/src/common/function/src/scalars.rs @@ -31,4 +31,5 @@ pub mod ip; pub(crate) mod test; pub(crate) mod timestamp; pub(crate) mod uddsketch_calc; +pub(crate) mod uddsketch_rank; pub mod udf; diff --git a/src/common/function/src/scalars/uddsketch_rank.rs b/src/common/function/src/scalars/uddsketch_rank.rs new file mode 100644 index 00000000000..d746b212ed7 --- /dev/null +++ b/src/common/function/src/scalars/uddsketch_rank.rs @@ -0,0 +1,224 @@ +// Copyright 2023 Greptime Team +// +// 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. + +//! Implementation of the scalar function `uddsketch_rank`. + +use std::fmt; +use std::fmt::Display; +use std::sync::Arc; + +use datafusion_common::DataFusionError; +use datafusion_common::arrow::array::{Array, AsArray, Float64Builder}; +use datafusion_common::arrow::datatypes::{DataType, Float64Type}; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility}; + +use crate::function::{Function, extract_args}; +use crate::function_registry::FunctionRegistry; + +const NAME: &str = "uddsketch_rank"; + +/// Implements the scalar function `uddsketch_rank`. +/// +/// It accepts a value and a serialized UDDSketch state, then estimates the +/// value's quantile rank by counting half of the matching bucket. Both current +/// and legacy state encodings are supported. Null arguments, empty or invalid +/// states, and invalid values produce null results. +#[derive(Debug)] +pub(crate) struct UddSketchRankFunction { + signature: Signature, +} + +impl UddSketchRankFunction { + pub fn register(registry: &FunctionRegistry) { + registry.register_scalar(UddSketchRankFunction::default()); + } +} + +impl Default for UddSketchRankFunction { + fn default() -> Self { + Self { + signature: Signature::exact( + vec![DataType::Float64, DataType::Binary], + Volatility::Immutable, + ), + } + } +} + +impl Display for UddSketchRankFunction { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{}", NAME.to_ascii_uppercase()) + } +} + +impl Function for UddSketchRankFunction { + fn name(&self) -> &str { + NAME + } + + fn return_type(&self, _: &[DataType]) -> datafusion_common::Result { + Ok(DataType::Float64) + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn invoke_with_args( + &self, + args: ScalarFunctionArgs, + ) -> datafusion_common::Result { + let [arg0, arg1] = extract_args(self.name(), &args)?; + + let Some(values) = arg0.as_primitive_opt::() else { + return Err(DataFusionError::Execution(format!( + "'{}' expects 1st argument to be Float64 datatype, got {}", + self.name(), + arg0.data_type() + ))); + }; + let Some(sketches) = arg1.as_binary_opt::() else { + return Err(DataFusionError::Execution(format!( + "'{}' expects 2nd argument to be Binary datatype, got {}", + self.name(), + arg1.data_type() + ))); + }; + let mut builder = Float64Builder::with_capacity(sketches.len()); + + for i in 0..sketches.len() { + if values.is_null(i) || sketches.is_null(i) { + builder.append_null(); + continue; + } + + match crate::uddsketch_compat::rank(sketches.value(i), values.value(i)) { + Ok(Some(rank)) => builder.append_value(rank), + Ok(None) => builder.append_null(), + Err(error) => { + common_telemetry::trace!("Failed to calculate UDDSketch rank: {}", error); + builder.append_null(); + } + } + } + + Ok(ColumnarValue::Array(Arc::new(builder.finish()))) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_schema::Field; + use datafusion_common::arrow::array::{Array, AsArray, BinaryArray, Float64Array}; + use datafusion_common::arrow::datatypes::{DataType, Float64Type}; + use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; + use uddsketch::UddSketch; + + use super::UddSketchRankFunction; + use crate::function::Function; + use crate::uddsketch_compat; + + fn invoke( + function: &UddSketchRankFunction, + values: Float64Array, + states: BinaryArray, + ) -> Float64Array { + let number_rows = values.len(); + let result = function + .invoke_with_args(ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(Arc::new(values)), + ColumnarValue::Array(Arc::new(states)), + ], + arg_fields: vec![], + number_rows, + return_field: Arc::new(Field::new("x", DataType::Float64, true)), + config_options: Arc::new(Default::default()), + }) + .unwrap(); + let ColumnarValue::Array(result) = result else { + unreachable!() + }; + result.as_primitive::().clone() + } + + #[test] + fn test_uddsketch_rank_function_name_and_return_type() { + let function = UddSketchRankFunction::default(); + + assert_eq!("uddsketch_rank", function.name()); + assert_eq!( + DataType::Float64, + function + .return_type(&[DataType::Float64, DataType::Binary]) + .unwrap() + ); + } + + #[test] + fn test_uddsketch_rank_function_canonical_state() { + let function = UddSketchRankFunction::default(); + let mut sketch = UddSketch::new(128, 0.01).unwrap(); + sketch + .add_batch(&[10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0]) + .unwrap(); + let values = [5.0, 10.0, 55.0, 100.0, 110.0]; + let expected = values.map(|value| sketch.rank(value).unwrap().unwrap()); + let encoded = sketch.encode().unwrap(); + let states = BinaryArray::from_iter_values((0..values.len()).map(|_| encoded.as_slice())); + + let result = invoke(&function, Float64Array::from(values.to_vec()), states); + + assert_eq!(result.values(), &expected); + } + + #[test] + fn test_uddsketch_rank_function_reads_legacy_state() { + let function = UddSketchRankFunction::default(); + let values = [f64::NEG_INFINITY, 0.0, 0.99, f64::INFINITY]; + let states = BinaryArray::from_iter_values( + (0..values.len()).map(|_| uddsketch_compat::LEGACY_STATE), + ); + + let result = invoke(&function, Float64Array::from(values.to_vec()), states); + + assert_eq!(result.null_count(), 0); + assert_eq!(result.values(), &[0.0, 0.375, 0.625, 1.0]); + } + + #[test] + fn test_uddsketch_rank_function_returns_null_for_invalid_rows() { + let function = UddSketchRankFunction::default(); + let empty = UddSketch::new(128, 0.01).unwrap().encode().unwrap(); + let mut populated = UddSketch::new(128, 0.01).unwrap(); + populated.add(1.0).unwrap(); + let populated = populated.encode().unwrap(); + let malformed = [1, 2, 3]; + let states = BinaryArray::from_iter([ + Some(populated.as_slice()), + None, + Some(empty.as_slice()), + Some(malformed.as_slice()), + Some(populated.as_slice()), + ]); + let values = + Float64Array::from(vec![None, Some(1.0), Some(1.0), Some(1.0), Some(f64::NAN)]); + + let result = invoke(&function, values, states); + + assert_eq!(result.null_count(), 5); + } +} diff --git a/src/common/function/src/uddsketch_compat.rs b/src/common/function/src/uddsketch_compat.rs index 733b4ab043a..3b14f3b91d5 100644 --- a/src/common/function/src/uddsketch_compat.rs +++ b/src/common/function/src/uddsketch_compat.rs @@ -61,6 +61,13 @@ struct LegacyUddSketchState { initial_error: f64, } +struct ValidatedLegacyUddSketch { + buckets: Vec<(LegacyBucketKey, u64)>, + alpha: f64, + gamma: f64, + count: u64, +} + pub(crate) fn decode(raw: &[u8]) -> Result { match UddSketch::decode(raw) { Ok(sketch) => Ok(sketch), @@ -85,6 +92,19 @@ pub(crate) fn quantile(raw: &[u8], quantile: f64) -> Result, String> } } +pub(crate) fn rank(raw: &[u8], value: f64) -> Result, String> { + match UddSketchRef::parse(raw) { + Ok(sketch) => sketch.rank(value).map_err(|error| error.to_string()), + Err(current_error) => decode_legacy_sketch(raw) + .and_then(|sketch| sketch.rank(value)) + .map_err(|legacy_error| { + format!( + "canonical decode failed: {current_error}; legacy decode failed: {legacy_error}" + ) + }), + } +} + fn validate_legacy_input(raw: &[u8]) -> Result<(), String> { if raw.len() > MAX_BYTES { return Err("input exceeds the legacy decode byte limit".to_string()); @@ -185,6 +205,69 @@ impl LegacyUddSketch { if !quantile.is_finite() || !(0.0..=1.0).contains(&quantile) { return Err("invalid quantile".to_string()); } + let validated = self.validate()?; + let count = validated.count; + if count == 0 { + return Ok(None); + } + + let target = if quantile == 1.0 { + count + } else { + ((count as f64 * quantile) as u64) + .saturating_add(1) + .min(count) + }; + let mut seen = 0_u64; + for (key, bucket_count) in validated.buckets { + seen += bucket_count; + if seen >= target { + return Ok(Some(legacy_bucket_value( + validated.alpha, + validated.gamma, + key, + )?)); + } + } + Err("legacy bucket counts do not cover the quantile rank".to_string()) + } + + fn rank(self, value: f64) -> Result, String> { + let validated = self.validate()?; + let count = validated.count; + if count == 0 { + return Ok(None); + } + if value.is_nan() { + return Err("invalid rank value".to_string()); + } + + let index = if value.is_infinite() { + i64::MAX + } else { + value.abs().log(validated.gamma).ceil() as i64 + }; + let target = if value == 0.0 { + LegacyBucketKey::Zero + } else if value.is_sign_negative() { + LegacyBucketKey::Negative(index) + } else { + LegacyBucketKey::Positive(index) + }; + let mut below = 0.0; + for (key, bucket_count) in validated.buckets { + if key == target { + return Ok(Some((below + bucket_count as f64 / 2.0) / count as f64)); + } + if key_lt(target, key) { + return Ok(Some(below / count as f64)); + } + below += bucket_count as f64; + } + Ok(Some(1.0)) + } + + fn validate(self) -> Result { if self.compactions >= 64 { return Err("legacy compaction count must be below 64".to_string()); } @@ -193,9 +276,6 @@ impl LegacyUddSketch { return Err("legacy maximum bucket count is outside supported limits".to_string()); } - let count = self.count; - let alpha = self.alpha; - let gamma = self.gamma; let buckets = self.buckets.ordered()?; if buckets.len() > self.max_buckets as usize { return Err("legacy populated bucket count exceeds its configured limit".to_string()); @@ -210,28 +290,16 @@ impl LegacyUddSketch { .checked_add(*count) .ok_or_else(|| "legacy bucket count sum overflows u64".to_string()) })?; - if decoded_count != count { + if decoded_count != self.count { return Err("legacy bucket counts do not match the value count".to_string()); } - if count == 0 { - return Ok(None); - } - let target = if quantile == 1.0 { - count - } else { - ((count as f64 * quantile) as u64) - .saturating_add(1) - .min(count) - }; - let mut seen = 0_u64; - for (key, bucket_count) in buckets { - seen += bucket_count; - if seen >= target { - return Ok(Some(legacy_bucket_value(alpha, gamma, key)?)); - } - } - Err("legacy bucket counts do not cover the quantile rank".to_string()) + Ok(ValidatedLegacyUddSketch { + buckets, + alpha: self.alpha, + gamma: self.gamma, + count: self.count, + }) } } @@ -471,6 +539,232 @@ pub(crate) const COMPACTED_LEGACY_SKETCH: &[u8] = &[ mod tests { use super::*; + fn legacy_rank_sketch() -> LegacyUddSketch { + LegacyUddSketch { + buckets: LegacyBucketStore { + map: HashMap::from([ + ( + LegacyBucketKey::Negative(2), + LegacyBucket { + count: 2, + next: LegacyBucketKey::Zero, + }, + ), + ( + LegacyBucketKey::Zero, + LegacyBucket { + count: 1, + next: LegacyBucketKey::Positive(2), + }, + ), + ( + LegacyBucketKey::Positive(2), + LegacyBucket { + count: 1, + next: LegacyBucketKey::Invalid, + }, + ), + ]), + head: LegacyBucketKey::Negative(2), + }, + alpha: 0.01, + gamma: (1.0 + 0.01) / (1.0 - 0.01), + compactions: 0, + max_buckets: 128, + count: 4, + sum: 0.0, + } + } + + #[test] + fn legacy_rank_matches_canonical_rank() { + let mut sketch = UddSketch::new(128, 0.01).unwrap(); + sketch + .add_batch(&[-10.0, -1.0, 0.0, 1.0, 10.0, f64::INFINITY]) + .unwrap(); + let encoded = sketch.encode().unwrap(); + + for value in [ + f64::NEG_INFINITY, + -10.0, + -0.0, + 0.0, + 1.0, + 10.0, + f64::INFINITY, + ] { + assert_eq!(rank(&encoded, value).unwrap(), sketch.rank(value).unwrap()); + } + } + + #[test] + fn legacy_rank_reads_wrapped_and_bare_legacy_states() { + let bare = &LEGACY_STATE[..LEGACY_STATE.len() - std::mem::size_of::()]; + + for value in [f64::NEG_INFINITY, 0.0, 0.99, f64::INFINITY] { + assert_eq!( + rank(LEGACY_STATE, value).unwrap(), + rank(bare, value).unwrap() + ); + } + assert_eq!(rank(LEGACY_STATE, f64::NEG_INFINITY).unwrap(), Some(0.0)); + assert_eq!(rank(LEGACY_STATE, f64::INFINITY).unwrap(), Some(1.0)); + } + + #[test] + fn legacy_rank_reads_compacted_bare_sketch() { + assert_eq!( + rank(COMPACTED_LEGACY_SKETCH, f64::NEG_INFINITY).unwrap(), + Some(0.0) + ); + assert!(rank(COMPACTED_LEGACY_SKETCH, 1.0).unwrap().is_some()); + assert_eq!( + rank(COMPACTED_LEGACY_SKETCH, f64::INFINITY).unwrap(), + Some(1.0) + ); + } + + #[test] + fn legacy_rank_handles_bucket_boundaries_and_empty_sketch() { + let gamma = legacy_rank_sketch().gamma; + for (value, expected) in [ + (f64::NEG_INFINITY, 0.0), + (-1.0, 0.5), + (-0.0, 0.625), + (0.0, 0.625), + (1.0, 0.75), + (gamma.powi(2), 0.875), + (f64::INFINITY, 1.0), + ] { + assert_eq!(legacy_rank_sketch().rank(value).unwrap(), Some(expected)); + } + + let mut empty = legacy_rank_sketch(); + empty.buckets.map.clear(); + empty.buckets.head = LegacyBucketKey::Invalid; + empty.count = 0; + assert_eq!(empty.rank(0.0).unwrap(), None); + } + + #[test] + fn legacy_rank_empty_sketch_returns_none_for_nan() { + let mut empty = legacy_rank_sketch(); + empty.buckets.map.clear(); + empty.buckets.head = LegacyBucketKey::Invalid; + empty.count = 0; + + assert_eq!(empty.rank(f64::NAN).unwrap(), None); + } + + #[test] + fn legacy_rank_accumulates_large_counts_with_canonical_rounding() { + let large_count = 1_u64 << 53; + let total_count = large_count + 3; + let sketch = LegacyUddSketch { + buckets: LegacyBucketStore { + map: HashMap::from([ + ( + LegacyBucketKey::Negative(3), + LegacyBucket { + count: large_count, + next: LegacyBucketKey::Negative(2), + }, + ), + ( + LegacyBucketKey::Negative(2), + LegacyBucket { + count: 1, + next: LegacyBucketKey::Negative(1), + }, + ), + ( + LegacyBucketKey::Negative(1), + LegacyBucket { + count: 1, + next: LegacyBucketKey::Zero, + }, + ), + ( + LegacyBucketKey::Zero, + LegacyBucket { + count: 1, + next: LegacyBucketKey::Invalid, + }, + ), + ]), + head: LegacyBucketKey::Negative(3), + }, + alpha: 0.01, + gamma: (1.0 + 0.01) / (1.0 - 0.01), + compactions: 0, + max_buckets: 128, + count: total_count, + sum: 0.0, + }; + let mut below = 0.0; + for count in [large_count, 1, 1] { + below += count as f64; + } + let expected = (below + 0.5) / total_count as f64; + + assert_eq!(sketch.rank(0.0).unwrap(), Some(expected)); + } + + #[test] + fn legacy_rank_rejects_nan_and_malformed_data() { + let bare_len = LEGACY_STATE.len() - std::mem::size_of::(); + let bare = &LEGACY_STATE[..bare_len]; + assert!(rank(bare, f64::NAN).is_err()); + + let mut invalid_mapping = bare.to_vec(); + let alpha_offset = bare_len - 44; + invalid_mapping[alpha_offset..alpha_offset + 8].copy_from_slice(&f64::NAN.to_le_bytes()); + let error = rank(&invalid_mapping, 0.0).unwrap_err(); + assert!(error.contains("canonical decode failed")); + assert!(error.contains("legacy decode failed")); + + let mut invalid_count = legacy_rank_sketch(); + invalid_count.count += 1; + assert!(invalid_count.rank(0.0).is_err()); + + let mut invalid_bucket = legacy_rank_sketch(); + let bucket = invalid_bucket + .buckets + .map + .remove(&LegacyBucketKey::Negative(2)) + .unwrap(); + invalid_bucket + .buckets + .map + .insert(LegacyBucketKey::Negative(i64::MIN), bucket); + invalid_bucket.buckets.head = LegacyBucketKey::Negative(i64::MIN); + assert!(invalid_bucket.rank(0.0).is_err()); + } + + #[test] + fn legacy_rank_accepts_saturated_mapping() { + let sketch = LegacyUddSketch { + buckets: LegacyBucketStore { + map: HashMap::from([( + LegacyBucketKey::Positive(0), + LegacyBucket { + count: 1, + next: LegacyBucketKey::Invalid, + }, + )]), + head: LegacyBucketKey::Positive(0), + }, + alpha: 1.0, + gamma: f64::INFINITY, + compactions: 63, + max_buckets: 7, + count: 1, + sum: f64::INFINITY, + }; + + assert_eq!(sketch.rank(1.0).unwrap(), Some(0.5)); + } + #[test] fn current_format_is_decoded_without_legacy_fallback() { let mut sketch = UddSketch::new(128, 0.01).unwrap(); diff --git a/tests/cases/standalone/common/aggregate/uddsketch.result b/tests/cases/standalone/common/aggregate/uddsketch.result index 6304ea93625..41f459a36af 100644 --- a/tests/cases/standalone/common/aggregate/uddsketch.result +++ b/tests/cases/standalone/common/aggregate/uddsketch.result @@ -52,6 +52,23 @@ select uddsketch_calc(0.95, uddsketch_state(128, 0.01, `value`)) from test_uddsk | 100.49456770856492 | +----------------------------------------------------------------------------------------------+ +SELECT + ROUND(uddsketch_rank(5, `state`), 4) AS below_min_rank, + ROUND(uddsketch_rank(10, `state`), 4) AS equal_min_rank, + ROUND(uddsketch_rank(55, `state`), 4) AS middle_rank, + ROUND(uddsketch_rank(100, `state`), 4) AS equal_max_rank, + ROUND(uddsketch_rank(110, `state`), 4) AS above_max_rank +FROM ( + SELECT uddsketch_state(128, 0.01, `value`) AS `state` + FROM test_uddsketch +); + ++----------------+----------------+-------------+----------------+----------------+ +| below_min_rank | equal_min_rank | middle_rank | equal_max_rank | above_max_rank | ++----------------+----------------+-------------+----------------+----------------+ +| 0.0 | 0.05 | 0.5 | 0.95 | 1.0 | ++----------------+----------------+-------------+----------------+----------------+ + CREATE TABLE grouped_uddsketch ( `state` BINARY, id_group INT PRIMARY KEY, @@ -72,6 +89,15 @@ SELECT uddsketch_calc(0.1, uddsketch_merge(128, 0.01, `state`)) FROM grouped_udd | 19.886670240866184 | +------------------------------------------------------------------------------------------------+ +SELECT ROUND(uddsketch_rank(55, uddsketch_merge(128, 0.01, `state`)), 4) AS merged_rank +FROM grouped_uddsketch; + ++-------------+ +| merged_rank | ++-------------+ +| 0.5 | ++-------------+ + -- should fail SELECT uddsketch_calc(0.1, uddsketch_merge(128, 0.1, `state`)) FROM grouped_uddsketch; @@ -82,6 +108,16 @@ SELECT uddsketch_calc(0.1, uddsketch_merge(64, 0.01, `state`)) FROM grouped_udds Error: 3001(EngineExecuteQuery), Error during planning: Merging UDDSketch with different parameters: arguments=(64, 0.01) vs actual input=(128, 0.01) +SELECT uddsketch_rank(55, uddsketch_state(128, 0.01, `value`)) IS NULL AS empty_rank_is_null +FROM test_uddsketch +WHERE `id` < 0; + ++--------------------+ +| empty_rank_is_null | ++--------------------+ +| true | ++--------------------+ + drop table test_uddsketch; Affected Rows: 0 diff --git a/tests/cases/standalone/common/aggregate/uddsketch.sql b/tests/cases/standalone/common/aggregate/uddsketch.sql index c959f33ee99..13deb88ddb1 100644 --- a/tests/cases/standalone/common/aggregate/uddsketch.sql +++ b/tests/cases/standalone/common/aggregate/uddsketch.sql @@ -24,6 +24,17 @@ select uddsketch_calc(0.75, uddsketch_state(128, 0.01, `value`)) from test_uddsk select uddsketch_calc(0.95, uddsketch_state(128, 0.01, `value`)) from test_uddsketch; +SELECT + ROUND(uddsketch_rank(5, `state`), 4) AS below_min_rank, + ROUND(uddsketch_rank(10, `state`), 4) AS equal_min_rank, + ROUND(uddsketch_rank(55, `state`), 4) AS middle_rank, + ROUND(uddsketch_rank(100, `state`), 4) AS equal_max_rank, + ROUND(uddsketch_rank(110, `state`), 4) AS above_max_rank +FROM ( + SELECT uddsketch_state(128, 0.01, `value`) AS `state` + FROM test_uddsketch +); + CREATE TABLE grouped_uddsketch ( `state` BINARY, id_group INT PRIMARY KEY, @@ -34,12 +45,19 @@ INSERT INTO grouped_uddsketch (`state`, id_group) SELECT uddsketch_state(128, 0. SELECT uddsketch_calc(0.1, uddsketch_merge(128, 0.01, `state`)) FROM grouped_uddsketch; +SELECT ROUND(uddsketch_rank(55, uddsketch_merge(128, 0.01, `state`)), 4) AS merged_rank +FROM grouped_uddsketch; + -- should fail SELECT uddsketch_calc(0.1, uddsketch_merge(128, 0.1, `state`)) FROM grouped_uddsketch; -- should fail SELECT uddsketch_calc(0.1, uddsketch_merge(64, 0.01, `state`)) FROM grouped_uddsketch; +SELECT uddsketch_rank(55, uddsketch_state(128, 0.01, `value`)) IS NULL AS empty_rank_is_null +FROM test_uddsketch +WHERE `id` < 0; + drop table test_uddsketch; drop table grouped_uddsketch; diff --git a/tests/compatibility/cases/uddsketch_legacy_state/case.toml b/tests/compatibility/cases/uddsketch_legacy_state/case.toml index b22670037d1..e68fdf3804b 100644 --- a/tests/compatibility/cases/uddsketch_legacy_state/case.toml +++ b/tests/compatibility/cases/uddsketch_legacy_state/case.toml @@ -1,5 +1,5 @@ name = "uddsketch_legacy_state" -reason = "Verify UDDSketch state bytes persisted by old binaries (wrapped legacy bincode UddSketchState) are read directly by uddsketch_calc and uddsketch_merge on the new binary, merged into canonical v1, and merged together with newly generated canonical states." +reason = "Verify UDDSketch state bytes persisted by old binaries (wrapped legacy bincode UddSketchState) are read directly by uddsketch_calc, uddsketch_rank, and uddsketch_merge on the new binary, merged into canonical v1, and merged together with newly generated canonical states." introduced_by = "PR #8867" topologies = ["distributed", "standalone"] from_range = [">=v1.0.0"] diff --git a/tests/compatibility/cases/uddsketch_legacy_state/verify.result b/tests/compatibility/cases/uddsketch_legacy_state/verify.result index b4c8a3d1566..d1a40d5a15b 100644 --- a/tests/compatibility/cases/uddsketch_legacy_state/verify.result +++ b/tests/compatibility/cases/uddsketch_legacy_state/verify.result @@ -11,7 +11,23 @@ ORDER BY grp; | 10 | 100.49456770856492 | +-----+--------------------+ --- (2) uddsketch_merge over the persisted legacy rows, then calc. +-- (2) direct uddsketch_rank of each persisted legacy row +SELECT grp, + ROUND(uddsketch_rank(30.0, state), 4) AS rank_at_30, + ROUND(uddsketch_rank(70.0, state), 4) AS rank_at_70, + ROUND(uddsketch_rank(100.0, state), 4) AS rank_at_100 +FROM uddsketch_states +ORDER BY grp; + ++-----+------------+------------+-------------+ +| grp | rank_at_30 | rank_at_70 | rank_at_100 | ++-----+------------+------------+-------------+ +| 0 | 0.625 | 1.0 | 1.0 | +| 5 | 0.0 | 0.5 | 1.0 | +| 10 | 0.0 | 0.0 | 0.5 | ++-----+------------+------------+-------------+ + +-- (3) uddsketch_merge over the persisted legacy rows, then calc. -- Succeeds only if every legacy row decodes to the exact declared -- parameters (128, 0.01); the merge parameter check is the guard. SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 @@ -23,7 +39,7 @@ FROM uddsketch_states; | 59.745049810145126 | +--------------------+ --- (3) generate fresh canonical states with the new binary and merge mixed legacy + canonical rows +-- (4) generate fresh canonical states with the new binary and merge mixed legacy + canonical rows CREATE TABLE uddsketch_new_states ( state BINARY, grp INT PRIMARY KEY, @@ -52,7 +68,7 @@ FROM ( | 59.745049810145126 | +--------------------+ --- (4) canonical-only control +-- (5) canonical-only control SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 FROM uddsketch_new_states; diff --git a/tests/compatibility/cases/uddsketch_legacy_state/verify.sql b/tests/compatibility/cases/uddsketch_legacy_state/verify.sql index 8cf84cfbbc5..dd3188f6712 100644 --- a/tests/compatibility/cases/uddsketch_legacy_state/verify.sql +++ b/tests/compatibility/cases/uddsketch_legacy_state/verify.sql @@ -3,13 +3,21 @@ SELECT grp, uddsketch_calc(0.5, state) AS p50 FROM uddsketch_states ORDER BY grp; --- (2) uddsketch_merge over the persisted legacy rows, then calc. +-- (2) direct uddsketch_rank of each persisted legacy row +SELECT grp, + ROUND(uddsketch_rank(30.0, state), 4) AS rank_at_30, + ROUND(uddsketch_rank(70.0, state), 4) AS rank_at_70, + ROUND(uddsketch_rank(100.0, state), 4) AS rank_at_100 +FROM uddsketch_states +ORDER BY grp; + +-- (3) uddsketch_merge over the persisted legacy rows, then calc. -- Succeeds only if every legacy row decodes to the exact declared -- parameters (128, 0.01); the merge parameter check is the guard. SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 FROM uddsketch_states; --- (3) generate fresh canonical states with the new binary and merge mixed legacy + canonical rows +-- (4) generate fresh canonical states with the new binary and merge mixed legacy + canonical rows CREATE TABLE uddsketch_new_states ( state BINARY, grp INT PRIMARY KEY, @@ -28,6 +36,6 @@ FROM ( SELECT state FROM uddsketch_new_states ) AS all_states; --- (4) canonical-only control +-- (5) canonical-only control SELECT uddsketch_calc(0.5, uddsketch_merge(128, 0.01, state)) AS p50 FROM uddsketch_new_states; From 1851f6bf4de7b6ae94fc9b0cb3048aa24877aea9 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 25 Aug 2026 12:34:53 +0000 Subject: [PATCH 84/92] fix(query): keep INSERT timestamp conversion out of the source query (#8911) * fix(query): keep INSERT timestamp conversion out of the source query Interpreting an INSERT's string timestamps used to work by pushing the conversion down into the source query, which changed what that query means. Two consequences: - Pushing through a UNION's DISTINCT moved the dedup key from the raw strings to parsed instants, so rows spelling the same instant differently collapsed into one. On an append-only table that is a silently dropped row. - A UNION branch that needed no conversion (a NULL, or an explicit cast) made the whole column give up, leaving sibling branches on UTC while the rest of the row used the session timezone. Convert at the assignment instead, by routing its cast through a timezone-carrying timestamp type and back. Arrow applies the timezone when a cast target carries one, and stripping it afterwards preserves the value. The source query is no longer touched, so both cases go away and the tree-walking rewrite (roughly 160 lines) is deleted. The rewrite reads source types, so it now runs TypeCoercion first: a UNION still carries its loose per-branch schema before coercion, and retargeting a cast whose input later becomes a timestamp would shift the value rather than reinterpret it. Signed-off-by: Dennis Zhuang * fix(query): address review on INSERT assignment rewrite - Clone the input `Arc` instead of the whole subtree, and only rebuild it when a `Values` row actually changes. - Defer cloning the cast source until the literal-folding path has been ruled out. - Move the UTC check onto `Timezone::is_utc`, replacing a bare string compare. - Cover a prepared `INSERT ... VALUES (?)`: an untyped placeholder types as `Null`, so the assignment cast is left for parameter substitution. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/common/time/src/timezone.rs | 9 + src/query/src/optimizer.rs | 1 + src/query/src/optimizer/insert_assignment.rs | 307 ++++++++++++++++ src/query/src/optimizer/type_conversion.rs | 11 +- .../type_conversion/insert_assignment.rs | 336 ------------------ src/query/src/planner.rs | 126 +++---- .../insert/insert_default_timezone.result | 58 +++ .../common/insert/insert_default_timezone.sql | 28 ++ 8 files changed, 472 insertions(+), 404 deletions(-) create mode 100644 src/query/src/optimizer/insert_assignment.rs delete mode 100644 src/query/src/optimizer/type_conversion/insert_assignment.rs diff --git a/src/common/time/src/timezone.rs b/src/common/time/src/timezone.rs index 41cc1f78428..dbd8dffcd49 100644 --- a/src/common/time/src/timezone.rs +++ b/src/common/time/src/timezone.rs @@ -108,6 +108,15 @@ impl Timezone { } } + /// A named zone merely sitting at +00:00 today is not UTC: it may have been + /// elsewhere at the timestamp being converted. + pub fn is_utc(&self) -> bool { + match self { + Self::Offset(offset) => offset.local_minus_utc() == 0, + Self::Named(tz) => matches!(tz, Tz::UTC), + } + } + /// Returns the number of seconds to add to convert from UTC to the local time. pub fn local_minus_utc(&self) -> i64 { match self { diff --git a/src/query/src/optimizer.rs b/src/query/src/optimizer.rs index 8827d48ed63..a6b36e6fdae 100644 --- a/src/query/src/optimizer.rs +++ b/src/query/src/optimizer.rs @@ -17,6 +17,7 @@ pub mod constant_term; pub mod count_nest_aggr; pub mod count_wildcard; pub mod global_limit; +pub(crate) mod insert_assignment; pub(crate) mod json_type_concretize; pub mod parallelize_scan; pub mod pass_distribution; diff --git a/src/query/src/optimizer/insert_assignment.rs b/src/query/src/optimizer/insert_assignment.rs new file mode 100644 index 00000000000..ad77736ce0a --- /dev/null +++ b/src/query/src/optimizer/insert_assignment.rs @@ -0,0 +1,307 @@ +// Copyright 2023 Greptime Team +// +// 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. + +use std::sync::Arc; + +use common_time::Timezone; +use datafusion::config::ConfigOptions; +use datafusion_common::{DFSchemaRef, Result, ScalarValue}; +use datafusion_expr::expr::{Alias, Cast}; +use datafusion_expr::{Distinct, Expr, ExprSchemable, LogicalPlan, Projection, Values}; +use datafusion_optimizer::analyzer::AnalyzerRule; +use datafusion_optimizer::analyzer::type_coercion::TypeCoercion; +use datatypes::arrow::datatypes::{DataType, TimeUnit}; +use session::context::QueryContextRef; + +use crate::optimizer::type_conversion::cast_string_to_timestamp; + +/// Interprets strings assigned to timestamp columns at an `INSERT` boundary +/// using the session timezone. `plan` is the assignment projection under a +/// `WriteOp::Insert`. +/// +/// DataFusion plans `INSERT` as a projection casting each source column to its +/// target column type, and that cast reads a naive string as UTC. Arrow does +/// apply a timezone when the cast target carries one, so the assignment is +/// routed through `Timestamp(unit, Some(tz))` and back. Stripping the timezone +/// afterwards is value-preserving — arrow only shifts values in the opposite +/// direction. +/// +/// The source query is left untouched. Reinterpreting a value where it is +/// *produced* would change what the source query means: pushing the conversion +/// below a `UNION`'s `DISTINCT`, for instance, moves the dedup key from the raw +/// strings to parsed instants and silently drops rows. +/// +/// # Why `TypeCoercion` runs here +/// +/// The rewrite reads source types, and those are only settled once a `UNION`'s +/// branch types have been reconciled: before coercion a union carries its loose +/// schema (the first branch's types), so a mixed +/// `SELECT 'string' UNION ALL SELECT CAST(.. AS TIMESTAMP)` still looks like a +/// string. Retargeting that cast would leave `Timestamp(None) -> +/// Timestamp(Some(tz))` behind once coercion retypes the union — the one +/// direction in which arrow shifts the value instead of relabelling it. +/// +/// Coercing here rather than deferring to the analyzer is forced by where an +/// INSERT is still identifiable: `exec_dml_statement` strips the `Dml` node and +/// executes its input, so by the time the analyzer runs, an assignment +/// projection is indistinguishable from any other projection. +/// +/// Explicit casts stay out of this: the SQL layer turns a user's +/// `CAST(x AS TIMESTAMP)` into an `arrow_cast` call, which only becomes an +/// `Expr::Cast` in the optimizer's `SimplifyExpressions`. Assignment casts are +/// therefore the only `Expr::Cast` reaching a timestamp column here. +/// +/// # Reach +/// +/// The emitted cast only carries its timezone where the expression is evaluated +/// on this node. Substrait drops the timezone name when a plan is pushed down — +/// it encodes any zoned timestamp as `PrecisionTimestampTz` and decodes it back +/// as UTC — so a source reading from a table falls back to UTC, the behaviour it +/// had before this rule existed. Sources that never leave this node (literals, +/// `VALUES`, and `UNION`s of them) keep the session timezone, and those are what +/// an INSERT's timestamp assignment is in practice. +pub(crate) fn rewrite_insert_assignments( + plan: LogicalPlan, + query_ctx: &QueryContextRef, + config: &ConfigOptions, +) -> Result { + let Some(timezone) = session_timezone(query_ctx) else { + return Ok(plan); + }; + + let plan = TypeCoercion::new().analyze(plan, config)?; + rewrite_assignment(plan, &timezone) +} + +/// Session timezone, in both forms the rewrite needs. +struct SessionTimezone { + /// Parses literals, matching the plain `INSERT ... VALUES` path. + parsed: Timezone, + /// Names the intermediate arrow cast target. + name: Arc, +} + +fn session_timezone(query_ctx: &QueryContextRef) -> Option { + let parsed = query_ctx.timezone(); + + // A UTC session already gets UTC semantics from the plain assignment cast. + if parsed.is_utc() { + return None; + } + + Some(SessionTimezone { + name: Arc::from(parsed.to_string()), + parsed, + }) +} + +fn rewrite_assignment(plan: LogicalPlan, timezone: &SessionTimezone) -> Result { + let LogicalPlan::Projection(assignment) = plan else { + return Ok(plan); + }; + + let mut exprs = assignment.expr.clone(); + let mut changed = false; + for expr in &mut exprs { + changed |= retarget_assignment_cast( + expr, + assignment.input.schema(), + Some(assignment.input.as_ref()), + timezone, + )?; + } + + // The planner types `VALUES` against the target table, so the assignment + // cast lands inside the `Values` rows instead of on the projection above. + let mut input = assignment.input.clone(); + if let LogicalPlan::Values(values) = assignment.input.as_ref() + && let Some(rewritten) = rewrite_values(values, timezone)? + { + input = Arc::new(LogicalPlan::Values(rewritten)); + changed = true; + } + + if !changed { + return Ok(LogicalPlan::Projection(assignment)); + } + Projection::try_new(exprs, input).map(LogicalPlan::Projection) +} + +fn rewrite_values(values: &Values, timezone: &SessionTimezone) -> Result> { + let mut rewritten = values.clone(); + let mut changed = false; + for row in &mut rewritten.values { + for expr in row.iter_mut() { + changed |= retarget_assignment_cast(expr, &values.schema, None, timezone)?; + } + } + + Ok(changed.then_some(rewritten)) +} + +/// Reinterprets one assignment cast, returning whether it was rewritten. +/// +/// `source_plan` is the projection's input, used to resolve a literal behind a +/// column reference; `Values` rows carry their expression inline and pass `None`. +fn retarget_assignment_cast( + expr: &mut Expr, + schema: &DFSchemaRef, + source_plan: Option<&LogicalPlan>, + timezone: &SessionTimezone, +) -> Result { + let expr = unalias_mut(expr); + let Expr::Cast(Cast { + expr: source, + data_type: DataType::Timestamp(unit, None), + }) = expr + else { + return Ok(false); + }; + let unit = *unit; + + if !matches!( + source.get_type(schema)?, + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) { + return Ok(false); + } + + // Fold literals with the same parser the plain `INSERT ... VALUES` path + // uses, so a given string means the same thing however it reaches a column. + // The parsers disagree on ambiguous local times: this one resolves them, + // arrow rejects them. + let folded = source_literal(source.as_ref(), source_plan) + .and_then(|literal| convert_literal(&literal, unit, &timezone.parsed)); + if let Some(folded) = folded { + *expr = folded; + return Ok(true); + } + + let source = source.as_ref().clone(); + *expr = Expr::Cast(Cast::new( + Box::new(Expr::Cast(Cast::new( + Box::new(source), + DataType::Timestamp(unit, Some(timezone.name.clone())), + ))), + DataType::Timestamp(unit, None), + )); + Ok(true) +} + +fn source_literal(source: &Expr, source_plan: Option<&LogicalPlan>) -> Option { + match source { + Expr::Literal(value, _) => Some(value.clone()), + Expr::Column(column) => { + let plan = source_plan?; + let index = plan.schema().maybe_index_of_column(column)?; + lineage_literal(plan, index).cloned() + } + _ => None, + } +} + +/// Resolves a literal when every row carries the same value at `output_idx`. +/// +/// Read-only: the literal is folded into the assignment above, so nodes that +/// drop, reorder or deduplicate rows can be traversed — none of them changes +/// the value a surviving row carries, and folding above them leaves their keys +/// on the original strings. +fn lineage_literal(plan: &LogicalPlan, output_idx: usize) -> Option<&ScalarValue> { + if output_idx >= plan.schema().fields().len() { + return None; + } + + match plan { + LogicalPlan::Projection(projection) => match unalias(&projection.expr[output_idx]) { + Expr::Literal(value, _) => Some(value), + Expr::Column(column) => { + let input_idx = projection.input.schema().maybe_index_of_column(column)?; + lineage_literal(projection.input.as_ref(), input_idx) + } + _ => None, + }, + LogicalPlan::Filter(_) + | LogicalPlan::Sort(_) + | LogicalPlan::Limit(_) + | LogicalPlan::SubqueryAlias(_) + | LogicalPlan::Distinct(Distinct::All(_)) => { + let inputs = plan.inputs(); + let [input] = inputs.as_slice() else { + return None; + }; + lineage_literal(input, output_idx) + } + _ => None, + } +} + +fn convert_literal(value: &ScalarValue, unit: TimeUnit, timezone: &Timezone) -> Option { + let ScalarValue::Utf8(Some(value)) = value else { + return None; + }; + cast_string_to_timestamp(value, &DataType::Timestamp(unit, None), Some(timezone)) + .ok() + .filter(|value| !value.is_null()) + .map(|value| Expr::Literal(value, None)) +} + +fn unalias(expr: &Expr) -> &Expr { + match expr { + Expr::Alias(Alias { expr, .. }) => unalias(expr), + expr => expr, + } +} + +fn unalias_mut(expr: &mut Expr) -> &mut Expr { + match expr { + Expr::Alias(Alias { expr, .. }) => unalias_mut(expr), + expr => expr, + } +} + +#[cfg(test)] +mod tests { + use datafusion_common::DFSchema; + use datafusion_expr::expr::Placeholder; + + use super::*; + + fn shanghai() -> SessionTimezone { + let parsed = Timezone::from_tz_string("Asia/Shanghai").unwrap(); + SessionTimezone { + name: Arc::from(parsed.to_string()), + parsed, + } + } + + /// A prepared `INSERT ... VALUES (?)` arrives here as a cast over an untyped + /// placeholder, which must survive for parameter substitution. + #[test] + fn test_untyped_placeholder_assignment_is_left_alone() { + let schema = Arc::new(DFSchema::empty()); + let mut expr = Expr::Cast(Cast::new( + Box::new(Expr::Placeholder(Placeholder::new_with_field( + "$1".to_string(), + None, + ))), + DataType::Timestamp(TimeUnit::Millisecond, None), + )); + let original = expr.clone(); + + let changed = retarget_assignment_cast(&mut expr, &schema, None, &shanghai()).unwrap(); + + assert!(!changed); + assert_eq!(expr, original); + } +} diff --git a/src/query/src/optimizer/type_conversion.rs b/src/query/src/optimizer/type_conversion.rs index 3941fedfe8a..40bb43e791e 100644 --- a/src/query/src/optimizer/type_conversion.rs +++ b/src/query/src/optimizer/type_conversion.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -mod insert_assignment; - use std::sync::Arc; use common_time::Timezone; @@ -31,7 +29,7 @@ use session::context::QueryContextRef; use crate::QueryEngineContext; use crate::optimizer::ExtensionAnalyzerRule; -use crate::optimizer::type_conversion::insert_assignment::rewrite_insert_assignments; +use crate::optimizer::insert_assignment::rewrite_insert_assignments; use crate::plan::ExtractExpr; /// TypeConversionRule converts some literal values in logical plan to other types according @@ -46,7 +44,7 @@ impl ExtensionAnalyzerRule for TypeConversionRule { &self, plan: LogicalPlan, ctx: &QueryEngineContext, - _config: &ConfigOptions, + config: &ConfigOptions, ) -> Result { plan.transform_up_with_subqueries(|plan| match plan { LogicalPlan::Filter(filter) => { @@ -129,7 +127,8 @@ impl ExtensionAnalyzerRule for TypeConversionRule { LogicalPlan::Dml(mut dml) if matches!(dml.op, WriteOp::Insert(_)) => { dml.input = Arc::new(rewrite_insert_assignments( dml.input.as_ref().clone(), - ctx.query_ctx(), + &ctx.query_ctx(), + config, )?); Ok(Transformed::yes(LogicalPlan::Dml(dml))) } @@ -324,7 +323,7 @@ fn timestamp_to_timestamp_ms_expr(val: i64, unit: TimeUnit) -> Expr { ) } -fn cast_string_to_timestamp( +pub(crate) fn cast_string_to_timestamp( string: &str, target_type: &DataType, timezone: Option<&Timezone>, diff --git a/src/query/src/optimizer/type_conversion/insert_assignment.rs b/src/query/src/optimizer/type_conversion/insert_assignment.rs deleted file mode 100644 index 11ab7e7829e..00000000000 --- a/src/query/src/optimizer/type_conversion/insert_assignment.rs +++ /dev/null @@ -1,336 +0,0 @@ -// Copyright 2023 Greptime Team -// -// 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. - -use std::sync::Arc; - -use datafusion_common::{Column, Result, ScalarValue}; -use datafusion_expr::expr::{Alias, Cast}; -use datafusion_expr::{Distinct, Expr, ExprSchemable, LogicalPlan, Projection, Union, Values}; -use datatypes::arrow::datatypes::DataType; -use session::context::QueryContextRef; - -use crate::optimizer::type_conversion::cast_string_to_timestamp; -use crate::plan::ExtractExpr; - -/// Rewrites string literals that feed timestamp columns at an INSERT boundary. -/// Constants are folded at the assignment to avoid changing source types; -/// `VALUES` and `UNION` inputs are rewritten per row or branch. Explicit casts -/// stay on DataFusion's existing path. -pub(super) fn rewrite_insert_assignments( - plan: LogicalPlan, - query_ctx: QueryContextRef, -) -> Result { - let LogicalPlan::Projection(assignment) = plan else { - return Ok(plan); - }; - - let converter = InsertAssignmentConverter { query_ctx }; - let mut exprs = assignment.expr.clone(); - let mut input = assignment.input.as_ref().clone(); - let mut changed = false; - for (output_idx, expr) in assignment.expr.iter().enumerate() { - let target_type = assignment.schema.field(output_idx).data_type(); - if !matches!(target_type, DataType::Timestamp(_, _)) { - continue; - } - - let Some(column) = assignment_input_column(expr) else { - continue; - }; - let Some(input_idx) = input.schema().maybe_index_of_column(column) else { - continue; - }; - - let literal = lineage_literal(&input, input_idx).cloned(); - if let Some(literal) = literal - && let Some(folded) = converter.convert_literal(&literal, target_type) - { - let (qualifier, field) = assignment.schema.qualified_field(output_idx); - exprs[output_idx] = folded.alias_qualified(qualifier.cloned(), field.name()); - changed = true; - continue; - } - - // A hand-built DML plan (e.g. from flow) may share one source column - // between targets; rewriting it in place would retype every reader. - if assignment - .expr - .iter() - .enumerate() - .any(|(idx, other)| idx != output_idx && other.column_refs().contains(column)) - { - continue; - } - if let Some(rewritten) = converter.rewrite_output_column(&input, input_idx, target_type)? { - input = rewritten; - changed = true; - } - } - - if !changed { - return Ok(LogicalPlan::Projection(assignment)); - } - Projection::try_new(exprs, Arc::new(input)).map(LogicalPlan::Projection) -} - -/// Resolves a literal when every row carries the same value at `output_idx`. -fn lineage_literal(plan: &LogicalPlan, output_idx: usize) -> Option<&ScalarValue> { - if output_idx >= plan.schema().fields().len() { - return None; - } - - match plan { - LogicalPlan::Projection(projection) => match unalias(&projection.expr[output_idx]) { - Expr::Literal(value, _) => Some(value), - Expr::Column(column) => { - let input_idx = projection.input.schema().maybe_index_of_column(column)?; - lineage_literal(projection.input.as_ref(), input_idx) - } - _ => None, - }, - // These nodes drop, reorder or deduplicate rows without touching the - // value a surviving row carries, and their schema stays positional. - LogicalPlan::Filter(_) - | LogicalPlan::Sort(_) - | LogicalPlan::Limit(_) - | LogicalPlan::SubqueryAlias(_) - | LogicalPlan::Distinct(Distinct::All(_)) => { - let inputs = plan.inputs(); - let [input] = inputs.as_slice() else { - return None; - }; - lineage_literal(input, output_idx) - } - _ => None, - } -} - -struct InsertAssignmentConverter { - query_ctx: QueryContextRef, -} - -impl InsertAssignmentConverter { - fn rewrite_output_column( - &self, - plan: &LogicalPlan, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - if output_idx >= plan.schema().fields().len() { - return Ok(None); - } - - match plan { - // DataFusion pushes INSERT assignment casts into Values, so inspect - // them even though the Values schema already has the target type. - LogicalPlan::Values(values) => self.rewrite_values(values, output_idx, target_type), - // Once the source query has produced the target type, its casts and - // coercions belong to the source query rather than INSERT assignment. - _ if plan.schema().field(output_idx).data_type() == target_type => Ok(None), - LogicalPlan::Projection(projection) => { - self.rewrite_projection(projection, output_idx, target_type) - } - LogicalPlan::Union(union) => self.rewrite_union(union, output_idx, target_type), - // Filter and Sort are excluded here: their predicates and keys read - // the retyped column, which would change the source query. Constants - // still reach the assignment through `lineage_literal`. Distinct::All - // holds no expressions; allowing it shifts dedup keys from raw - // strings to parsed instants. - LogicalPlan::Limit(_) - | LogicalPlan::SubqueryAlias(_) - | LogicalPlan::Distinct(Distinct::All(_)) => { - self.rewrite_passthrough(plan, output_idx, target_type) - } - _ => Ok(None), - } - } - - fn rewrite_projection( - &self, - projection: &Projection, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - match unalias(&projection.expr[output_idx]) { - Expr::Literal(value, _) => { - let Some(converted) = self.convert_literal(value, target_type) else { - return Ok(None); - }; - let (qualifier, field) = projection.schema.qualified_field(output_idx); - let mut exprs = projection.expr.clone(); - exprs[output_idx] = converted.alias_qualified(qualifier.cloned(), field.name()); - - Projection::try_new(exprs, projection.input.clone()) - .map(LogicalPlan::Projection) - .map(Some) - } - Expr::Column(column) => { - // Retyping the input column affects every output column reading - // it, so only follow lineage with a single consumer. - if projection - .expr - .iter() - .enumerate() - .any(|(idx, other)| idx != output_idx && other.column_refs().contains(column)) - { - return Ok(None); - } - - let input = projection.input.as_ref(); - let Some(input_idx) = input.schema().maybe_index_of_column(column) else { - return Ok(None); - }; - let Some(rewritten) = self.rewrite_output_column(input, input_idx, target_type)? - else { - return Ok(None); - }; - - Projection::try_new(projection.expr.clone(), Arc::new(rewritten)) - .map(LogicalPlan::Projection) - .map(Some) - } - _ => Ok(None), - } - } - - fn rewrite_values( - &self, - values: &Values, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - let mut rewritten = values.clone(); - let mut changed = false; - for row in &mut rewritten.values { - let Some(expr) = row.get_mut(output_idx) else { - return Ok(None); - }; - - if let Expr::Cast(Cast { - expr: inner, - data_type, - }) = expr - && data_type == target_type - && let Expr::Literal(value, _) = inner.as_ref() - && let Some(value) = self.convert_literal(value, target_type) - { - *expr = value; - changed = true; - continue; - } - - if expr.get_type(values.schema.as_ref()).ok().as_ref() != Some(target_type) { - return Ok(None); - } - } - - Ok(changed.then_some(LogicalPlan::Values(rewritten))) - } - - fn rewrite_union( - &self, - union: &Union, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - let mut inputs = Vec::with_capacity(union.inputs.len()); - // A union has one shared schema, so rewrite it only when the same output - // column is an INSERT-assigned literal in every branch. - for input in &union.inputs { - let Some(rewritten) = self.rewrite_output_column(input, output_idx, target_type)? - else { - return Ok(None); - }; - inputs.push(Arc::new(rewritten)); - } - - // Untouched columns may still have mismatched branch types before - // TypeCoercion runs, so rebuild loosely like the SQL planner does. - Union::try_new_with_loose_types(inputs) - .map(LogicalPlan::Union) - .map(Some) - } - - fn rewrite_passthrough( - &self, - plan: &LogicalPlan, - output_idx: usize, - target_type: &DataType, - ) -> Result> { - let inputs = plan.inputs(); - let [input] = inputs.as_slice() else { - return Ok(None); - }; - let Some(rewritten_input) = self.rewrite_output_column(input, output_idx, target_type)? - else { - return Ok(None); - }; - - plan.with_new_exprs(plan.expressions_consider_join(), vec![rewritten_input]) - .map(Some) - } - - fn convert_literal(&self, value: &ScalarValue, target_type: &DataType) -> Option { - let ScalarValue::Utf8(Some(value)) = value else { - return None; - }; - cast_string_to_timestamp(value, target_type, Some(&self.query_ctx.timezone())) - .ok() - .filter(|value| !value.is_null()) - .map(|value| Expr::Literal(value, None)) - } -} - -fn assignment_input_column(expr: &Expr) -> Option<&Column> { - let expr = match unalias(expr) { - Expr::Cast(Cast { expr, .. }) => expr.as_ref(), - expr => expr, - }; - let Expr::Column(column) = expr else { - return None; - }; - Some(column) -} - -fn unalias(expr: &Expr) -> &Expr { - match expr { - Expr::Alias(Alias { expr, .. }) => unalias(expr), - expr => expr, - } -} - -#[cfg(test)] -mod tests { - use datafusion_common::ScalarValue; - use datafusion_common::arrow::datatypes::TimeUnit; - use session::context::QueryContext; - - use super::*; - - #[test] - fn test_convert_literal_falls_back_for_unsupported_literal() { - let converter = InsertAssignmentConverter { - query_ctx: QueryContext::arc(), - }; - let target_type = DataType::Timestamp(TimeUnit::Nanosecond, None); - - for literal in ["1970-01-01", "-8-01-01 00:00:01.5"] { - assert_eq!( - converter - .convert_literal(&ScalarValue::Utf8(Some(literal.to_string())), &target_type,), - None - ); - } - } -} diff --git a/src/query/src/planner.rs b/src/query/src/planner.rs index a7d4a84cb28..c3871a96afa 100644 --- a/src/query/src/planner.rs +++ b/src/query/src/planner.rs @@ -866,6 +866,30 @@ mod tests { engine.planner().plan(&stmt, query_ctx).await.unwrap() } + /// Plans `sql` and runs the DataFusion analyzer, which is where + /// `InsertAssignmentRule` sits. Planning alone stops short of it, so these + /// assertions would not see the assignment rewrite at all. + async fn analyze_insert( + engine: &QueryEngineRef, + sql: &str, + query_ctx: &QueryContextRef, + ) -> String { + let stmt = QueryLanguageParser::parse_sql(sql, query_ctx).unwrap(); + let plan = engine + .planner() + .plan(&stmt, query_ctx.clone()) + .await + .unwrap(); + let context = engine.engine_context(query_ctx.clone()); + let state = context.state(); + state + .analyzer() + .execute_and_check(plan, state.config_options(), |_, _| {}) + .unwrap() + .display_indent_schema() + .to_string() + } + #[tokio::test] async fn test_insert_timestamp_literals_use_query_timezone() { let query_ctx = Arc::new( @@ -898,20 +922,6 @@ mod tests { ) AS source", &[1_785_902_400_001_i64][..], ), - ( - "INSERT INTO timestamps (ts, st) \ - SELECT '2026-08-06 12:00:00.001', now() \ - UNION ALL \ - SELECT '2026-08-07 12:00:00.001', now()", - &[1_785_988_800_001_i64, 1_786_075_200_001_i64][..], - ), - ( - "INSERT INTO timestamps (ts, st) \ - SELECT '2026-08-16 12:00:00.001', now() \ - UNION \ - SELECT '2026-08-17 12:00:00.001', now()", - &[1_786_852_800_001_i64, 1_786_939_200_001_i64][..], - ), ( "INSERT INTO timestamps (ts, st) \ SELECT '2026-08-18 12:00:00.001', max(st) \ @@ -940,14 +950,7 @@ mod tests { &[1_786_766_400_001_i64][..], ), ] { - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx.clone()) - .await - .unwrap() - .display_indent() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; for expected_timestamp in expected_timestamps { assert!( @@ -968,19 +971,19 @@ mod tests { let engine = create_timestamp_test_engine().await; let sql = "INSERT INTO timestamps (ts, st) \ VALUES (CAST('2026-08-08 12:00:00.001' AS TIMESTAMP), now())"; - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx) - .await - .unwrap() - .display_indent() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; + // An explicit cast reaches the analyzer as an `arrow_cast` call rather + // than an `Expr::Cast`, which is how it stays out of the rewrite. assert!( plan.contains("arrow_cast(Utf8(\"2026-08-08 12:00:00.001\")"), "{plan}" ); + // 12:00:00.001 read as Shanghai local time; the source query keeps UTC. + assert!( + !plan.contains("TimestampMillisecond(1786104000001, None)"), + "{plan}" + ); } #[tokio::test] @@ -1017,14 +1020,7 @@ mod tests { ][..], ), ] { - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx.clone()) - .await - .unwrap() - .display_indent_schema() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; for expected in expected { assert!(plan.contains(expected), "{plan}"); @@ -1033,30 +1029,37 @@ mod tests { } #[tokio::test] - async fn test_insert_union_tolerates_uncoerced_untouched_column() { + async fn test_insert_union_converts_via_assignment_cast() { let query_ctx = Arc::new( QueryContextBuilder::default() .timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap()) .build(), ); let engine = create_timestamp_test_engine().await; - // `st` stays Timestamp vs Null across branches until TypeCoercion runs. - let sql = "INSERT INTO timestamps (ts, st) \ - SELECT '2026-08-06 12:00:00.001', now() \ - UNION ALL \ - SELECT '2026-08-07 12:00:00.001', NULL"; - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx) - .await - .unwrap() - .display_indent() - .to_string(); + // Branches disagree, so the conversion stays as a cast on the + // assignment instead of folding. One cast covers every branch, which is + // why a NULL branch no longer cancels the conversion for the column and + // why UNION's dedup keys stay on the original strings. + for sql in [ + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-06 12:00:00.001', now() \ + UNION ALL \ + SELECT '2026-08-07 12:00:00.001', NULL", + "INSERT INTO timestamps (ts, st) \ + SELECT '2026-08-16 12:00:00.001', now() \ + UNION \ + SELECT '2026-08-17 12:00:00.001', now()", + ] { + let plan = analyze_insert(&engine, sql, &query_ctx).await; - for expected_timestamp in [1_785_988_800_001_i64, 1_786_075_200_001_i64] { assert!( - plan.contains(&format!("TimestampMillisecond({expected_timestamp}, None)")), + plan.contains("AS Timestamp(ms, \"Asia/Shanghai\")"), + "{plan}" + ); + // The branches themselves are untouched. + assert!( + plan.contains("Utf8(\"2026-08-07 12:00:00.001\")") + || plan.contains("Utf8(\"2026-08-17 12:00:00.001\")"), "{plan}" ); } @@ -1074,14 +1077,8 @@ mod tests { SELECT '2026-08-10 12:00:00.001', now() \ UNION ALL \ SELECT CAST('2026-08-11 12:00:00.001' AS TIMESTAMP), now()"; - let stmt = QueryLanguageParser::parse_sql(sql, &query_ctx).unwrap(); - let plan = engine - .planner() - .plan(&stmt, query_ctx) - .await - .unwrap() - .display_indent_schema() - .to_string(); + let plan = analyze_insert(&engine, sql, &query_ctx).await; + assert!( !plan.contains("TimestampMillisecond(1786334400001, None)"), "{plan}" @@ -1090,6 +1087,11 @@ mod tests { plan.contains("arrow_cast(Utf8(\"2026-08-11 12:00:00.001\")"), "{plan}" ); + // TypeCoercion has already settled this union to timestamp, so the + // assignment has nothing left to reinterpret. Retargeting the cast here + // would leave a Timestamp(None) -> Timestamp(Some(tz)) step behind, + // which shifts the value instead of relabelling it. + assert!(!plan.contains("Asia/Shanghai"), "{plan}"); } #[tokio::test] diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.result b/tests/cases/standalone/common/insert/insert_default_timezone.result index 413eb5d46ae..41d8ca6cba8 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.result +++ b/tests/cases/standalone/common/insert/insert_default_timezone.result @@ -114,6 +114,27 @@ INSERT INTO test3 (ts, st, ts_ns) VALUES ( Affected Rows: 1 +-- a NULL branch must not cancel the conversion for the whole column +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-13 12:00:00.001' AS a, '2026-08-13 12:00:00.123456789' AS b +UNION ALL +SELECT '2026-08-14 12:00:00.001', NULL; + +Affected Rows: 2 + +-- NULL in the first branch: the union's schema starts out as Null +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-15 12:00:00.001' AS a, NULL AS b +UNION ALL +SELECT '2026-08-19 12:00:00.001', '2026-08-19 12:00:00.123456789'; + +Affected Rows: 2 + +-- the assignment cast also lands on non-literal VALUES expressions +INSERT INTO test3 (ts, st) VALUES (concat('2026-08-20 ', '12:00:00.001'), now()); + +Affected Rows: 1 + SELECT ts, ts_ns FROM test3 ORDER BY ts; +-------------------------+-------------------------------+ @@ -130,10 +151,43 @@ SELECT ts, ts_ns FROM test3 ORDER BY ts; | 2026-08-10T12:00:00.001 | | | 2026-08-11T12:00:00.001 | | | 2026-08-12T04:00:00.123 | 2026-08-12T04:00:00.123456789 | +| 2026-08-13T04:00:00.001 | 2026-08-13T04:00:00.123456789 | +| 2026-08-14T04:00:00.001 | | +| 2026-08-15T04:00:00.001 | | | 2026-08-16T04:00:00.001 | | | 2026-08-17T04:00:00.001 | | +| 2026-08-19T04:00:00.001 | 2026-08-19T04:00:00.123456789 | +| 2026-08-20T04:00:00.001 | | +-------------------------+-------------------------------+ +-- UNION dedup keys must stay on the source strings: these two spell the same +-- instant differently, so the source query yields two rows and both are kept. +CREATE TABLE test4 (ts TIMESTAMP TIME INDEX) WITH ('append_mode'='true'); + +Affected Rows: 0 + +INSERT INTO test4 (ts) +SELECT '2026-08-06 04:00:00' UNION SELECT '2026-08-06 04:00:00.000'; + +Affected Rows: 2 + +SELECT count(*) FROM test4; + ++----------+ +| count(*) | ++----------+ +| 2 | ++----------+ + +SELECT ts FROM test4 ORDER BY ts; + ++---------------------+ +| ts | ++---------------------+ +| 2026-08-05T20:00:00 | +| 2026-08-05T20:00:00 | ++---------------------+ + SET time_zone = 'UTC'; Affected Rows: 0 @@ -150,3 +204,7 @@ DROP TABLE test3; Affected Rows: 0 +DROP TABLE test4; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/insert/insert_default_timezone.sql b/tests/cases/standalone/common/insert/insert_default_timezone.sql index 7cbef15d6d9..f7c196699b7 100644 --- a/tests/cases/standalone/common/insert/insert_default_timezone.sql +++ b/tests/cases/standalone/common/insert/insert_default_timezone.sql @@ -56,8 +56,34 @@ INSERT INTO test3 (ts, st, ts_ns) VALUES ( '2026-08-09 12:00:00.123456789' ); +-- a NULL branch must not cancel the conversion for the whole column +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-13 12:00:00.001' AS a, '2026-08-13 12:00:00.123456789' AS b +UNION ALL +SELECT '2026-08-14 12:00:00.001', NULL; + +-- NULL in the first branch: the union's schema starts out as Null +INSERT INTO test3 (ts, ts_ns) +SELECT '2026-08-15 12:00:00.001' AS a, NULL AS b +UNION ALL +SELECT '2026-08-19 12:00:00.001', '2026-08-19 12:00:00.123456789'; + +-- the assignment cast also lands on non-literal VALUES expressions +INSERT INTO test3 (ts, st) VALUES (concat('2026-08-20 ', '12:00:00.001'), now()); + SELECT ts, ts_ns FROM test3 ORDER BY ts; +-- UNION dedup keys must stay on the source strings: these two spell the same +-- instant differently, so the source query yields two rows and both are kept. +CREATE TABLE test4 (ts TIMESTAMP TIME INDEX) WITH ('append_mode'='true'); + +INSERT INTO test4 (ts) +SELECT '2026-08-06 04:00:00' UNION SELECT '2026-08-06 04:00:00.000'; + +SELECT count(*) FROM test4; + +SELECT ts FROM test4 ORDER BY ts; + SET time_zone = 'UTC'; DROP TABLE test1; @@ -65,3 +91,5 @@ DROP TABLE test1; DROP TABLE test2; DROP TABLE test3; + +DROP TABLE test4; From 6d86e6ff068fe8471c3d7a28e70f9144e7684e15 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Tue, 25 Aug 2026 13:08:44 +0000 Subject: [PATCH 85/92] feat: synthesize OTLP resource descriptor for the semantic entity graph (#8904) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(servers): compose OTLP metrics job from service.namespace/service.name The OTel Prometheus compatibility spec defines job as "/" when the namespace is present. The OTLP metrics path only used the bare service.name, so the job tag diverged from target_info produced by Prometheus-side exporters for the same resource. Compose the namespace form, and keep not fabricating a job when service.name is absent. Behavior change: resources carrying service.namespace now get "namespace/name" as their job tag value. Signed-off-by: Dennis Zhuang * feat(otlp): synthesize otel_resource_info at OTLP metrics ingestion Ordinary OTLP metrics scatter filtered resource attributes as tags over every logical metric table, so metrics-only services contribute nothing to the semantic entity graph. Each request now also projects its distinct resources into one info-metric-shaped mito table, otel_resource_info: a fixed allowlist of identity-relevant attributes under their raw OTel keys (independent of the label translation strategy and the promote/ignore headers) plus derived job/instance compatibility columns, value 1.0, and the newest data-point timestamp. The descriptor is written after the main insert is committed; a failure there (conflicting pre-existing table, auto-create disabled) degrades to an OTLP partial_success warning with rejected_data_points = 0 instead of failing the request and triggering client retries of already-accepted data. A request writing a metric named otel_resource_info suppresses synthesis. Legacy mode is unchanged. Signed-off-by: Dennis Zhuang * feat(operator): otel info-metric conventions with host/container entities Whitelist the ingestion-synthesized otel_resource_info descriptor via a new otel_info_metrics conventions map, gated on source=opentelemetry (the existing gate hardcoded source=prometheus). Its declarations use explicit descriptive lists instead of descriptive_rest so identifying attributes of other entities do not leak into service.instance. Conventions tightened per the Astronomy Shop findings: host identity is host.id with host.name descriptive only (host.name is not stable across SDKs and resource detectors), a generic container entity (new entity type) is declared only when container.id is present, and trace-v1 tables now synthesize host/container from their flattened resource attributes too. New co-declared edges: service.instance runs_on container, container runs_on host. Signed-off-by: Dennis Zhuang * test(otlp): cover the resource descriptor in integration tests Covers the descriptor's raw-key columns and info-metric options through the HTTP path, the namespace/name job composition end-to-end, column names being independent of the translation strategy, the allowlist excluding unlisted resource attributes, auto-create after a drop, the metric-name collision suppressing synthesis, and the partial-success warning (rejected_data_points = 0) when a pre-existing incompatible table fails the descriptor write while metric data is accepted. Signed-off-by: Dennis Zhuang * chore: cargo fmt Signed-off-by: Dennis Zhuang * fix(frontend): degrade descriptor permission denial to a warning A table-level permission policy denying otel_resource_info would have failed the whole OTLP metrics request because the descriptor's permission check ran before the main insert. The descriptor is derived enrichment: check its permission in the degrade path so a denial skips the write and surfaces as the partial-success warning, like any other descriptor write failure. Signed-off-by: Dennis Zhuang * fix(otlp): guard descriptor writes with semantic ownership markers A pre-existing schema-compatible table named otel_resource_info would silently receive descriptor rows while its missing semantic stamps kept it out of the entity graph. The descriptor write now requires the auto-created table's ownership markers (mito engine + signal_type + source + metric.type=info + metadata_quality=declared) and otherwise degrades to the partial-success warning; the entity-graph gate for the otel whitelist likewise requires metric.type=info, so a user table stamped with only signal/source no longer picks up implicit declarations. Also fold the descriptor write cost into the response and surface the degrade warning through the otel-arrow BatchStatus status_message. Integration tests pin the full marker set on auto-create and that an existing owned descriptor keeps accepting writes without degrading — a missing marker would otherwise silently stop every descriptor write after the first request. Signed-off-by: Dennis Zhuang * perf(otlp): build descriptor rows without the per-resource BTreeMap Projecting a resource allocated a BTreeMap and then collected it into the row key, and every attribute was matched against the allowlist by linear scan. Collect the tags into a Vec and sort once, and match the allowlist instead of scanning it. Measured on the conversion path: descriptor work drops 16-18%, from 10.6% to 8.9% of conversion CPU on the worst shape (1000 resources with 4 data points each), where the cost tracks resource count rather than data-point count. Also trims the comments and tests added with the descriptor to what carries information. Signed-off-by: Dennis Zhuang * test(otlp): pin the descriptor permission-denial degrade path A policy denying the descriptor table must not fail the metrics request, which the fix in 2401b3dd9c does but nothing covered. Verified as a regression guard by mutation: moving the permission check back before the main write makes this test fail. Signed-off-by: Dennis Zhuang * test(otlp): keep legacy mode covered after trimming the unit tests Trimming the descriptor tests dropped the only assertion that legacy mode skips the job/instance remap and the promote filter. Both alter the columns of tables already in use, so fold the check into the legacy conversion test rather than leaving it uncovered. Signed-off-by: Dennis Zhuang * fix(semantic-graph): stop encoding column names into composite entity ids A composite entity id rendered the identifying columns as sorted `col=value` pairs, so the same identity split into one entity per signal: a trace table names its columns service_name and resource_attributes.service.instance.id where a metric table names them job and instance. One service instance became two nodes with two parallel edge sets, breaking the walk from a trace to that instance's metrics. Render an id as its values in declared order instead, escaping the separator so components stay distinguishable, which is what single-column ids already did by keeping only the value. entity_id_attrs still carries the structured form. Values alone are not enough for a namespaced service: the metric side folds service.namespace into job while traces keep the bare name. Add qualified_by to the conventions so the trace declarations compose the namespace the same way, per the OTel rule that job is / or the bare name when the namespace is empty. A table without the namespace column keeps the unqualified identity rather than losing the declaration. Conventions validation now rejects one entity type declared with a different number of id columns by two sources, which would silently produce ids that can never match. Signed-off-by: Dennis Zhuang * style(otlp): import the parent module by crate path check-super-imports.py, part of the CI format gate, rejects a file-level `use super::`. Signed-off-by: Dennis Zhuang * feat(otlp): gate the resource descriptor, and fix what review found Synthesizing greptime_otel_resource_info creates and writes a table the user never sent, so it is now off unless otlp.experimental_enable_resource_info says otherwise. With it off the request costs exactly what it did before the descriptor existed: nothing is projected, no table is created, no write and no permission check happen. Tests run with it on. StandaloneOptions carried no otlp field, so the whole [otlp] section was silently dropped in standalone mode; map it through, or the new option (and trace_ingest_chunk_size before it) would do nothing there. Renamed from otel_resource_info: the greptime_ prefix marks the table as engine-managed and makes a collision with a user metric unlikely, which is what the pre-existing-table ownership check and its per-request catalog lookup were defending against. Both are gone. A request may carry data for several graph windows, but the descriptor folded every data-point time into one row at the newest of them, leaving the earlier windows with metric rows and no entities. Key the rows by window as well, and take the times from the data points the encoder actually writes: it drops exponential histograms, and a resource carrying nothing else was being described as an entity with no measurements. Projecting a resource cloned its attributes once per data point. Nest the windows under the attributes instead, so they are moved once per resource, and walk the data-point times through a visitor rather than collecting a Vec per metric. Also documents what the two maps key and hold, and lifts the projected attribute names to constants beside KEY_SERVICE_NAME. Signed-off-by: Dennis Zhuang * fix(otlp): skip the descriptor's work entirely when it is disabled The collision scan over the request's output tables ran even with the feature off. Short-circuit on the option instead, and update the config snapshot the new [otlp] section changed. Also drops the doc comment orphaned by the deleted ownership check: it had attached itself to the trait impl and described a check that no longer exists. Signed-off-by: Dennis Zhuang * refactor(semantic-graph): drop the expect and name the service identity The CASE is built through Case directly rather than the fallible when().otherwise() builder, so the non-test path no longer carries an expect (architecture-invariants $4). service_identity returned two same-typed Options that both call sites destructured positionally; a named struct makes a swap fail to compile. Also records that id-column order is part of the identity, where the option docs and the conventions authors will read it. Signed-off-by: Dennis Zhuang * chore(semantic-graph): drop comments that narrate the code Signed-off-by: Dennis Zhuang * fix(semantic-graph): cast duration_nano before the trace-table union Trace tables written before the signed-integer ingest change hold duration_nano as UInt64 and later ones as Int64. The calls derivation unions the per-table selects, and the two have no common integer type, so a deployment holding both shapes could not build the plan. The cross-table test now spans both. Signed-off-by: Dennis Zhuang * refactor(semantic-graph): drop the redundant duration_nano casts Signed-off-by: Dennis Zhuang * fix(otlp): decide exponential histogram acceptance in one place The resource descriptor mirrored only the experimental gate, so with both experimental flags on a resource whose only metric is a delta exponential histogram was described as an entity with no measurements. The encoder's whole-metric rules move into exponential_histogram_gate, which both call, and the descriptor takes its timestamps through exponential_histogram_value so per-point rejections drop out too. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- config/config.md | 2 + config/frontend.example.toml | 3 + config/standalone.example.toml | 3 + ...25-entity-relationships-and-graph-query.md | 8 +- src/common/catalog/src/consts.rs | 5 + src/frontend/src/instance.rs | 1 + src/frontend/src/instance/builder.rs | 1 + src/frontend/src/instance/entity_graph.rs | 230 ++++++++- src/frontend/src/instance/otlp.rs | 47 +- src/frontend/src/metrics.rs | 8 + src/frontend/src/service_config/otlp.rs | 7 + src/operator/src/statement/semantic_graph.rs | 202 ++++++-- .../statement/semantic_graph/conventions.rs | 74 ++- .../statement/semantic_graph/conventions.yaml | 50 +- .../statement/semantic_graph/relationships.rs | 79 ++- src/servers/src/http/otlp.rs | 5 +- src/servers/src/otlp/metrics.rs | 467 ++++++++++++++---- src/servers/src/otlp/metrics/resource_info.rs | 393 +++++++++++++++ src/servers/src/otlp/trace.rs | 8 + src/session/src/protocol_ctx.rs | 3 + src/table/src/requests/semantic.rs | 4 + tests-integration/src/instance.rs | 115 +++++ tests-integration/src/standalone.rs | 5 + tests-integration/tests/http.rs | 220 ++++++++- .../common/system/semantic_graph.result | 261 ++++++++-- .../common/system/semantic_graph.sql | 119 +++++ tests/conf/frontend-test.toml.template | 3 + tests/conf/standalone-test.toml.template | 3 + 28 files changed, 2093 insertions(+), 233 deletions(-) create mode 100644 src/servers/src/otlp/metrics/resource_info.rs diff --git a/config/config.md b/config/config.md index 91cb4fe7669..992296d4827 100644 --- a/config/config.md +++ b/config/config.md @@ -76,6 +76,7 @@ | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | | `otlp.experimental_enable_exponential_histogram` | Bool | `false` | Experimental: enable cumulative OTLP exponential histogram ingestion. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | +| `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric
resource attributes, so metrics-only services reach the semantic graph. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | | `prom_store.with_metric_engine` | Bool | `true` | Whether to store the data from Prometheus remote write in metric engine. | @@ -319,6 +320,7 @@ | `otlp.enable` | Bool | `true` | Whether to enable OpenTelemetry protocol in HTTP API. | | `otlp.experimental_enable_exponential_histogram` | Bool | `false` | Experimental: enable cumulative OTLP exponential histogram ingestion. | | `otlp.trace_ingest_chunk_size` | Integer | `512` | Maximum spans per trace ingest chunk. Set to 0 to disable splitting. | +| `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric
resource attributes, so metrics-only services reach the semantic graph. | | `prom_store` | -- | -- | Prometheus remote storage options | | `prom_store.enable` | Bool | `true` | Whether to enable Prometheus remote write and read in HTTP API. | | `prom_store.with_metric_engine` | Bool | `true` | Whether to store the data from Prometheus remote write in metric engine. | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 3635a63f063..80534855721 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -241,6 +241,9 @@ enable = true experimental_enable_exponential_histogram = false ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 +## Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric +## resource attributes, so metrics-only services reach the semantic graph. +experimental_enable_resource_info = false ## Prometheus remote storage options [prom_store] diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 72ba992cce0..70f74c7f946 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -208,6 +208,9 @@ enable = true experimental_enable_exponential_histogram = false ## Maximum spans per trace ingest chunk. Set to 0 to disable splitting. trace_ingest_chunk_size = 512 +## Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric +## resource attributes, so metrics-only services reach the semantic graph. +experimental_enable_resource_info = false ## Prometheus remote storage options [prom_store] diff --git a/docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md b/docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md index 995cf6cab7e..1c78e91f252 100644 --- a/docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md +++ b/docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md @@ -149,7 +149,7 @@ The decisions behind the shape: - **`provenance` is part of edge identity**, so a declared edge and a derived edge between the same pair coexist, and a user-declared edge survives when derivation is off. - **`confidence` expresses derivation certainty, not statistical completeness**: `1.0` for a successfully paired or declared edge, `< 1.0` for virtual-node or agent-inferred edges. It does not correct for trace sampling. - **RED metrics describe the observed span-pair population** — the `service_graph` connector semantics. Under sampling, counts understate true traffic, and ratios (error rate, mean duration) are representative only when sampling is unbiased with respect to status and latency — tail sampling that keeps errors and slow traces skews both. -- **Endpoint encoding (v1).** A single-attribute identity is the value verbatim; a composite identity is the attributes sorted by key, rendered `k1=v1,k2=v2` — a human-readable traversal key, not a collision-free canonical encoding. Nodes carry the structured `entity_id_attrs`; edges carry only the string key. Known v1 limitation: composite components containing unescaped `,` or `=` are not guaranteed to be distinguished. +- **Endpoint encoding.** An identity renders as its values in declared order (broad to narrow), joined by `,` with `,` and `\` escaped in each value. The identifying attribute *keys* are deliberately absent: the same identity reaches the graph under different column names per signal — a trace table's `service_name` against a metric table's `job` — and encoding the keys would split one entity into one entity per signal. Where a source carries a qualifier separately (`service.namespace` beside `service.name`) the convention folds it in as `/`, reproducing the identity a source that pre-composed it (Prometheus `job`) already emits. Nodes carry the structured `entity_id_attrs`; edges carry only the string key. - **Hand-declared edges** (`provenance = 'declared'`) live in a physical table of the same shape (plus a business validity window) in `greptime_private` and are unioned into the computed `semantic_relationships`. The relationship vocabulary is small, typed, and inverse-paired; stored direction is `src -> dst`, the inverse is a query concern: @@ -197,7 +197,7 @@ GROUP BY 1, src_id, dst_id; Edge endpoints are built from each trace table's `service` entity declaration (the same id expression the registry uses), so edges land on exactly the entity ids the registry emits, self-calls are judged on the full endpoint id, and a table whose declaration is unusable contributes no edges rather than silently switching identity. The join is bounded by a time-proximity window and the results of **all** trace tables combine into one edge per window with merged RED metrics. One assumption to note: a deployment can route traces to multiple tables (`x-greptime-trace-table-name`), so one distributed trace is not guaranteed to live in a single table — pairing must account for that, and how is an implementation concern. Uninstrumented peers become **virtual nodes**: a client span with no matching server span yields an edge to a synthetic node named from `peer.service`/`db.name`/`server.address`. -**The entity registry.** Per declaring table: filter to the window, project (time bin, casts, JSON assembly), `DISTINCT`. The per-table results are unioned without any cross-source merge: no descriptive-conflict or JSON-merging policy lives in the derivation. Rows whose identity columns are NULL identify nothing and are filtered out; single-column ids are used verbatim, composite ids render the sorted `k=v` form plus the `entity_id_attrs` JSON. Consumers deduplicate with `SELECT DISTINCT entity_type, entity_id`; a unique node set is the snapshot relation's job (see *Querying the graph*). +**The entity registry.** Per declaring table: filter to the window, project (time bin, casts, JSON assembly), `DISTINCT`. The per-table results are unioned without any cross-source merge: no descriptive-conflict or JSON-merging policy lives in the derivation. Rows whose identity columns are NULL identify nothing and are filtered out; ids render as their escaped values joined in declared order, with the `entity_id_attrs` JSON alongside once an id is assembled from more than one column. Consumers deduplicate with `SELECT DISTINCT entity_type, entity_id`; a unique node set is the snapshot relation's job (see *Querying the graph*). **Attribute edges.** Shared attributes provide join keys; co-declaration or a relationship template provides the relationship semantics. A shared value alone determines neither direction nor edge type, so the derivation is rule-based, not join-everything. A table declaring both a `service.instance` and a `host` identity on the same rows derives `runs_on` between them: the row itself witnesses the relationship, and the built-in vocabulary fixes the direction. Cross-table edges require a declared relationship template (endpoint mappings, `rel_type`, direction): a derivation *rule* declared on schema, not per-instance edge data. @@ -284,7 +284,7 @@ Drawbacks: - Read-time derivation costs a self-join over trace tables on every scan of the computed tables. The default one-hour window bounds it, but very large deployments will want materialisation (Future Work). - `semantic_entities` exposes per-observation rows; deduplication falls on the consumer (`DISTINCT`) until the snapshot relation lands in M2. -- The v1 `k=v` id encoding can collide on hostile values (unescaped `,`/`=`); hardening is an open question. +- The id encoding is readable rather than canonical: escaping keeps components distinguishable, but a value containing the qualifier separator (`/`) is still ambiguous against a qualified name. - The graph is only as connected as the identity columns tables actually share: the `job`/`service_name` misalignment above is the canonical example, and enrichment discipline gates connectivity. # Implementation Plan @@ -306,7 +306,7 @@ Drawbacks: # Open Questions -1. **Endpoint encoding hardening.** The v1 `k=v` rendering (with `entity_id_attrs` JSON as the source of truth) is readable but collides when values contain `,`/`=`. Open: a fixed-width `entity_id_hash` join key (New-Relic-style); typed value encoding; the criteria for introducing the `semantic_key=column` mapping form; `schema_url`-gated entity merge. Related and unresolved: whether the identifying attribute *key* should participate in entity equality — a single-attribute id keeps only the value, so `job=api` and `service_name=api` currently name the same entity even when the keys differ semantically. +1. **Endpoint encoding hardening.** The rendering (with `entity_id_attrs` JSON as the source of truth) is readable rather than canonical. Open: a fixed-width `entity_id_hash` join key (New-Relic-style); typed value encoding; the criteria for introducing the `semantic_key=column` mapping form; `schema_url`-gated entity merge. Settled: identifying attribute *keys* do not participate in entity equality, so `job=api` and `service_name=api` name the same entity — the property that lets one service reach the graph from several signals. 2. **Identity reconciliation across signals.** Trace edges key on the trace table's service declaration; host edges key on `host.id`/`k8s.pod.uid`. Linking them relies on tables carrying both columns; how aggressively to enrich (a `k8sattributes`-equivalent) gates how connected the graph is. 3. **`GRAPH_TABLE` surface scope.** How much of SQL/PGQ to implement vs. a minimal single-hop subset. Start minimal and grow with demand. 4. **Visibility beyond table grants.** The derivation contract filters by table-read permission; whether topology needs finer policies (row-level, tenant-scoped visibility of cross-schema edges) is open. diff --git a/src/common/catalog/src/consts.rs b/src/common/catalog/src/consts.rs index b976cbc256e..63aa103b83d 100644 --- a/src/common/catalog/src/consts.rs +++ b/src/common/catalog/src/consts.rs @@ -220,6 +220,11 @@ pub const SEMANTIC_RELATIONSHIPS_TABLE_NAME: &str = "semantic_relationships"; /// `semantic_relationships`; see [`is_ddl_reserved_table`] for its lifecycle. pub const SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME: &str = "semantic_relationships_declared"; +/// Width of the window the graph derivation bins observations into. Sources +/// synthesizing observations must land a row in every window they describe, +/// so this is shared rather than restated per crate. +pub const SEMANTIC_GRAPH_WINDOW_NANOS: i64 = 60 * 1_000_000_000; + // Column names of the graph tables: `catalog` exposes these schemas and the // read-time plans in `operator` must project exactly them. pub const OBSERVED_AT_COLUMN: &str = "observed_at"; diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 78b701a3442..412ff495f65 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -134,6 +134,7 @@ pub struct Instance { slow_query_options: SlowQueryOptions, influxdb_default_merge_mode: InfluxdbMergeMode, trace_ingest_chunk_size: usize, + otlp_resource_info: bool, suspend: Arc, // cache for otlp metrics diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index c49a8394beb..0f87dc424f0 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -367,6 +367,7 @@ impl FrontendBuilder { slow_query_options: self.options.slow_query.clone(), influxdb_default_merge_mode: self.options.influxdb.default_merge_mode, trace_ingest_chunk_size: self.options.otlp.trace_ingest_chunk_size, + otlp_resource_info: self.options.otlp.experimental_enable_resource_info, suspend: Arc::new(AtomicBool::new(false)), }) } diff --git a/src/frontend/src/instance/entity_graph.rs b/src/frontend/src/instance/entity_graph.rs index 14f18ed1b79..a0817cf2806 100644 --- a/src/frontend/src/instance/entity_graph.rs +++ b/src/frontend/src/instance/entity_graph.rs @@ -23,7 +23,7 @@ //! the query engine. Injected into the catalog manager after the engine is built, //! breaking the `catalog -> query` cycle. -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::sync::Weak; use async_trait::async_trait; @@ -60,8 +60,9 @@ use table::TableRef; use table::metadata::TableInfo; use table::predicate::{TimeRangeExtraction, extract_time_range_strict}; use table::requests::{ - EntityRole, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_METRIC, SOURCE_PROMETHEUS, - is_trace_v1_table, parse_entity_columns, parse_entity_option_key, + EntityRole, SEMANTIC_METRIC_TYPE, SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_METRIC, + SOURCE_OPENTELEMETRY, SOURCE_PROMETHEUS, is_trace_v1_table, parse_entity_columns, + parse_entity_option_key, }; use crate::error; @@ -177,6 +178,7 @@ impl EntityGraphProviderImpl { time_index: time_index.clone(), entity_type, id_columns, + id_qualifier: None, descriptive_columns, scope_columns, }) @@ -188,10 +190,11 @@ impl EntityGraphProviderImpl { /// All entity declarations of one table: the explicit options plus the /// zero-configuration conventions (`otlp_trace_entities` for trace-v1 /// tables — including the `service` identity of tables created before the - /// ingest-side auto-stamp — and the Prometheus descriptor whitelist). An - /// explicit declaration of a type always suppresses the implicit one, even - /// when the explicit declaration is invalid and skipped: silently falling - /// back would change entity identity behind the user's back. + /// ingest-side auto-stamp — and the Prometheus/OTel descriptor + /// whitelists). An explicit declaration of a type always suppresses the + /// implicit one, even when the explicit declaration is invalid and + /// skipped: silently falling back would change entity identity behind the + /// user's back. fn declarations_for( table_info: &TableInfo, conventions: &Conventions, @@ -204,7 +207,20 @@ impl EntityGraphProviderImpl { &mut declarations, ); } - Self::extend_with_prometheus_conventions(table_info, conventions, &mut declarations); + Self::extend_with_info_metric_conventions( + table_info, + &conventions.prometheus_info_metrics, + SOURCE_PROMETHEUS, + None, + &mut declarations, + ); + Self::extend_with_info_metric_conventions( + table_info, + &conventions.otel_info_metrics, + SOURCE_OPENTELEMETRY, + Some(servers::semantic::METRIC_TYPE_INFO), + &mut declarations, + ); declarations } @@ -216,29 +232,34 @@ impl EntityGraphProviderImpl { }) } - /// Implicit declarations of the well-known Prometheus entity-descriptor - /// metrics (the `prometheus_info_metrics` whitelist of `conventions.yaml`), - /// gated on the ingest-stamped `signal_type=metric` + `source=prometheus` - /// options and keyed by table name. The metric engine's physical table + /// Implicit declarations of the well-known entity-descriptor metrics + /// (the `prometheus_info_metrics` / `otel_info_metrics` whitelists of + /// `conventions.yaml`), gated on the ingest-stamped `signal_type=metric` + /// option plus the whitelist's expected `source`; OTel descriptors also + /// require `metric.type=info`. The metric engine's physical table /// aggregates every logical table's columns and must not contribute a /// duplicate source. - fn extend_with_prometheus_conventions( + fn extend_with_info_metric_conventions( table_info: &TableInfo, - conventions: &Conventions, + whitelist: &BTreeMap>, + expected_source: &str, + expected_metric_type: Option<&str>, declarations: &mut Vec, ) { - let Some(implicit_entities) = conventions.prometheus_info_metrics.get(&table_info.name) - else { + let Some(implicit_entities) = whitelist.get(&table_info.name) else { return; }; let options = &table_info.meta.options.extra_options; if options.get(SEMANTIC_SIGNAL_TYPE).map(String::as_str) != Some(SIGNAL_TYPE_METRIC) - || options.get(SEMANTIC_SOURCE).map(String::as_str) != Some(SOURCE_PROMETHEUS) + || options.get(SEMANTIC_SOURCE).map(String::as_str) != Some(expected_source) + || expected_metric_type.is_some_and(|expected| { + options.get(SEMANTIC_METRIC_TYPE).map(String::as_str) != Some(expected) + }) || table_info.is_physical_table() { debug!( - "Table `{}` matches the info-metric whitelist but is not a prometheus logical \ - metric table; skipping its implicit declarations", + "Table `{}` matches the info-metric whitelist but is not an eligible \ + `{expected_source}` info-metric source; skipping its implicit declarations", table_info.name ); return; @@ -296,12 +317,19 @@ impl EntityGraphProviderImpl { .cloned() .collect() }; + // A table predating the qualifier column keeps the unqualified + // identity rather than losing the declaration. + let id_qualifier = implicit + .qualified_by + .clone() + .filter(|c| schema.column_schema_by_name(c).is_some()); declarations.push(EntityDeclaration { schema: table_info.schema_name.clone(), table: table_info.name.clone(), time_index: time_index.clone(), entity_type: implicit.entity.clone(), id_columns: implicit.id.clone(), + id_qualifier, descriptive_columns, scope_columns: vec![], }); @@ -832,6 +860,24 @@ mod tests { declarations[2].id_columns, vec!["service_name", "resource_attributes.service.instance.id"] ); + assert_eq!(declarations[1].id_qualifier, None); + + let namespaced = table_info( + &[ + "service_name", + "resource_attributes.service.namespace", + "resource_attributes.service.instance.id", + ], + &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)], + ); + for declaration in sorted_declarations(&namespaced) { + assert_eq!( + declaration.id_qualifier.as_deref(), + Some("resource_attributes.service.namespace"), + "{} must qualify its identity like the metric side's job", + declaration.entity_type + ); + } // Missing uid column: no pod entity synthesized, no name-based guess. let no_uid = table_info( @@ -856,6 +902,47 @@ mod tests { assert!(sorted_declarations(&invalid_explicit).is_empty()); } + #[test] + fn trace_table_host_and_container_require_stable_ids() { + let with_ids = table_info( + &[ + "service_name", + "resource_attributes.host.id", + "resource_attributes.host.name", + "resource_attributes.container.id", + ], + &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)], + ); + let declarations = sorted_declarations(&with_ids); + let types: Vec<&str> = declarations + .iter() + .map(|d| d.entity_type.as_str()) + .collect(); + assert_eq!(types, vec!["container", "host", "service"]); + assert_eq!( + declarations[1].id_columns, + vec!["resource_attributes.host.id"] + ); + assert_eq!( + declarations[1].descriptive_columns, + vec!["resource_attributes.host.name"] + ); + + let names_only = table_info( + &[ + "service_name", + "resource_attributes.host.name", + "resource_attributes.container.name", + ], + &[(TABLE_DATA_MODEL, TABLE_DATA_MODEL_TRACE_V1)], + ); + let types: Vec = sorted_declarations(&names_only) + .into_iter() + .map(|d| d.entity_type) + .collect(); + assert_eq!(types, vec!["service"]); + } + #[test] fn prometheus_implicit_declarations_are_gated() { let labels: &[&str] = &["namespace", "pod", "node"]; @@ -900,6 +987,111 @@ mod tests { assert_eq!(declarations[1].id_columns, vec!["pod"]); } + const OTEL_STAMPS: &[(&str, &str)] = &[ + (SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC), + (SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY), + (SEMANTIC_METRIC_TYPE, servers::semantic::METRIC_TYPE_INFO), + ]; + + #[test] + fn greptime_otel_resource_info_gets_implicit_declarations() { + let info = prom_table_info( + "greptime_otel_resource_info", + &[ + "job", + "instance", + "service.name", + "service.namespace", + "host.id", + "host.name", + "container.id", + "container.name", + "k8s.pod.uid", + "k8s.pod.name", + "k8s.namespace.name", + ], + OTEL_STAMPS, + ); + let declarations = sorted_declarations(&info); + let types: Vec<&str> = declarations + .iter() + .map(|d| d.entity_type.as_str()) + .collect(); + assert_eq!( + types, + vec![ + "container", + "host", + "k8s.pod", + "service", + "service.instance" + ] + ); + assert_eq!(declarations[1].id_columns, vec!["host.id"]); + assert_eq!(declarations[1].descriptive_columns, vec!["host.name"]); + assert_eq!(declarations[3].id_columns, vec!["job"]); + assert_eq!( + declarations[3].descriptive_columns, + vec!["service.name", "service.namespace"] + ); + assert_eq!(declarations[4].id_columns, vec!["job", "instance"]); + assert!(declarations[4].descriptive_columns.is_empty()); + + let partial = prom_table_info( + "greptime_otel_resource_info", + &["job", "service.name", "host.id"], + OTEL_STAMPS, + ); + let types: Vec = sorted_declarations(&partial) + .into_iter() + .map(|d| d.entity_type) + .collect(); + assert_eq!(types, vec!["host", "service"]); + } + + #[test] + fn otel_implicit_declarations_are_gated() { + let labels: &[&str] = &["job", "instance", "host.id"]; + assert!( + sorted_declarations(&prom_table_info( + "greptime_otel_resource_info", + labels, + PROM_STAMPS + )) + .is_empty() + ); + let mut stamps = OTEL_STAMPS.to_vec(); + stamps.push((PHYSICAL_TABLE_METADATA_KEY, "true")); + assert!( + sorted_declarations(&prom_table_info( + "greptime_otel_resource_info", + labels, + &stamps + )) + .is_empty() + ); + assert!( + conventions() + .unwrap() + .otel_info_metrics + .contains_key(servers::otlp::metrics::OTEL_RESOURCE_INFO_TABLE_NAME) + ); + + let wrong_type = [ + (SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC), + (SEMANTIC_SOURCE, SOURCE_OPENTELEMETRY), + (SEMANTIC_METRIC_TYPE, servers::semantic::METRIC_TYPE_GAUGE), + ]; + assert!( + sorted_declarations(&prom_table_info( + "greptime_otel_resource_info", + labels, + &wrong_type + )) + .is_empty() + ); + } + #[test] fn target_info_descriptive_rest_covers_remaining_tags() { let info = prom_table_info( diff --git a/src/frontend/src/instance/otlp.rs b/src/frontend/src/instance/otlp.rs index 08086198536..ccc09c9978c 100644 --- a/src/frontend/src/instance/otlp.rs +++ b/src/frontend/src/instance/otlp.rs @@ -50,7 +50,7 @@ use table::requests::{ use self::trace_ingest::trace_conventions; use crate::instance::Instance; -use crate::metrics::{OTLP_LOGS_ROWS, OTLP_METRICS_ROWS}; +use crate::metrics::{OTLP_LOGS_ROWS, OTLP_METRICS_ROWS, OTLP_RESOURCE_INFO_WRITE_ERRORS}; fn trace_permission_targets( table_name: &str, @@ -119,9 +119,15 @@ impl OpenTelemetryProtocolHandler for Instance { .cloned() .unwrap_or_default(); metric_ctx.is_legacy = is_legacy; + metric_ctx.resource_info = self.otlp_resource_info; - let (requests, rows, semantic_index, mut outcome) = - otlp::metrics::to_grpc_insert_requests(request, &mut metric_ctx)?; + let otlp::metrics::MetricsConversion { + requests, + rows, + semantic_index, + resource_info, + mut outcome, + } = otlp::metrics::to_grpc_insert_requests(request, &mut metric_ctx)?; if outcome.rejected_data_points > 0 { warn!( "Rejected {} OTLP exponential histogram data points: {}", @@ -156,7 +162,7 @@ impl OpenTelemetryProtocolHandler for Instance { // OTLP tables have one sample field in both the legacy and physical paths. let output = if metric_ctx.is_legacy || !metric_ctx.with_metric_engine { - self.handle_row_inserts(requests, ctx, false, true) + self.handle_row_inserts(requests, ctx.clone(), false, true) .await .map_err(BoxedError::new) .context(error::ExecuteGrpcQuerySnafu) @@ -165,12 +171,43 @@ impl OpenTelemetryProtocolHandler for Instance { .extension(PHYSICAL_TABLE_PARAM) .unwrap_or(GREPTIME_PHYSICAL_TABLE) .to_string(); - self.handle_metric_row_inserts(requests, ctx, physical_table) + self.handle_metric_row_inserts(requests, ctx.clone(), physical_table) .await .map_err(BoxedError::new) .context(error::ExecuteGrpcQuerySnafu) }?; outcome.write_cost = output.meta.cost; + + // Derived enrichment, written after the metric data is committed: + // failing here would make the client retry data the server already + // accepted, so every failure degrades to a warning instead. + if let Some(resource_info) = resource_info { + let written = match self.check_row_insert_permission( + &resource_info, + &ctx, + PermissionReq::Action(OTLP_WRITE), + ) { + Ok(_) => self + .handle_row_inserts(resource_info, ctx, false, false) + .await + .map_err(BoxedError::new) + .map_err(|e| e.to_string()), + Err(e) => Err(e.to_string()), + }; + match written { + Ok(descriptor_output) => outcome.write_cost += descriptor_output.meta.cost, + Err(e) => { + OTLP_RESOURCE_INFO_WRITE_ERRORS.inc(); + warn!("Failed to write the OTLP resource descriptor table: {e}"); + outcome.error_message.get_or_insert(format!( + "metric data was accepted, but writing the resource \ + descriptor table `{}` failed: {e}", + otlp::metrics::OTEL_RESOURCE_INFO_TABLE_NAME + )); + } + } + } + Ok(outcome) } diff --git a/src/frontend/src/metrics.rs b/src/frontend/src/metrics.rs index aba33637cfc..f649db0b46f 100644 --- a/src/frontend/src/metrics.rs +++ b/src/frontend/src/metrics.rs @@ -45,6 +45,14 @@ lazy_static! { ) .unwrap(); + /// Failed writes of the synthesized OTLP resource descriptor table; these + /// surface as OTLP partial-success warnings, not request failures. + pub static ref OTLP_RESOURCE_INFO_WRITE_ERRORS: IntCounter = register_int_counter!( + "greptime_frontend_otlp_resource_info_write_errors", + "frontend otlp resource descriptor write errors" + ) + .unwrap(); + /// The number of OpenTelemetry traces send by frontend node. pub static ref OTLP_TRACES_ROWS: IntCounter = register_int_counter!( "greptime_frontend_otlp_traces_rows", diff --git a/src/frontend/src/service_config/otlp.rs b/src/frontend/src/service_config/otlp.rs index 5dcab2c8ce2..d6a36f702e4 100644 --- a/src/frontend/src/service_config/otlp.rs +++ b/src/frontend/src/service_config/otlp.rs @@ -23,6 +23,11 @@ pub struct OtlpOptions { pub experimental_enable_exponential_histogram: bool, /// Maximum spans per trace ingest chunk. Set to 0 to disable splitting. pub trace_ingest_chunk_size: usize, + /// Whether to synthesize the `greptime_otel_resource_info` descriptor + /// table from the resource attributes of OTLP metrics, so metrics-only + /// services reach the semantic graph. Off by default: it creates and + /// writes a table the user did not send. + pub experimental_enable_resource_info: bool, } impl Default for OtlpOptions { @@ -31,6 +36,7 @@ impl Default for OtlpOptions { enable: true, experimental_enable_exponential_histogram: false, trace_ingest_chunk_size: DEFAULT_TRACE_INGEST_CHUNK_SIZE, + experimental_enable_resource_info: false, } } } @@ -45,6 +51,7 @@ mod tests { assert!(default.enable); assert!(!default.experimental_enable_exponential_histogram); assert_eq!(default.trace_ingest_chunk_size, 512); + assert!(!default.experimental_enable_resource_info); let options: OtlpOptions = toml::from_str("enable = false").unwrap(); assert!(!options.enable); diff --git a/src/operator/src/statement/semantic_graph.rs b/src/operator/src/statement/semantic_graph.rs index 7014a94f8ac..86827cf8c7a 100644 --- a/src/operator/src/statement/semantic_graph.rs +++ b/src/operator/src/statement/semantic_graph.rs @@ -41,8 +41,9 @@ use common_catalog::consts::{ ENTITY_DESCRIPTIVE_COLUMN, ENTITY_ID_ATTRS_COLUMN, ENTITY_ID_COLUMN, ENTITY_SCOPE_COLUMN, ENTITY_TYPE_COLUMN, ERROR_COUNT_COLUMN, FRESH_UNTIL_COLUMN, GENERATION_ID_COLUMN, OBSERVED_AT_COLUMN, PROVENANCE_COLUMN, REL_TYPE_COLUMN, REQUEST_COUNT_COLUMN, - SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SOURCE_TABLES_COLUMN, SRC_ID_COLUMN, - SRC_TYPE_COLUMN, VALID_FROM_COLUMN, VALID_UNTIL_COLUMN, WINDOW_END_COLUMN, WINDOW_START_COLUMN, + SEMANTIC_GRAPH_WINDOW_NANOS, SEMANTIC_RELATIONSHIPS_DECLARED_TABLE_NAME, SOURCE_TABLES_COLUMN, + SRC_ID_COLUMN, SRC_TYPE_COLUMN, VALID_FROM_COLUMN, VALID_UNTIL_COLUMN, WINDOW_END_COLUMN, + WINDOW_START_COLUMN, }; use common_function::function::FunctionContext; use common_function::function_registry::FUNCTION_REGISTRY; @@ -60,7 +61,7 @@ use datafusion::dataframe::DataFrame; use datafusion::functions::{core as core_fns, datetime as datetime_fns, string as string_fns}; use datafusion::functions_nested::expr_fn::make_array; use datafusion_common::{Column, Result as DfResult, ScalarValue}; -use datafusion_expr::{Expr, LogicalPlan, ScalarUDF, cast, ident, lit}; +use datafusion_expr::{Case, Expr, LogicalPlan, ScalarUDF, cast, ident, lit}; pub use relationships::{ CallsSource, CoDeclaredSource, DeclaredSource, RelationshipSources, build_relationships_plan, }; @@ -122,9 +123,10 @@ pub fn declared_relationships_schema_matches(table_info: &table::metadata::Table }) } -/// Bin width for the temporal window of derived rows: 60s buckets, matching the -/// service-graph convention. -const BIN_NANOS: i64 = 60 * 1_000_000_000; +/// Bin width for the temporal window of derived rows, matching the +/// service-graph convention. Shared so ingestion-synthesized observations +/// land in the same buckets. +const BIN_NANOS: i64 = SEMANTIC_GRAPH_WINDOW_NANOS; /// Default retention for the declared-edge table; expiry slides the topology window. const DEFAULT_DECLARED_RELATIONSHIPS_TTL: &str = "90d"; @@ -284,8 +286,12 @@ pub struct EntityDeclaration { /// The table's time index column, used for the temporal window filter. pub time_index: String, pub entity_type: String, - /// Identifying columns (>= 1). One column → id verbatim; several → composite. + /// Identifying columns (>= 1), ordered broad to narrow. pub id_columns: Vec, + /// Optional column qualifying the first identity component, so a source + /// carrying the parts separately (a trace table's `service.namespace`) + /// yields the same id as one carrying them pre-composed (`job`). + pub id_qualifier: Option, /// Descriptive columns snapshotted into the `descriptive` JSON (may be empty). pub descriptive_columns: Vec, /// Scope columns (namespace/environment). One column → scope verbatim; @@ -417,28 +423,91 @@ fn cast_string_or_empty(column: &str) -> Expr { core_fns::coalesce().call(vec![cast(ident(column), DataType::Utf8), lit("")]) } -/// The canonical entity-id expression for `id_columns`: the value verbatim for -/// a single column, the sorted `k=v,k=v` rendering for a composite. `col` -/// constructs the column reference (unqualified for registry branches, -/// join-side-qualified for the calls derivation). /// A row identifies an entity only when every identity component is present /// and non-empty: kube-state-metrics descriptors emit empty-string labels (an /// unscheduled pod's `node`, an owner-less pod's `owner_*`), and an empty -/// string is never a meaningful entity id. +/// string is never a meaningful entity id. The qualifier is optional by +/// construction and so is not guarded. pub(crate) fn identifies(column: &str) -> Expr { ident(column) .is_not_null() .and(cast(ident(column), DataType::Utf8).not_eq(lit(""))) } -fn entity_id_expr(id_columns: &[String], col: &dyn Fn(&str) -> Expr) -> Expr { - if let [id] = id_columns { - cast(col(id), DataType::Utf8) - } else { - let mut cols = id_columns.to_vec(); - cols.sort(); - sorted_kv_expr_with(&cols, false, col) +const ID_SEPARATOR: &str = ","; +const ID_ESCAPE: &str = "\\"; +/// Left unescaped: `/` is how Prometheus renders `job`. +const ID_QUALIFIER_SEPARATOR: &str = "/"; + +/// Escapes so a composite id decodes back to its components. The escape +/// character goes first, or it would double the escapes the separator pass +/// introduces. +fn escaped_id_value(value: Expr) -> Expr { + let escape_escape = string_fns::replace().call(vec![ + value, + lit(ID_ESCAPE), + lit(format!("{ID_ESCAPE}{ID_ESCAPE}")), + ]); + string_fns::replace().call(vec![ + escape_escape, + lit(ID_SEPARATOR), + lit(format!("{ID_ESCAPE}{ID_SEPARATOR}")), + ]) +} + +/// The identity values in declared order (broad to narrow), escaped and +/// joined. +/// +/// The identifying *column names* are deliberately absent: the same identity +/// reaches us under different names per source — a trace table's +/// `service_name` against a metric table's `job` — and encoding them would +/// split one entity into one per signal. `entity_id_attrs` keeps the +/// structured form. +/// +/// `col` constructs the column reference (unqualified for registry branches, +/// join-side-qualified for the calls derivation). +fn entity_id_expr( + id_columns: &[String], + qualifier: Option<&str>, + col: &dyn Fn(&str) -> Expr, +) -> Expr { + let mut parts = Vec::with_capacity(id_columns.len() * 2); + for (index, column) in id_columns.iter().enumerate() { + if index > 0 { + parts.push(lit(ID_SEPARATOR)); + } + let value = escaped_id_value(cast(col(column), DataType::Utf8)); + parts.push(match qualifier { + Some(qualifier) if index == 0 => qualified_id_expr(qualifier, value, col), + _ => value, + }); } + if let [single] = parts.as_slice() { + single.clone() + } else { + concat_expr(parts) + } +} + +/// `/`, or bare `value` when the qualifier is empty on this +/// row — the spec's rule composing `job` from `service.namespace` and +/// `service.name`. +fn qualified_id_expr(qualifier: &str, value: Expr, col: &dyn Fn(&str) -> Expr) -> Expr { + let qualifier = escaped_id_value( + core_fns::coalesce().call(vec![cast(col(qualifier), DataType::Utf8), lit("")]), + ); + Expr::Case(Case::new( + None, + vec![( + Box::new(qualifier.clone().eq(lit(""))), + Box::new(value.clone()), + )], + Some(Box::new(concat_expr(vec![ + qualifier, + lit(ID_QUALIFIER_SEPARATOR), + value, + ]))), + )) } /// The `parse_json` UDF, shared by all derivation plans. Resolved from the @@ -583,14 +652,22 @@ fn registry_source( let mut rows = Vec::with_capacity(1 + rest.len()); for decl in std::iter::once(first).chain(rest) { // CAST even a single-column id: id columns need not be strings, and - // the computed table declares entity_id STRING. Composite ids - // additionally carry a JSON object of the id columns in - // entity_id_attrs. - let entity_id = entity_id_expr(&decl.id_columns, &|c| ident(c)); - let entity_id_attrs = if decl.id_columns.len() == 1 { + // the computed table declares entity_id STRING. An id assembled from + // more than one column additionally carries its parts as a JSON + // object in entity_id_attrs, the structured form of the id string. + let entity_id = entity_id_expr(&decl.id_columns, decl.id_qualifier.as_deref(), &|c| { + ident(c) + }); + let id_parts = decl + .id_qualifier + .iter() + .chain(&decl.id_columns) + .cloned() + .collect::>(); + let entity_id_attrs = if id_parts.len() == 1 { null_json() } else { - json_object_expr(&decl.id_columns) + json_object_expr(&id_parts) }; let scope = match decl.scope_columns.as_slice() { @@ -811,11 +888,74 @@ mod tests { time_index: "ts".to_string(), entity_type: entity_type.to_string(), id_columns: id_columns.iter().map(|s| s.to_string()).collect(), + id_qualifier: None, descriptive_columns: vec![], scope_columns: vec![], } } + /// The id string must decode back to its components, and a qualifier must + /// reproduce the identity a source that pre-composed it already emits. + #[tokio::test] + async fn registry_id_escaping_and_qualifier() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("service_name", DataType::Utf8, false), + Field::new("instance", DataType::Utf8, false), + Field::new("namespace", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000, 2_000, 3_000])) as ArrayRef, + Arc::new(StringArray::from(vec!["cart", "a,b", "we\\ird"])), + Arc::new(StringArray::from(vec!["i-1", "c", "i-3"])), + Arc::new(StringArray::from(vec![Some("shop"), None, Some("")])), + ], + ) + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_table( + "svc", + Arc::new(MemTable::try_new(schema, vec![vec![batch]]).unwrap()), + ) + .unwrap(); + + let mut declaration = decl("service.instance", &["service_name", "instance"]); + declaration.id_qualifier = Some("namespace".to_string()); + declaration.time_index = "ts".to_string(); + let plan = build_registry_plan( + vec![RegistrySource { + declarations: vec![declaration], + scan: ctx.table("svc").await.unwrap(), + }], + &test_window(), + ) + .unwrap() + .unwrap(); + + let mut ids: Vec = collect(&ctx, plan) + .await + .iter() + .flat_map(|batch| strings(batch, 5)) + .collect(); + ids.sort(); + assert_eq!( + ids, + vec![ + // an absent qualifier leaves the identity bare, and a value + // holding the separator stays distinguishable from two values + "a\\,b,c".to_string(), + "shop/cart,i-1".to_string(), + "we\\\\ird,i-3".to_string(), + ] + ); + } + #[tokio::test] async fn registry_single_column_identity() { let ctx = metric_table_ctx(); @@ -900,24 +1040,24 @@ mod tests { .collect(); rows.sort(); - // Composite id -> sorted `k=v,k=v` plus a JSON object of the id columns; - // descriptive JSON keeps `\`, `"` and control characters intact in - // runtime values, NULL -> "". + // Composite id -> values in declared order plus a JSON object of the + // id columns; descriptive JSON keeps `\`, `"` and control characters + // intact in runtime values, NULL -> "". assert_eq!( rows, vec![ ( - "pid=42,service_name=cart".to_string(), + "cart,42".to_string(), Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()), Some(r#"{"host":""}"#.to_string()), ), ( - "pid=42,service_name=cart".to_string(), + "cart,42".to_string(), Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()), Some(r#"{"host":"h2"}"#.to_string()), ), ( - "pid=42,service_name=cart".to_string(), + "cart,42".to_string(), Some(r#"{"pid":"42","service_name":"cart"}"#.to_string()), Some(r#"{"host":"we\"ird\\\nhost"}"#.to_string()), ), diff --git a/src/operator/src/statement/semantic_graph/conventions.rs b/src/operator/src/statement/semantic_graph/conventions.rs index 36d058921b5..9ec6f52e158 100644 --- a/src/operator/src/statement/semantic_graph/conventions.rs +++ b/src/operator/src/statement/semantic_graph/conventions.rs @@ -19,7 +19,7 @@ //! with the binary, not an operator-editable configuration surface; explicit //! `greptime.semantic.entity.*` declarations always override it. -use std::collections::{BTreeMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::LazyLock; use serde::Deserialize; @@ -43,15 +43,20 @@ pub struct VirtualDstCandidate { pub connection_type: String, } -/// One implicit entity declaration: of a whitelisted Prometheus info metric, -/// or of a trace-v1 table's flattened resource attributes. +/// One implicit entity declaration: of a whitelisted Prometheus or OTel info +/// metric, or of a trace-v1 table's flattened resource attributes. #[derive(Debug, Deserialize)] #[serde(deny_unknown_fields)] pub struct ImplicitEntity { pub entity: String, - /// Identifying label columns; every one must exist on the table for the - /// declaration to apply. + /// Identifying label columns, ordered broad to narrow; every one must + /// exist on the table for the declaration to apply. pub id: Vec, + /// Column qualifying `id[0]` as `/`, so a source + /// carrying the parts separately matches one carrying them pre-composed. + /// Skipped when the column is absent or empty. + #[serde(default)] + pub qualified_by: Option, /// Descriptive label columns, filtered to those present (kube-state-metrics /// label sets vary across versions). #[serde(default)] @@ -70,7 +75,12 @@ pub struct Conventions { pub trace_co_declared_edges: Vec, pub virtual_dst_candidates: Vec, pub otlp_trace_entities: Vec, + /// Table name -> the entities that table declares, for Prometheus-sourced + /// descriptor metrics (`source = prometheus`). pub prometheus_info_metrics: BTreeMap>, + /// Table name -> the entities that table declares, for OTLP-sourced + /// descriptor tables (`source = opentelemetry`). + pub otel_info_metrics: BTreeMap>, } /// The built-in entity-type vocabulary. User-declared types are open-ended; @@ -78,6 +88,7 @@ pub struct Conventions { pub const ENTITY_TYPE_SERVICE: &str = "service"; pub const ENTITY_TYPE_SERVICE_INSTANCE: &str = "service.instance"; pub const ENTITY_TYPE_HOST: &str = "host"; +pub const ENTITY_TYPE_CONTAINER: &str = "container"; pub const ENTITY_TYPE_PROCESS: &str = "process"; pub const ENTITY_TYPE_K8S_POD: &str = "k8s.pod"; pub const ENTITY_TYPE_K8S_NODE: &str = "k8s.node"; @@ -88,10 +99,11 @@ pub const ENTITY_TYPE_GEN_AI_AGENT: &str = "gen_ai.agent"; pub const ENTITY_TYPE_GEN_AI_MODEL: &str = "gen_ai.model"; pub const ENTITY_TYPE_GEN_AI_TOOL: &str = "gen_ai.tool"; -const ENTITY_TYPES: [&str; 12] = [ +const ENTITY_TYPES: [&str; 13] = [ ENTITY_TYPE_SERVICE, ENTITY_TYPE_SERVICE_INSTANCE, ENTITY_TYPE_HOST, + ENTITY_TYPE_CONTAINER, ENTITY_TYPE_PROCESS, ENTITY_TYPE_K8S_POD, ENTITY_TYPE_K8S_NODE, @@ -200,11 +212,16 @@ fn validate(conventions: &Conventions) -> Result<(), String> { let per_table = conventions .prometheus_info_metrics .iter() + .chain(&conventions.otel_info_metrics) .map(|(table, entities)| (table.as_str(), entities)) .chain(std::iter::once(( "otlp traces", &conventions.otlp_trace_entities, ))); + // An entity type declared with a different number of id columns by two + // sources yields ids that can never match, silently splitting one entity + // in two. + let mut arity = HashMap::new(); for (table, entities) in per_table { let mut seen_types = HashSet::new(); for implicit in entities { @@ -226,6 +243,12 @@ fn validate(conventions: &Conventions) -> Result<(), String> { implicit.entity )); } + if implicit.qualified_by.as_ref().is_some_and(String::is_empty) { + return Err(format!( + "entity `{}` of info metric `{table}` has an empty qualified_by", + implicit.entity + )); + } if implicit.descriptive_rest && !implicit.descriptive.is_empty() { return Err(format!( "entity `{}` of info metric `{table}` sets both descriptive and \ @@ -233,6 +256,17 @@ fn validate(conventions: &Conventions) -> Result<(), String> { implicit.entity )); } + match arity.insert(implicit.entity.as_str(), implicit.id.len()) { + Some(previous) if previous != implicit.id.len() => { + return Err(format!( + "entity `{}` is declared with {previous} id columns elsewhere but \ + {} for `{table}`", + implicit.entity, + implicit.id.len() + )); + } + _ => {} + } } } Ok(()) @@ -249,25 +283,28 @@ mod tests { /// Each case must fail on exactly the rule it names, so the shared /// boilerplate is valid and the mutated part is inside the vocabulary. - fn broken(edges: &str, info_metrics: &str) -> String { + fn broken(edges: &str, info_metrics: &str, otel_metrics: &str) -> String { format!( "co_declared_edges: [{edges}]\ntrace_co_declared_edges: []\n\ virtual_dst_candidates: []\notlp_trace_entities: []\n\ - prometheus_info_metrics: {{{info_metrics}}}" + prometheus_info_metrics: {{{info_metrics}}}\n\ + otel_info_metrics: {{{otel_metrics}}}" ) } #[test] fn validation_rejects_broken_conventions() { - let err = |edges, info| parse(&broken(edges, info)).unwrap_err(); + let err = |edges, info, otel| parse(&broken(edges, info, otel)).unwrap_err(); - assert!(err("{src: host, dst: service, rel: pets}", "").contains("unknown rel_type")); + assert!(err("{src: host, dst: service, rel: pets}", "", "").contains("unknown rel_type")); assert!( - err("{src: k8s.pods, dst: k8s.node, rel: runs_on}", "").contains("unknown entity type") + err("{src: k8s.pods, dst: k8s.node, rel: runs_on}", "", "") + .contains("unknown entity type") ); assert!( err( "{src: host, dst: service, rel: uses}, {src: host, dst: service, rel: uses}", + "", "" ) .contains("duplicate edge rule") @@ -275,14 +312,27 @@ mod tests { assert!( err( "", - "t: [{entity: host, id: [x], descriptive: [y], descriptive_rest: true}]" + "t: [{entity: host, id: [x], descriptive: [y], descriptive_rest: true}]", + "" ) .contains("descriptive_rest") ); + // the otel map runs through the same per-table validation + assert!(err("", "", "t: [{entity: hosts, id: [x]}]").contains("unknown entity type")); + // ids of a different arity for one type can never match each other + assert!( + err( + "", + "a: [{entity: host, id: [x]}]", + "b: [{entity: host, id: [x, y]}]" + ) + .contains("id columns") + ); // Unknown YAML keys are rejected, catching typos in the embedded file. assert!( parse(&broken( "{src: host, dst: service, rel: uses, direction: down}", + "", "" )) .is_err() diff --git a/src/operator/src/statement/semantic_graph/conventions.yaml b/src/operator/src/statement/semantic_graph/conventions.yaml index 9fc37fe2a02..808fe0f2658 100644 --- a/src/operator/src/statement/semantic_graph/conventions.yaml +++ b/src/operator/src/statement/semantic_graph/conventions.yaml @@ -2,6 +2,10 @@ # versioned with the binary (include_str!), not an operator-editable surface; # non-standard deployments override per table with explicit # greptime.semantic.entity.* declarations, which always win. +# +# Every `id` list is ordered broad to narrow, and that order is part of the +# identity: the id renders as the values joined in this order, so all +# declarations of one entity type must agree on it or they name two entities. co_declared_edges: - { src: service.instance, dst: host, rel: runs_on } @@ -11,6 +15,8 @@ co_declared_edges: - { src: service.instance, dst: service, rel: part_of } - { src: k8s.pod, dst: k8s.workload, rel: part_of } - { src: service.instance, dst: k8s.pod, rel: runs_on } + - { src: service.instance, dst: container, rel: runs_on } + - { src: container, dst: host, rel: runs_on } # Applied only to trace sources; the endpoints are Greptime entity types # derived from GenAI semantic-convention attributes. @@ -30,12 +36,25 @@ virtual_dst_candidates: - { column: span_attributes.server.address, connection_type: virtual_node } # Implicit declarations for greptime_trace_v1 tables; the pod UID bridges -# trace-side pods onto the kube-state-metrics entities below. +# trace-side pods onto the kube-state-metrics entities below. host.name and +# container.name vary by SDK and resource detector, so identity is the id. +# qualified_by folds the namespace into the service identity the same way the +# metric sources' `job` already carries it, so both signals name one entity. otlp_trace_entities: - entity: service id: [service_name] + qualified_by: resource_attributes.service.namespace - entity: service.instance id: [service_name, resource_attributes.service.instance.id] + qualified_by: resource_attributes.service.namespace + - entity: host + id: [resource_attributes.host.id] + descriptive: + - resource_attributes.host.name + - entity: container + id: [resource_attributes.container.id] + descriptive: + - resource_attributes.container.name - entity: k8s.pod id: [resource_attributes.k8s.pod.uid] descriptive: @@ -113,3 +132,32 @@ prometheus_info_metrics: - entity: service.instance id: [job, instance] descriptive_rest: true + +# Implicit declarations for the OTLP resource descriptor synthesized at +# metrics ingestion, gated on source=opentelemetry and metric.type=info. +# job/instance are the derived Prometheus-compatible identity, so services +# line up with target_info. Unlike target_info the attribute set is known, so +# descriptive columns are listed instead of swept up by descriptive_rest, +# which would misfile other entities' ids under service.instance. +otel_info_metrics: + greptime_otel_resource_info: + - entity: service + id: [job] + descriptive: + - service.name + - service.namespace + - entity: service.instance + id: [job, instance] + - entity: host + id: [host.id] + descriptive: + - host.name + - entity: container + id: [container.id] + descriptive: + - container.name + - entity: k8s.pod + id: [k8s.pod.uid] + descriptive: + - k8s.pod.name + - k8s.namespace.name diff --git a/src/operator/src/statement/semantic_graph/relationships.rs b/src/operator/src/statement/semantic_graph/relationships.rs index 3b5c6d37b78..257c844a9a8 100644 --- a/src/operator/src/statement/semantic_graph/relationships.rs +++ b/src/operator/src/statement/semantic_graph/relationships.rs @@ -200,9 +200,9 @@ fn co_declared_branch( bin.clone() + bin_interval(), bin.clone() + bin_interval(), lit(src.entity_type.as_str()), - entity_id_expr(&src.id_columns, &|c| ident(c)), + entity_id_expr(&src.id_columns, src.id_qualifier.as_deref(), &|c| ident(c)), lit(dst.entity_type.as_str()), - entity_id_expr(&dst.id_columns, &|c| ident(c)), + entity_id_expr(&dst.id_columns, dst.id_qualifier.as_deref(), &|c| ident(c)), lit(*rel_type), lit(*provenance), lit(1.0_f64), @@ -501,6 +501,13 @@ fn span_predicate(service: &EntityDeclaration, window: &GraphQueryWindow, strict } /// One trace table's client spans, normalized to +/// A trace table's `duration_nano` is UInt64 on tables created before the +/// signed-integer ingest change and Int64 after it. The per-table selects are +/// unioned, and those two have no common integer type. +fn duration_nano_expr() -> Expr { + cast(ident(DURATION_NANO_COLUMN), DataType::Int64).alias(DURATION_NANO_COLUMN) +} + /// `(timestamp, trace_id, span_id, src_id, status_code, duration_nano, /// virtual_dst, virtual_conn)`. `src_id` is built from the table's `service` /// declaration, so edges land on exactly the entity ids the registry emits (a @@ -553,9 +560,12 @@ fn client_spans( ident(SPAN_ID_COLUMN), // The cast inside entity_id_expr also normalizes tag columns, which // come out of the storage engine dictionary-encoded. - entity_id_expr(&service.id_columns, &|c| ident(c)).alias(SRC_ID_COLUMN), + entity_id_expr(&service.id_columns, service.id_qualifier.as_deref(), &|c| { + ident(c) + }) + .alias(SRC_ID_COLUMN), ident(SPAN_STATUS_CODE_COLUMN), - ident(DURATION_NANO_COLUMN), + duration_nano_expr(), virtual_dst.alias("virtual_dst"), virtual_conn.alias("virtual_conn"), ]) @@ -578,9 +588,12 @@ fn server_spans( ident(TRACE_TIMESTAMP_COLUMN), ident(TRACE_ID_COLUMN), ident(PARENT_SPAN_ID_COLUMN), - entity_id_expr(&service.id_columns, &|c| ident(c)).alias(DST_ID_COLUMN), + entity_id_expr(&service.id_columns, service.id_qualifier.as_deref(), &|c| { + ident(c) + }) + .alias(DST_ID_COLUMN), ident(SPAN_STATUS_CODE_COLUMN), - ident(DURATION_NANO_COLUMN), + duration_nano_expr(), ]) } @@ -609,7 +622,10 @@ fn agent_calls_branch( ident(TRACE_TIMESTAMP_COLUMN), ident(TRACE_ID_COLUMN), ident(SPAN_ID_COLUMN), - entity_id_expr(&agent.id_columns, &|c| ident(c)).alias(SRC_ID_COLUMN), + entity_id_expr(&agent.id_columns, agent.id_qualifier.as_deref(), &|c| { + ident(c) + }) + .alias(SRC_ID_COLUMN), ])?; let child = trace .scan @@ -619,9 +635,12 @@ fn agent_calls_branch( ident(TRACE_TIMESTAMP_COLUMN), ident(TRACE_ID_COLUMN), ident(PARENT_SPAN_ID_COLUMN), - entity_id_expr(&agent.id_columns, &|c| ident(c)).alias(DST_ID_COLUMN), + entity_id_expr(&agent.id_columns, agent.id_qualifier.as_deref(), &|c| { + ident(c) + }) + .alias(DST_ID_COLUMN), ident(SPAN_STATUS_CODE_COLUMN), - ident(DURATION_NANO_COLUMN), + duration_nano_expr(), ])?; parents = union_all(parents, parent)?; children = union_all(children, child)?; @@ -732,6 +751,16 @@ mod tests { name: &str, extra: &[(&str, &[Option<&str>])], spans: &[Span<'_>], + ) { + register_typed_trace_table(ctx, name, extra, spans, DataType::UInt64) + } + + fn register_typed_trace_table( + ctx: &SessionContext, + name: &str, + extra: &[(&str, &[Option<&str>])], + spans: &[Span<'_>], + duration_type: DataType, ) { let mut fields = vec![ Field::new( @@ -745,7 +774,7 @@ mod tests { Field::new(SPAN_KIND_COLUMN, DataType::Utf8, false), Field::new(SPAN_STATUS_CODE_COLUMN, DataType::Utf8, false), Field::new(SERVICE_NAME_COLUMN, DataType::Utf8, false), - Field::new(DURATION_NANO_COLUMN, DataType::UInt64, false), + Field::new(DURATION_NANO_COLUMN, duration_type.clone(), false), ]; for (column, _) in extra { fields.push(Field::new(*column, DataType::Utf8, true)); @@ -774,9 +803,14 @@ mod tests { Arc::new(StringArray::from( spans.iter().map(|s| s.6).collect::>(), )), - Arc::new(UInt64Array::from( - spans.iter().map(|s| s.7).collect::>(), - )), + match duration_type { + DataType::Int64 => Arc::new(Int64Array::from( + spans.iter().map(|s| s.7 as i64).collect::>(), + )) as ArrayRef, + _ => Arc::new(UInt64Array::from( + spans.iter().map(|s| s.7).collect::>(), + )), + }, ]; for (_, values) in extra { columns.push(Arc::new(StringArray::from(values.to_vec()))); @@ -843,6 +877,7 @@ mod tests { time_index: TRACE_TIMESTAMP_COLUMN.to_string(), entity_type: "service".to_string(), id_columns: id_columns.iter().map(|s| s.to_string()).collect(), + id_qualifier: None, descriptive_columns: vec![], scope_columns: vec![], } @@ -972,7 +1007,7 @@ mod tests { ), ], ); - register_trace_table( + register_typed_trace_table( &ctx, "trace_b", &[], @@ -989,6 +1024,7 @@ mod tests { 1_500_000_000, ), ], + DataType::Int64, ); let a = ctx.table("trace_a").await.unwrap(); let b = ctx.table("trace_b").await.unwrap(); @@ -1294,6 +1330,7 @@ mod tests { time_index: "ts".to_string(), entity_type: entity_type.to_string(), id_columns: id_columns.iter().map(|s| s.to_string()).collect(), + id_qualifier: None, descriptive_columns: vec![], scope_columns: vec![], } @@ -1544,16 +1581,10 @@ mod tests { let total: usize = batches.iter().map(|b| b.num_rows()).sum(); assert_eq!(total, 1); let batch = &batches[0]; - // Composite service identity renders the same sorted `k=v` form the - // registry emits, so edges land on registry entity ids. - assert_eq!( - strings(batch, 5), - vec!["service_name=frontend,service_namespace=ns1"] - ); - assert_eq!( - strings(batch, 7), - vec!["service_name=cart,service_namespace=ns1"] - ); + // Composite service identity renders exactly as the registry emits it, + // so edges land on registry entity ids. + assert_eq!(strings(batch, 5), vec!["frontend,ns1"]); + assert_eq!(strings(batch, 7), vec!["cart,ns1"]); } fn build_relationships_plan_for_test( diff --git a/src/servers/src/http/otlp.rs b/src/servers/src/http/otlp.rs index d04022cdf88..3ecaba74a4e 100644 --- a/src/servers/src/http/otlp.rs +++ b/src/servers/src/http/otlp.rs @@ -118,8 +118,9 @@ pub async fn metrics( promote_scope_attrs: http_opts.promote_scope_attrs, with_metric_engine, experimental_enable_exponential_histogram, - // set is_legacy later + // set by the frontend from its config is_legacy: false, + resource_info: false, metric_type: MetricType::Init, metric_translation_strategy: http_opts.metric_translation_strategy, })); @@ -128,7 +129,7 @@ pub async fn metrics( handler.metrics(request, query_ctx).await.map(|outcome| { if outcome.accepted_data_points == 0 && outcome.rejected_data_points > 0 { OtlpMetricsResponse::Failure(outcome) - } else if outcome.rejected_data_points > 0 { + } else if outcome.rejected_data_points > 0 || outcome.error_message.is_some() { OtlpMetricsResponse::PartialSuccess(outcome) } else { OtlpMetricsResponse::FullSuccess(outcome) diff --git a/src/servers/src/otlp/metrics.rs b/src/servers/src/otlp/metrics.rs index bd0bfb2ce77..d65e117e07c 100644 --- a/src/servers/src/otlp/metrics.rs +++ b/src/servers/src/otlp/metrics.rs @@ -25,6 +25,7 @@ use common_query::native_histogram::{ }; use common_query::prelude::{GREPTIME_COUNT, greptime_timestamp, greptime_value}; use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; +use common_telemetry::warn; use lazy_static::lazy_static; use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, KeyValue, any_value}; @@ -36,7 +37,7 @@ use table::requests::{ }; use crate::error::{self, Result}; -use crate::otlp::trace::{KEY_SERVICE_INSTANCE_ID, KEY_SERVICE_NAME}; +use crate::otlp::trace::{KEY_SERVICE_INSTANCE_ID, KEY_SERVICE_NAME, KEY_SERVICE_NAMESPACE}; use crate::query_handler::MetricsIngestOutcome; use crate::row_writer::{self, MultiTableData, TableData}; pub use crate::semantic::SemanticIndex; @@ -45,8 +46,11 @@ use crate::semantic::{ METRIC_TYPE_UPDOWN_COUNTER, }; +mod resource_info; mod translator; +pub use resource_info::OTEL_RESOURCE_INFO_TABLE_NAME; +use resource_info::ResourceInfoData; pub use translator::legacy_normalize_otlp_name; pub(crate) use translator::ucum_to_openmetrics_unit; use translator::{translate_label_name, translate_metric_name}; @@ -96,28 +100,43 @@ const MIN_EXPONENTIAL_HISTOGRAM_SCALE: i32 = -4; const MAX_EXPONENTIAL_HISTOGRAM_SCALE: i32 = 8; const MAX_REJECTION_MESSAGE_BYTES: usize = 512; +/// Result of converting one OTLP metrics request. +#[derive(Debug)] +pub struct MetricsConversion { + pub requests: RowInsertRequests, + /// Row count of `requests`; the resource descriptor is not counted. + pub rows: usize, + /// Per-table semantic index for the auto-create path to stamp as table + /// options; covers the descriptor table too. + pub semantic_index: SemanticIndex, + /// The synthesized resource descriptor, written separately from the + /// metric data. See [`resource_info`]. + pub resource_info: Option, + pub outcome: MetricsIngestOutcome, +} + /// Convert OpenTelemetry metrics to GreptimeDB insert requests /// /// See /// /// for data structure of OTLP metrics. -/// -/// Returns `InsertRequests`, total number of rows to ingest, and the per-table -/// semantic index for the auto-create path to stamp as table options. pub fn to_grpc_insert_requests( request: ExportMetricsServiceRequest, metric_ctx: &mut OtlpMetricCtx, -) -> Result<( - RowInsertRequests, - usize, - SemanticIndex, - MetricsIngestOutcome, -)> { +) -> Result { let mut table_writer = MultiTableData::default(); let mut semantic_index = SemanticIndex::default(); let mut outcome = MetricsIngestOutcome::default(); + let mut resource_info = ResourceInfoData::default(); for resource in &request.resource_metrics { + if metric_ctx.resource_info + && !metric_ctx.is_legacy + && let Some(r) = resource.resource.as_ref() + { + resource_info.observe(&r.attributes, resource, metric_ctx); + } + let resource_attrs = resource.resource.as_ref().map(|r| { let mut attrs = r.attributes.clone(); process_resource_attrs(&mut attrs, metric_ctx); @@ -149,8 +168,46 @@ pub fn to_grpc_insert_requests( } let (requests, rows) = table_writer.into_row_insert_requests(); + validate_sample_kinds(&requests)?; - Ok((requests, rows, semantic_index, outcome)) + + // The metric and the descriptor would fight over the same table across + // two engines. The user's metric wins. + let resource_info = if !metric_ctx.resource_info { + None + } else if requests + .inserts + .iter() + .any(|r| r.table_name == OTEL_RESOURCE_INFO_TABLE_NAME) + { + warn!( + "Skipping OTLP resource descriptor synthesis: the request writes \ + a metric named `{OTEL_RESOURCE_INFO_TABLE_NAME}`" + ); + None + } else { + resource_info.into_row_insert_requests()? + }; + if resource_info.is_some() { + semantic_index.record_scalar( + OTEL_RESOURCE_INFO_TABLE_NAME, + SEMANTIC_METRIC_TYPE, + crate::semantic::METRIC_TYPE_INFO, + ); + semantic_index.record_scalar( + OTEL_RESOURCE_INFO_TABLE_NAME, + SEMANTIC_METRIC_METADATA_QUALITY, + METADATA_QUALITY_DECLARED, + ); + } + + Ok(MetricsConversion { + requests, + rows, + semantic_index, + outcome, + resource_info, + }) } fn validate_sample_kinds(requests: &RowInsertRequests) -> Result<()> { @@ -287,30 +344,61 @@ fn from_metric_type(data: &metric::Data) -> MetricType { } } +/// Non-scalar values (bool, arrays, maps, bytes) are not representable as tags. +fn scalar_value_string(value: Option<&AnyValue>) -> Option { + match value.and_then(|v| v.value.as_ref())? { + any_value::Value::StringValue(s) => Some(s.clone()), + any_value::Value::IntValue(v) => Some(v.to_string()), + any_value::Value::DoubleValue(v) => Some(v.to_string()), + _ => None, + } +} + +/// Prometheus-style `(job, instance)` identity. Per the OTel Prometheus +/// compatibility spec `job` folds in `service.namespace` when present; a +/// resource without `service.name` gets no job rather than a fabricated one. +/// Both fields are optional strings, so a named type keeps a swapped +/// destructuring from type-checking at the call sites. +pub(crate) struct ServiceIdentity { + pub job: Option, + pub instance: Option, +} + +pub(crate) fn service_identity(attrs: &[KeyValue]) -> ServiceIdentity { + let mut name = None; + let mut namespace = None; + let mut instance = None; + for kv in attrs { + match kv.key.as_str() { + KEY_SERVICE_NAME => name = scalar_value_string(kv.value.as_ref()), + KEY_SERVICE_NAMESPACE => namespace = scalar_value_string(kv.value.as_ref()), + KEY_SERVICE_INSTANCE_ID => instance = scalar_value_string(kv.value.as_ref()), + _ => {} + } + } + let job = name.map(|name| match namespace { + Some(ns) if !ns.is_empty() => format!("{ns}/{name}"), + _ => name, + }); + ServiceIdentity { job, instance } +} + +fn string_key_value(key: &str, value: String) -> KeyValue { + KeyValue { + key: key.to_string(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(value)), + }), + } +} + fn process_resource_attrs(attrs: &mut Vec, metric_ctx: &OtlpMetricCtx) { if metric_ctx.is_legacy { return; } - // remap service.name and service.instance.id to job and instance - let mut tmp = Vec::with_capacity(2); - for kv in attrs.iter() { - match &kv.key as &str { - KEY_SERVICE_NAME => { - tmp.push(KeyValue { - key: JOB_KEY.to_string(), - value: kv.value.clone(), - }); - } - KEY_SERVICE_INSTANCE_ID => { - tmp.push(KeyValue { - key: INSTANCE_KEY.to_string(), - value: kv.value.clone(), - }); - } - _ => {} - } - } + // remap the service identity attributes to job and instance + let ServiceIdentity { job, instance } = service_identity(attrs); // if promote all, then exclude the list, else, include the list if metric_ctx.promote_all_resource_attrs { @@ -322,7 +410,12 @@ fn process_resource_attrs(attrs: &mut Vec, metric_ctx: &OtlpMetricCtx) }); } - attrs.extend(tmp); + if let Some(job) = job { + attrs.push(string_key_value(JOB_KEY, job)); + } + if let Some(instance) = instance { + attrs.push(string_key_value(INSTANCE_KEY, instance)); + } } fn process_scope_attrs(scope: &ScopeMetrics, metric_ctx: &OtlpMetricCtx) -> Option> { @@ -536,35 +629,13 @@ fn encode_exponential_histogram( metric_ctx: &OtlpMetricCtx, outcome: &mut MetricsIngestOutcome, ) -> Result { - if !metric_ctx.experimental_enable_exponential_histogram { + if let Err(rejection) = exponential_histogram_gate(histogram, metric_ctx) { reject_data_points(outcome, histogram.data_points.len(), || { - format!( - "metric `{name}` uses OTLP exponential histograms; set otlp.experimental_enable_exponential_histogram = true to enable ingestion" - ) + rejection.message(name) })?; return Ok(false); } - match AggregationTemporality::try_from(histogram.aggregation_temporality) { - Ok(AggregationTemporality::Cumulative) => {} - Ok(AggregationTemporality::Delta) => { - reject_data_points(outcome, histogram.data_points.len(), || { - format!( - "metric `{name}` uses delta OTLP exponential histograms; only cumulative temporality is supported" - ) - })?; - return Ok(false); - } - _ => { - reject_data_points(outcome, histogram.data_points.len(), || { - format!( - "metric `{name}` has unspecified OTLP exponential histogram temporality; cumulative temporality is required" - ) - })?; - return Ok(false); - } - } - let column_schema = native_histogram_column_schema().map_err(|error| { error::InternalSnafu { err_msg: error.to_string(), @@ -611,7 +682,46 @@ fn encode_exponential_histogram( Ok(emitted) } -fn exponential_histogram_value( +pub(crate) enum ExponentialHistogramRejection { + Disabled, + DeltaTemporality, + UnspecifiedTemporality, +} + +impl ExponentialHistogramRejection { + fn message(&self, name: &str) -> String { + match self { + Self::Disabled => format!( + "metric `{name}` uses OTLP exponential histograms; set otlp.experimental_enable_exponential_histogram = true to enable ingestion" + ), + Self::DeltaTemporality => format!( + "metric `{name}` uses delta OTLP exponential histograms; only cumulative temporality is supported" + ), + Self::UnspecifiedTemporality => format!( + "metric `{name}` has unspecified OTLP exponential histogram temporality; cumulative temporality is required" + ), + } + } +} + +/// Whole-metric acceptance, decided once for the encoder and for the resource +/// descriptor, which must not describe a resource whose data was rejected. +/// Individual points can still fail [`exponential_histogram_value`]. +pub(crate) fn exponential_histogram_gate( + histogram: &ExponentialHistogram, + metric_ctx: &OtlpMetricCtx, +) -> std::result::Result<(), ExponentialHistogramRejection> { + if !metric_ctx.experimental_enable_exponential_histogram { + return Err(ExponentialHistogramRejection::Disabled); + } + match AggregationTemporality::try_from(histogram.aggregation_temporality) { + Ok(AggregationTemporality::Cumulative) => Ok(()), + Ok(AggregationTemporality::Delta) => Err(ExponentialHistogramRejection::DeltaTemporality), + _ => Err(ExponentialHistogramRejection::UnspecifiedTemporality), + } +} + +pub(crate) fn exponential_histogram_value( data_point: &ExponentialHistogramDataPoint, ) -> std::result::Result<(ValueData, i64), String> { if data_point.start_time_unix_nano > data_point.time_unix_nano { @@ -823,29 +933,21 @@ fn write_attributes( }; let tags = attrs.iter().filter_map(|attr| { - attr.value - .as_ref() - .and_then(|v| v.value.as_ref()) - .and_then(|val| { - let key = match attribute_type { - AttributeType::Resource | AttributeType::DataPoint => { - translate_label_name(&attr.key, metric_ctx.metric_translation_strategy) - } - AttributeType::Scope => { - format!( - "otel_scope_{}", - translate_label_name(&attr.key, metric_ctx.metric_translation_strategy) - ) - } - AttributeType::Legacy => legacy_normalize_otlp_name(&attr.key), - }; - match val { - any_value::Value::StringValue(s) => Some((key, s.clone())), - any_value::Value::IntValue(v) => Some((key, v.to_string())), - any_value::Value::DoubleValue(v) => Some((key, v.to_string())), - _ => None, // TODO(sunng87): allow different type of values - } - }) + // TODO(sunng87): allow different type of values + let value = scalar_value_string(attr.value.as_ref())?; + let key = match attribute_type { + AttributeType::Resource | AttributeType::DataPoint => { + translate_label_name(&attr.key, metric_ctx.metric_translation_strategy) + } + AttributeType::Scope => { + format!( + "otel_scope_{}", + translate_label_name(&attr.key, metric_ctx.metric_translation_strategy) + ) + } + AttributeType::Legacy => legacy_normalize_otlp_name(&attr.key), + }; + Some((key, value)) }); row_writer::write_tags(writer, tags, row)?; @@ -1295,6 +1397,169 @@ mod tests { } } + fn descriptor_ctx() -> OtlpMetricCtx { + OtlpMetricCtx { + resource_info: true, + ..Default::default() + } + } + + fn attr_value(attrs: &[KeyValue], key: &str) -> Option { + attrs + .iter() + .find(|kv| kv.key == key) + .and_then(|kv| scalar_value_string(kv.value.as_ref())) + } + + fn gauge_request( + resource_attrs: Vec, + metric_name: &str, + ) -> ExportMetricsServiceRequest { + use otel_arrow_rust::proto::opentelemetry::resource::v1::Resource; + ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + resource: Some(Resource { + attributes: resource_attrs, + ..Default::default() + }), + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + name: metric_name.to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + time_unix_nano: 1_000_000, + value: Some(Value::AsInt(1)), + ..Default::default() + }], + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + } + } + + fn column_names(request: &RowInsertRequests, table: &str) -> Vec { + request + .inserts + .iter() + .find(|r| r.table_name == table) + .unwrap_or_else(|| panic!("missing table {table}")) + .rows + .as_ref() + .unwrap() + .schema + .iter() + .map(|c| c.column_name.clone()) + .collect() + } + + #[test] + fn test_conversion_synthesizes_resource_descriptor() { + set_default_prefix(None).unwrap(); + let request = gauge_request( + vec![keyvalue("service.name", "api"), keyvalue("host.id", "h-1")], + "my_gauge", + ); + let conversion = to_grpc_insert_requests(request, &mut descriptor_ctx()).unwrap(); + + // descriptor keeps raw OTel keys while the metric table's labels went + // through the (default underscore-escaping) translation strategy + let resource_info = conversion.resource_info.expect("descriptor synthesized"); + let descriptor_cols = column_names(&resource_info, OTEL_RESOURCE_INFO_TABLE_NAME); + assert!(descriptor_cols.contains(&"host.id".to_string())); + assert!(descriptor_cols.contains(&"service.name".to_string())); + assert!(descriptor_cols.contains(&"job".to_string())); + let metric_cols = column_names(&conversion.requests, "my_gauge"); + assert!(metric_cols.contains(&"service_name".to_string())); + assert!(!metric_cols.contains(&"service.name".to_string())); + // host.id is not in the promote list: only the descriptor keeps it + assert!(!metric_cols.contains(&"host_id".to_string())); + + let decoded = decode(&conversion.semantic_index); + let t = &decoded[OTEL_RESOURCE_INFO_TABLE_NAME]; + assert_eq!( + t.get(SEMANTIC_METRIC_TYPE).map(String::as_str), + Some("info") + ); + assert_eq!( + t.get(SEMANTIC_METRIC_METADATA_QUALITY).map(String::as_str), + Some("declared") + ); + } + + #[test] + fn test_conversion_skips_descriptor_for_legacy_mode() { + set_default_prefix(None).unwrap(); + let request = gauge_request( + vec![keyvalue("service.name", "api"), keyvalue("host.id", "h-1")], + "my_gauge", + ); + let mut ctx = OtlpMetricCtx { + is_legacy: true, + ..descriptor_ctx() + }; + let conversion = to_grpc_insert_requests(request, &mut ctx).unwrap(); + assert!(conversion.resource_info.is_none()); + + // legacy tables predate job/instance and the promote filter; adding + // either would alter the schema of tables already in use + let cols = column_names(&conversion.requests, "my_gauge"); + assert!(!cols.contains(&"job".to_string())); + assert!(cols.contains(&"host_id".to_string())); + } + + #[test] + fn test_conversion_skips_descriptor_on_metric_name_collision() { + set_default_prefix(None).unwrap(); + let request = gauge_request( + vec![keyvalue("service.name", "api")], + OTEL_RESOURCE_INFO_TABLE_NAME, + ); + let conversion = to_grpc_insert_requests(request, &mut descriptor_ctx()).unwrap(); + assert!(conversion.resource_info.is_none()); + // the metric itself still goes through the main path + assert!( + conversion + .requests + .inserts + .iter() + .any(|r| r.table_name == OTEL_RESOURCE_INFO_TABLE_NAME) + ); + } + + #[test] + fn test_job_composition_follows_service_namespace() { + let mut attrs = vec![ + keyvalue("service.name", "api"), + keyvalue("service.namespace", "shop"), + ]; + process_resource_attrs(&mut attrs, &OtlpMetricCtx::default()); + assert_eq!(attr_value(&attrs, "job").as_deref(), Some("shop/api")); + + let mut attrs = vec![keyvalue("service.name", "api")]; + process_resource_attrs(&mut attrs, &OtlpMetricCtx::default()); + assert_eq!(attr_value(&attrs, "job").as_deref(), Some("api")); + + let mut attrs = vec![ + keyvalue("service.name", "api"), + keyvalue("service.namespace", ""), + ]; + process_resource_attrs(&mut attrs, &OtlpMetricCtx::default()); + assert_eq!(attr_value(&attrs, "job").as_deref(), Some("api")); + + // no service.name: no job is fabricated + let mut attrs = vec![ + keyvalue("service.namespace", "shop"), + keyvalue("service.instance.id", "inst-1"), + ]; + process_resource_attrs(&mut attrs, &OtlpMetricCtx::default()); + assert_eq!(attr_value(&attrs, "job"), None); + assert_eq!(attr_value(&attrs, "instance").as_deref(), Some("inst-1")); + } + #[test] fn test_encode_gauge() { let mut tables = MultiTableData::default(); @@ -2104,8 +2369,12 @@ mod tests { AggregationTemporality::Cumulative, ), ]); - let (requests, _, semantic_index, outcome) = - to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + let MetricsConversion { + requests, + semantic_index, + outcome, + .. + } = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); assert_eq!(outcome.accepted_data_points, 1); assert_eq!(outcome.rejected_data_points, 1); @@ -2127,8 +2396,9 @@ mod tests { vec![], AggregationTemporality::Cumulative, )]); - let (_, _, _, outcome) = - to_grpc_insert_requests(empty, &mut OtlpMetricCtx::default()).unwrap(); + let outcome = to_grpc_insert_requests(empty, &mut OtlpMetricCtx::default()) + .unwrap() + .outcome; assert_eq!(outcome.rejected_data_points, 0); assert_eq!(outcome.error_message, None); } @@ -2208,8 +2478,12 @@ mod tests { .collect(), }; - let (requests, rows, _, outcome) = - to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + let MetricsConversion { + requests, + rows, + outcome, + .. + } = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); assert_eq!(outcome.accepted_data_points, 2); assert_eq!(rows, 6); @@ -2243,8 +2517,13 @@ mod tests { experimental_enable_exponential_histogram: true, ..Default::default() }; - let (requests, rows, semantic_index, outcome) = - to_grpc_insert_requests(request, &mut ctx).unwrap(); + let MetricsConversion { + requests, + rows, + semantic_index, + outcome, + .. + } = to_grpc_insert_requests(request, &mut ctx).unwrap(); assert_eq!(outcome.accepted_data_points, 0); assert_eq!(outcome.rejected_data_points, 1); @@ -2269,14 +2548,17 @@ mod tests { experimental_enable_exponential_histogram: true, ..Default::default() }; - let (new_requests, _, _, _) = - to_grpc_insert_requests(request.clone(), &mut new_ctx).unwrap(); + let new_requests = to_grpc_insert_requests(request.clone(), &mut new_ctx) + .unwrap() + .requests; let mut legacy_ctx = OtlpMetricCtx { experimental_enable_exponential_histogram: true, is_legacy: true, ..Default::default() }; - let (legacy_requests, _, _, _) = to_grpc_insert_requests(request, &mut legacy_ctx).unwrap(); + let legacy_requests = to_grpc_insert_requests(request, &mut legacy_ctx) + .unwrap() + .requests; let new_insert = &new_requests.inserts[0]; let legacy_insert = &legacy_requests.inserts[0]; @@ -2325,8 +2607,9 @@ mod tests { vec![ExponentialHistogramDataPoint::default()], AggregationTemporality::Cumulative, )]); - let (_, _, _, outcome) = - to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()).unwrap(); + let outcome = to_grpc_insert_requests(request, &mut OtlpMetricCtx::default()) + .unwrap() + .outcome; assert_eq!(outcome.rejected_data_points, 1); assert!(outcome.error_message.unwrap().len() <= MAX_REJECTION_MESSAGE_BYTES); diff --git a/src/servers/src/otlp/metrics/resource_info.rs b/src/servers/src/otlp/metrics/resource_info.rs new file mode 100644 index 00000000000..cd3427f7163 --- /dev/null +++ b/src/servers/src/otlp/metrics/resource_info.rs @@ -0,0 +1,393 @@ +// Copyright 2023 Greptime Team +// +// 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. + +//! Resource descriptor synthesized from OTLP metrics requests, so the entity +//! graph reads one table instead of scanning every logical metric table for +//! attributes the promote filter may have dropped. +//! +//! Columns are a fixed allowlist under the raw OTel attribute names: the +//! conventions whitelist matches fixed names, so they must not follow the +//! per-request label translation strategy or the promote/ignore headers. + +use std::collections::BTreeMap; + +use api::v1::RowInsertRequests; +use common_catalog::consts::SEMANTIC_GRAPH_WINDOW_NANOS; +use common_grpc::precision::Precision; +use common_query::prelude::{greptime_timestamp, greptime_value}; +use otel_arrow_rust::proto::opentelemetry::common::v1::KeyValue; +use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ResourceMetrics, metric}; +use session::protocol_ctx::OtlpMetricCtx; + +use crate::error::Result; +use crate::otlp::metrics::{ + INSTANCE_KEY, JOB_KEY, ServiceIdentity, exponential_histogram_gate, + exponential_histogram_value, scalar_value_string, service_identity, +}; +use crate::otlp::trace::{ + KEY_CONTAINER_ID, KEY_CONTAINER_NAME, KEY_HOST_ID, KEY_HOST_NAME, KEY_K8S_NAMESPACE_NAME, + KEY_K8S_POD_NAME, KEY_K8S_POD_UID, KEY_SERVICE_NAME, KEY_SERVICE_NAMESPACE, +}; +use crate::row_writer::{self, MultiTableData}; + +/// Prefixed like the other engine-managed tables, so a user metric is +/// unlikely to claim the name. +pub const OTEL_RESOURCE_INFO_TABLE_NAME: &str = "greptime_otel_resource_info"; + +/// Attributes projected under their raw OTel keys. `service.instance.id` is +/// absent on purpose: it lands in `instance`. +/// +/// Matched instead of scanned: this runs for every attribute of every +/// resource, and the compiler turns it into a length-and-prefix dispatch. +fn is_projected_attr(key: &str) -> bool { + matches!( + key, + KEY_SERVICE_NAME + | KEY_SERVICE_NAMESPACE + | KEY_HOST_ID + | KEY_HOST_NAME + | KEY_CONTAINER_ID + | KEY_CONTAINER_NAME + | KEY_K8S_POD_UID + | KEY_K8S_POD_NAME + | KEY_K8S_NAMESPACE_NAME + ) +} + +/// Upper bound of [`is_projected_attr`] plus the derived `job`/`instance`, +/// used to size the per-row buffers. +const MAX_PROJECTED_TAGS: usize = 11; + +/// Projected attributes (sorted `(name, value)` pairs) -> graph window -> +/// the newest data-point time seen in that window, which is what the row for +/// that window is stamped with. +/// +/// Windows are an inner map so the attributes are stored, and moved, once per +/// resource. They are keyed separately because one request may carry data for +/// several of them: a single row per resource would describe only the newest +/// window, leaving the earlier ones with metric rows but no entities. +#[derive(Debug, Default)] +pub struct ResourceInfoData { + rows: BTreeMap, BTreeMap>, +} + +impl ResourceInfoData { + /// Takes the raw attributes, before the promote filter runs on them. + pub fn observe( + &mut self, + raw_attrs: &[KeyValue], + resource: &ResourceMetrics, + metric_ctx: &OtlpMetricCtx, + ) { + let mut tags = Vec::with_capacity(MAX_PROJECTED_TAGS); + let ServiceIdentity { job, instance } = service_identity(raw_attrs); + if let Some(job) = job { + tags.push((JOB_KEY.to_string(), job)); + } + if let Some(instance) = instance { + tags.push((INSTANCE_KEY.to_string(), instance)); + } + for kv in raw_attrs { + if is_projected_attr(&kv.key) + && let Some(value) = scalar_value_string(kv.value.as_ref()) + { + tags.push((kv.key.clone(), value)); + } + } + if tags.is_empty() { + return; + } + // Sorted so equal attribute sets share a key, and so the emitted + // columns keep a stable order. + tags.sort_unstable(); + + let mut observed: BTreeMap = BTreeMap::new(); + for_each_encoded_time(resource, metric_ctx, |ts| { + let window = ts - ts.rem_euclid(SEMANTIC_GRAPH_WINDOW_NANOS); + observed + .entry(window) + .and_modify(|newest| *newest = (*newest).max(ts)) + .or_insert(ts); + }); + if observed.is_empty() { + return; + } + + let windows = self.rows.entry(tags).or_default(); + for (window, newest) in observed { + windows + .entry(window) + .and_modify(|seen| *seen = (*seen).max(newest)) + .or_insert(newest); + } + } + + /// Every projected attribute becomes a tag, so auto-create puts it in the + /// primary key where the conventions expect it. + pub fn into_row_insert_requests(self) -> Result> { + if self.rows.is_empty() { + return Ok(None); + } + + let mut writer = MultiTableData::default(); + let table = writer.get_or_default_table_data( + OTEL_RESOURCE_INFO_TABLE_NAME, + MAX_PROJECTED_TAGS + 2, + self.rows.values().map(BTreeMap::len).sum(), + ); + for (tags, windows) in &self.rows { + for ts_nanos in windows.values().copied() { + let mut row = table.alloc_one_row(); + row_writer::write_tags(table, tags.iter().cloned(), &mut row)?; + row_writer::write_f64(table, greptime_value(), 1.0, &mut row)?; + row_writer::write_ts_to_millis( + table, + greptime_timestamp(), + Some(ts_nanos), + Precision::Nanosecond, + &mut row, + )?; + table.add_row(row); + } + } + + let (requests, _) = writer.into_row_insert_requests(); + Ok(Some(requests)) + } +} + +/// Visits the times of the data points the encoder writes rows for, so a +/// resource is described exactly where it is measured rather than wherever +/// its request happens to reach. +fn for_each_encoded_time( + resource: &ResourceMetrics, + metric_ctx: &OtlpMetricCtx, + mut visit: impl FnMut(i64), +) { + fn visit_all(points: impl Iterator, visit: &mut impl FnMut(i64)) { + for ts in points { + visit(ts as i64); + } + } + for scope in &resource.scope_metrics { + for m in &scope.metrics { + match &m.data { + Some(metric::Data::Gauge(g)) => { + visit_all(g.data_points.iter().map(|p| p.time_unix_nano), &mut visit) + } + Some(metric::Data::Sum(s)) => { + visit_all(s.data_points.iter().map(|p| p.time_unix_nano), &mut visit) + } + Some(metric::Data::Histogram(h)) => { + visit_all(h.data_points.iter().map(|p| p.time_unix_nano), &mut visit) + } + Some(metric::Data::Summary(s)) => { + visit_all(s.data_points.iter().map(|p| p.time_unix_nano), &mut visit) + } + Some(metric::Data::ExponentialHistogram(h)) + if exponential_histogram_gate(h, metric_ctx).is_ok() => + { + for point in &h.data_points { + if let Ok((_, ts)) = exponential_histogram_value(point) { + visit(ts); + } + } + } + Some(metric::Data::ExponentialHistogram(_)) => {} + None => {} + } + } + } +} + +#[cfg(test)] +mod tests { + use api::v1::SemanticType; + use api::v1::value::ValueData; + use common_query::prelude::set_default_prefix; + use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, any_value}; + use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ + AggregationTemporality, ExponentialHistogram, ExponentialHistogramDataPoint, Gauge, Metric, + NumberDataPoint, ScopeMetrics, + }; + + use super::*; + + fn kv(key: &str, value: &str) -> KeyValue { + KeyValue { + key: key.into(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(value.into())), + }), + } + } + + fn gauge_at(times: &[i64]) -> ResourceMetrics { + ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + data: Some(metric::Data::Gauge(Gauge { + data_points: times + .iter() + .map(|ts| NumberDataPoint { + time_unix_nano: *ts as u64, + ..Default::default() + }) + .collect(), + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + } + } + + #[test] + fn observe_projects_allowlist_and_dedups_per_request() { + let mut data = ResourceInfoData::default(); + let attrs = vec![ + kv("service.name", "api"), + kv("service.namespace", "shop"), + kv("service.instance.id", "inst-1"), + kv("host.id", "h-1"), + kv("os.type", "linux"), + ]; + data.observe(&attrs, &gauge_at(&[100, 50]), &OtlpMetricCtx::default()); + assert_eq!(data.rows.len(), 1); + let (tags, windows) = data.rows.iter().next().unwrap(); + assert_eq!(windows.values().copied().collect::>(), vec![100]); + assert!(tags.contains(&("job".to_string(), "shop/api".to_string()))); + assert!(tags.contains(&("instance".to_string(), "inst-1".to_string()))); + assert!(tags.contains(&("service.name".to_string(), "api".to_string()))); + assert!( + tags.iter() + .all(|(k, _)| k != "os.type" && k != "service.instance.id") + ); + + data.observe( + &[kv("host.id", "h-2")], + &gauge_at(&[10]), + &OtlpMetricCtx::default(), + ); + assert_eq!(data.rows.len(), 2); + + let mut empty = ResourceInfoData::default(); + empty.observe( + &[kv("os.type", "linux")], + &gauge_at(&[100]), + &OtlpMetricCtx::default(), + ); + assert!(empty.into_row_insert_requests().unwrap().is_none()); + } + + /// Earlier windows would keep their metric rows but lose their entities. + #[test] + fn observe_keeps_one_row_per_graph_window() { + let window = SEMANTIC_GRAPH_WINDOW_NANOS; + let mut data = ResourceInfoData::default(); + data.observe( + &[kv("service.name", "api")], + &gauge_at(&[window + 1, window + 2, 3 * window + 7]), + &OtlpMetricCtx::default(), + ); + + let windows = data.rows.values().next().unwrap(); + assert_eq!( + windows.iter().collect::>(), + vec![(&window, &(window + 2)), (&(3 * window), &(3 * window + 7))] + ); + } + + /// Describing a resource whose only data the encoder drops invents an + /// entity with no measurements. + #[test] + fn observe_ignores_data_the_encoder_drops() { + let exponential = |temporality: AggregationTemporality| ResourceMetrics { + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + data: Some(metric::Data::ExponentialHistogram(ExponentialHistogram { + data_points: vec![ExponentialHistogramDataPoint { + time_unix_nano: 100, + ..Default::default() + }], + aggregation_temporality: temporality as i32, + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }; + let enabled = OtlpMetricCtx { + experimental_enable_exponential_histogram: true, + ..Default::default() + }; + + for (resource, ctx) in [ + ( + exponential(AggregationTemporality::Cumulative), + OtlpMetricCtx::default(), + ), + (exponential(AggregationTemporality::Delta), enabled), + ] { + let mut data = ResourceInfoData::default(); + data.observe(&[kv("service.name", "api")], &resource, &ctx); + assert!(data.into_row_insert_requests().unwrap().is_none()); + } + } + + #[test] + fn rows_carry_raw_key_tags_value_and_millis_timestamp() { + set_default_prefix(None).unwrap(); + let mut data = ResourceInfoData::default(); + data.observe( + &[kv("service.name", "api"), kv("host.id", "h-1")], + &gauge_at(&[1_700_000_000_123_456_789]), + &OtlpMetricCtx::default(), + ); + let requests = data.into_row_insert_requests().unwrap().unwrap(); + assert_eq!(requests.inserts.len(), 1); + let insert = &requests.inserts[0]; + assert_eq!(insert.table_name, OTEL_RESOURCE_INFO_TABLE_NAME); + + let rows = insert.rows.as_ref().unwrap(); + let names = rows + .schema + .iter() + .map(|c| c.column_name.as_str()) + .collect::>(); + assert_eq!( + names, + vec![ + "host.id", + "job", + "service.name", + greptime_value(), + greptime_timestamp() + ] + ); + for column in &rows.schema[..3] { + assert_eq!(column.semantic_type, SemanticType::Tag as i32); + } + + assert_eq!(rows.rows.len(), 1); + let values = &rows.rows[0].values; + assert_eq!(values[3].value_data, Some(ValueData::F64Value(1.0))); + assert_eq!( + values[4].value_data, + Some(ValueData::TimestampMillisecondValue(1_700_000_000_123)) + ); + } +} diff --git a/src/servers/src/otlp/trace.rs b/src/servers/src/otlp/trace.rs index a8439cdb2c4..9ee16507bda 100644 --- a/src/servers/src/otlp/trace.rs +++ b/src/servers/src/otlp/trace.rs @@ -46,7 +46,15 @@ pub const TRACE_STATE_COLUMN: &str = "trace_state"; // const keys pub const KEY_SERVICE_NAME: &str = "service.name"; +pub const KEY_SERVICE_NAMESPACE: &str = "service.namespace"; pub const KEY_SERVICE_INSTANCE_ID: &str = "service.instance.id"; +pub const KEY_HOST_ID: &str = "host.id"; +pub const KEY_HOST_NAME: &str = "host.name"; +pub const KEY_CONTAINER_ID: &str = "container.id"; +pub const KEY_CONTAINER_NAME: &str = "container.name"; +pub const KEY_K8S_POD_UID: &str = "k8s.pod.uid"; +pub const KEY_K8S_POD_NAME: &str = "k8s.pod.name"; +pub const KEY_K8S_NAMESPACE_NAME: &str = "k8s.namespace.name"; pub const KEY_SPAN_KIND: &str = "span.kind"; // jaeger const keys, not sure if they are general diff --git a/src/session/src/protocol_ctx.rs b/src/session/src/protocol_ctx.rs index 537dec5c414..d9a4a386f91 100644 --- a/src/session/src/protocol_ctx.rs +++ b/src/session/src/protocol_ctx.rs @@ -56,6 +56,9 @@ pub struct OtlpMetricCtx { pub with_metric_engine: bool, pub experimental_enable_exponential_histogram: bool, pub is_legacy: bool, + /// Set from the server's `otlp.experimental_enable_resource_info`; off + /// means the resource descriptor is not synthesized at all. + pub resource_info: bool, pub metric_type: MetricType, pub metric_translation_strategy: OtlpMetricTranslationStrategy, } diff --git a/src/table/src/requests/semantic.rs b/src/table/src/requests/semantic.rs index 2daa440c11d..471be8b891c 100644 --- a/src/table/src/requests/semantic.rs +++ b/src/table/src/requests/semantic.rs @@ -98,6 +98,10 @@ pub const SEMANTIC_ENTITY_SERVICE_ID: &str = "greptime.semantic.entity.service.i /// ([`has_stable_string_form`]). #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum EntityRole { + /// Identifying columns. Their **order is part of the identity**: the + /// entity id is their values joined in the declared order, so two tables + /// naming the same entity must list them the same way (broad to narrow) + /// or they name two entities. Id, Descriptive, Scope, diff --git a/tests-integration/src/instance.rs b/tests-integration/src/instance.rs index 08d2a1aff15..9f48d01a467 100644 --- a/tests-integration/src/instance.rs +++ b/tests-integration/src/instance.rs @@ -691,4 +691,119 @@ mod tests { "denied source leaked into entities:\n{pretty_print}" ); } + + /// Denying the descriptor's table must not fail the request: the metric + /// data is already committed, so the client would retry what the server + /// accepted. + #[tokio::test(flavor = "multi_thread")] + async fn test_otlp_descriptor_permission_denial_degrades() { + use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest; + use otel_arrow_rust::proto::opentelemetry::common::v1::{AnyValue, KeyValue, any_value}; + use otel_arrow_rust::proto::opentelemetry::metrics::v1::number_data_point::Value; + use otel_arrow_rust::proto::opentelemetry::metrics::v1::{ + Gauge, Metric, NumberDataPoint, ResourceMetrics, ScopeMetrics, metric, + }; + use otel_arrow_rust::proto::opentelemetry::resource::v1::Resource; + use servers::query_handler::OpenTelemetryProtocolHandler; + + common_telemetry::init_default_ut_logging(); + + /// Denies writes to the descriptor table only. + struct DenyResourceInfo; + + impl auth::PermissionChecker for DenyResourceInfo { + fn check_permission( + &self, + _user_info: auth::UserInfoRef, + _req: auth::PermissionReq, + ) -> auth::error::Result { + Ok(auth::PermissionResp::Allow) + } + + fn check_permission_with_table_targets( + &self, + _user_info: auth::UserInfoRef, + _req: auth::PermissionReq, + targets: auth::PermissionTableTargets, + ) -> auth::error::Result { + if let auth::PermissionTableTargets::Resolved(targets) = targets + && targets + .iter() + .any(|target| target.table == "greptime_otel_resource_info") + { + return Ok(auth::PermissionResp::Reject); + } + Ok(auth::PermissionResp::Allow) + } + } + + let plugins = Plugins::new(); + plugins.insert::(Arc::new(DenyResourceInfo)); + + let standalone = GreptimeDbStandaloneBuilder::new("test_otlp_descriptor_denied") + .with_plugin(plugins) + .build() + .await; + let instance = standalone.fe_instance().clone(); + + let attr = |key: &str, value: &str| KeyValue { + key: key.to_string(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(value.to_string())), + }), + }; + let request = ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + resource: Some(Resource { + attributes: vec![attr("service.name", "api"), attr("host.id", "h-1")], + ..Default::default() + }), + scope_metrics: vec![ScopeMetrics { + metrics: vec![Metric { + name: "denied_descriptor_gauge".to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + time_unix_nano: 1_700_000_000_000_000_000, + value: Some(Value::AsDouble(1.0)), + ..Default::default() + }], + })), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }], + }; + + let outcome = + OpenTelemetryProtocolHandler::metrics(&*instance, request, QueryContext::arc()) + .await + .unwrap(); + let warning = outcome + .error_message + .expect("denial must surface as a warning"); + assert!( + warning.contains("greptime_otel_resource_info"), + "unexpected warning: {warning}" + ); + + // the metric data is committed, the descriptor table was never created + let sql = "select table_name from information_schema.tables \ + where table_schema = 'public' order by table_name"; + let output = query(&instance, sql).await; + let OutputData::Stream(s) = output.data else { + unreachable!() + }; + let batches = common_recordbatch::util::collect_batches(s).await.unwrap(); + let pretty_print = batches.pretty_print().unwrap(); + assert!( + pretty_print.contains("denied_descriptor_gauge"), + "metric data was not committed:\n{pretty_print}" + ); + assert!( + !pretty_print.contains("greptime_otel_resource_info"), + "denied descriptor table was created:\n{pretty_print}" + ); + } } diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 1085c21983e..16fc2e65634 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -369,6 +369,11 @@ impl GreptimeDbStandaloneBuilder { grpc: GrpcOptions::default().with_server_addr("127.0.0.1:4001"), slow_query: self.slow_query_options.clone(), auto_create_table: self.auto_create_table, + // Tests cover the descriptor, so they run with it enabled. + otlp: frontend::service_config::OtlpOptions { + experimental_enable_resource_info: true, + ..Default::default() + }, ..StandaloneOptions::default() }; diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 7a787833376..4a8f96260e8 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -164,6 +164,7 @@ macro_rules! http_tests { test_otlp_metrics_new, test_otlp_exponential_histogram, test_otlp_metric_translation_strategies, + test_otlp_metrics_resource_info_conflicts, test_otlp_traces_v0, test_otlp_traces_v1, test_otlp_traces_v1_entity_graph, @@ -2208,6 +2209,7 @@ enable = true enable = true experimental_enable_exponential_histogram = false trace_ingest_chunk_size = 512 +experimental_enable_resource_info = true [prom_store] enable = true @@ -6500,7 +6502,7 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { .await; assert_eq!(StatusCode::OK, res.status()); - let expected = "[[\"claude_code_cost_usage_USD_total\"],[\"claude_code_token_usage_tokens_total\"],[\"demo\"],[\"greptime_physical_table\"],[\"numbers\"]]"; + let expected = "[[\"claude_code_cost_usage_USD_total\"],[\"claude_code_token_usage_tokens_total\"],[\"demo\"],[\"greptime_otel_resource_info\"],[\"greptime_physical_table\"],[\"numbers\"]]"; validate_data("otlp_metrics_all_tables", &client, "show tables;", expected).await; // Metric-engine logical table carries the semantic identity. Match substrings @@ -6567,6 +6569,33 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { ) .await; + // The synthesized resource descriptor: a plain mito info table keyed by + // the raw OTel attribute keys (service.name is the only allowlisted + // resource attribute in this payload), independent of the promote/ignore + // headers. The timestamp is the newest data-point time. + validate_data( + "otlp_metrics_resource_info", + &client, + "select job, \"service.name\", greptime_value, greptime_timestamp from greptime_otel_resource_info;", + "[[\"claude-code\",\"claude-code\",1.0,1753780559836]]", + ) + .await; + // All four ownership markers must land on auto-create: the descriptor + // write path refuses tables that miss any of them, so a missing stamp + // here would silently stop every follow-up descriptor write. + validate_data( + "otlp_metrics_resource_info_options", + &client, + "select count(*) from information_schema.tables where table_name = 'greptime_otel_resource_info' \ + and engine = 'mito' \ + and create_options like '%greptime.semantic.metric.type=info%' \ + and create_options like '%greptime.semantic.metric.metadata_quality=declared%' \ + and create_options like '%greptime.semantic.signal_type=metric%' \ + and create_options like '%greptime.semantic.source=opentelemetry%';", + "[[1]]", + ) + .await; + // drop table let res = client .get("/v1/sql?sql=drop table `claude_code_cost_usage_USD_total`;") @@ -6578,6 +6607,11 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { .send() .await; assert_eq!(res.status(), StatusCode::OK); + let res = client + .get("/v1/sql?sql=drop table greptime_otel_resource_info;") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); // write metrics data // with scope attrs @@ -6704,6 +6738,16 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { ) .await; + // the descriptor was auto-created again after the earlier drop + validate_data( + "otlp_metrics_resource_info_recreated", + &client, + "select count(*) from information_schema.tables \ + where table_name = 'greptime_otel_resource_info' and engine = 'mito';", + "[[1]]", + ) + .await; + // drop table let res = client .get("/v1/sql?sql=drop table `claude_code_cost_usage_USD_total`;") @@ -6715,6 +6759,44 @@ pub async fn test_otlp_metrics_new(store_type: StorageType) { .send() .await; assert_eq!(res.status(), StatusCode::OK); + let res = client + .get("/v1/sql?sql=drop table greptime_otel_resource_info;") + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + + // job composition end-to-end: service.namespace folds into "/" + // on both the descriptor and the metric tables + let content = r#" +{"resourceMetrics":[{"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"api"}},{"key":"service.namespace","value":{"stringValue":"shop"}},{"key":"host.id","value":{"stringValue":"h-1"}}]},"scopeMetrics":[{"scope":{"name":"s"},"metrics":[{"name":"ns_gauge","gauge":{"dataPoints":[{"timeUnixNano":"1753780559836000000","asDouble":1.0}]}}]}]}]} + "#; + let req: ExportMetricsServiceRequest = serde_json::from_str(content).unwrap(); + let res = send_req( + &client, + vec![( + HeaderName::from_static("content-type"), + HeaderValue::from_static("application/x-protobuf"), + )], + "/v1/otlp/v1/metrics", + req.encode_to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + validate_data( + "otlp_metrics_namespace_job_descriptor", + &client, + "select job, \"service.name\", \"service.namespace\", \"host.id\" from greptime_otel_resource_info;", + "[[\"shop/api\",\"api\",\"shop\",\"h-1\"]]", + ) + .await; + validate_data( + "otlp_metrics_namespace_job_metric_table", + &client, + "select job from ns_gauge;", + "[[\"shop/api\"]]", + ) + .await; guard.remove_all().await; } @@ -6912,6 +6994,24 @@ pub async fn test_otlp_metric_translation_strategies(store_type: StorageType) { ) .await; + // the descriptor's columns are the fixed raw keys regardless of the + // translation strategy, and only allowlisted attributes are projected + validate_data( + "otlp_metric_translation_strategy_resource_info", + &client, + "select job, \"service.name\" from greptime_otel_resource_info;", + "[[\"strategy-service\",\"strategy-service\"]]", + ) + .await; + validate_data( + "otlp_metric_translation_strategy_resource_info_allowlist", + &client, + "select count(*) from information_schema.columns \ + where table_name = 'greptime_otel_resource_info' and column_name = 'resource.attr';", + "[[0]]", + ) + .await; + let res = send_req( &client, vec![ @@ -6934,6 +7034,122 @@ pub async fn test_otlp_metric_translation_strategies(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_otlp_metrics_resource_info_conflicts(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = + setup_test_http_app_with_frontend(store_type, "test_otlp_metrics_resource_info_conflicts") + .await; + let client = TestClient::new(app).await; + + let content_type = || { + ( + HeaderName::from_static("content-type"), + HeaderValue::from_static("application/x-protobuf"), + ) + }; + + // A metric named like the descriptor wins the table name: it goes through + // the normal metric-engine path and no descriptor is synthesized. + let content = r#" +{"resourceMetrics":[{"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"api"}},{"key":"host.id","value":{"stringValue":"h-1"}}]},"scopeMetrics":[{"scope":{"name":"s"},"metrics":[{"name":"greptime_otel_resource_info","gauge":{"dataPoints":[{"timeUnixNano":"1753780559836000000","asDouble":1.0}]}}]}]}]} + "#; + let req: ExportMetricsServiceRequest = serde_json::from_str(content).unwrap(); + let res = send_req( + &client, + vec![content_type()], + "/v1/otlp/v1/metrics", + req.encode_to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + let body = ExportMetricsServiceResponse::decode(res.bytes().await).unwrap(); + assert!(body.partial_success.is_none()); + validate_data( + "otlp_metrics_name_collision_engine", + &client, + "select count(*) from information_schema.tables \ + where table_name = 'greptime_otel_resource_info' and engine = 'metric';", + "[[1]]", + ) + .await; + let res = execute_sql(&client, "drop table greptime_otel_resource_info;").await; + assert_eq!(res.status(), StatusCode::OK); + + // A pre-existing table with an incompatible schema fails the descriptor + // write, which degrades to an OTLP partial-success warning while the + // metric data itself is accepted. + let res = execute_sql( + &client, + "create table greptime_otel_resource_info (ts timestamp time index, job double, greptime_value double);", + ) + .await; + assert_eq!(res.status(), StatusCode::OK); + + let content = r#" +{"resourceMetrics":[{"resource":{"attributes":[{"key":"service.name","value":{"stringValue":"api"}}]},"scopeMetrics":[{"scope":{"name":"s"},"metrics":[{"name":"conflict_gauge","gauge":{"dataPoints":[{"timeUnixNano":"1753780559836000000","asDouble":7.0}]}}]}]}]} + "#; + let req: ExportMetricsServiceRequest = serde_json::from_str(content).unwrap(); + let res = send_req( + &client, + vec![content_type()], + "/v1/otlp/v1/metrics", + req.encode_to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + let body = ExportMetricsServiceResponse::decode(res.bytes().await).unwrap(); + let partial_success = body.partial_success.as_ref().unwrap(); + assert_eq!(partial_success.rejected_data_points, 0); + assert!( + partial_success + .error_message + .contains("greptime_otel_resource_info"), + "unexpected warning: {}", + partial_success.error_message + ); + validate_data( + "otlp_metrics_conflict_main_data_accepted", + &client, + "select greptime_value from conflict_gauge;", + "[[7.0]]", + ) + .await; + + // The descriptor is auto-created again once the conflicting table is + // gone, and a follow-up write into it does not degrade. + let res = execute_sql(&client, "drop table greptime_otel_resource_info;").await; + assert_eq!(res.status(), StatusCode::OK); + for round in 0..2 { + let req: ExportMetricsServiceRequest = serde_json::from_str(content).unwrap(); + let res = send_req( + &client, + vec![content_type()], + "/v1/otlp/v1/metrics", + req.encode_to_vec(), + false, + ) + .await; + assert_eq!(StatusCode::OK, res.status()); + let body = ExportMetricsServiceResponse::decode(res.bytes().await).unwrap(); + assert!( + body.partial_success.is_none(), + "round {round} unexpectedly degraded: {:?}", + body.partial_success + ); + } + validate_data( + "otlp_metrics_owned_descriptor_keeps_accepting", + &client, + "select count(*) from greptime_otel_resource_info;", + "[[1]]", + ) + .await; + + guard.remove_all().await; +} + pub async fn test_otlp_traces_v0(store_type: StorageType) { // init common_telemetry::init_default_ut_logging(); @@ -7136,7 +7352,7 @@ pub async fn test_otlp_traces_v1_entity_graph(store_type: StorageType) { let res = send_trace_v1_req(&client, "graph_traces", req, false).await; assert_eq!(StatusCode::OK, res.status()); - let expected = r#"[["service","frontend","service","cart","calls"],["service.instance","resource_attributes.service.instance.id=binst-1,service_name=batch","service","batch","part_of"],["service.instance","resource_attributes.service.instance.id=inst-1,service_name=frontend","service","frontend","part_of"],["service.instance","resource_attributes.service.instance.id=inst-1,service_name=frontend","k8s.pod","uid-9","runs_on"]]"#; + let expected = r#"[["service","frontend","service","cart","calls"],["service.instance","batch,binst-1","service","batch","part_of"],["service.instance","frontend,inst-1","service","frontend","part_of"],["service.instance","frontend,inst-1","k8s.pod","uid-9","runs_on"]]"#; validate_data( "otlp_traces_entity_graph", &client, diff --git a/tests/cases/standalone/common/system/semantic_graph.result b/tests/cases/standalone/common/system/semantic_graph.result index 6af28cfa2b5..f8894319842 100644 --- a/tests/cases/standalone/common/system/semantic_graph.result +++ b/tests/cases/standalone/common/system/semantic_graph.result @@ -94,26 +94,26 @@ select entity_type, entity_id, entity_id_attrs, scope, descriptive, source_table from greptime_private.semantic_entities order by entity_type, entity_id; -+------------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ -| entity_type | entity_id | entity_id_attrs | scope | descriptive | source_tables | -+------------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ -| host | h1 | | | | ["public.graph_app_metrics"] | -| process | host=h1,service_name=cart | {"host":"h1","service_name":"cart"} | | {"env":"us-east"} | ["public.graph_app_metrics"] | -| service | cart | | us-east | | ["public.graph_app_metrics"] | -| service.instance | cart-0 | | | | ["public.graph_app_metrics"] | -+------------------+---------------------------+-------------------------------------+---------+-------------------+------------------------------+ ++------------------+-----------+-------------------------------------+---------+-------------------+------------------------------+ +| entity_type | entity_id | entity_id_attrs | scope | descriptive | source_tables | ++------------------+-----------+-------------------------------------+---------+-------------------+------------------------------+ +| host | h1 | | | | ["public.graph_app_metrics"] | +| process | cart,h1 | {"host":"h1","service_name":"cart"} | | {"env":"us-east"} | ["public.graph_app_metrics"] | +| service | cart | | us-east | | ["public.graph_app_metrics"] | +| service.instance | cart-0 | | | | ["public.graph_app_metrics"] | ++------------------+-----------+-------------------------------------+---------+-------------------+------------------------------+ select src_type, src_id, dst_type, dst_id, rel_type, provenance, confidence from greptime_private.semantic_relationships order by rel_type, src_id; -+------------------+---------------------------+----------+--------+----------+------------+------------+ -| src_type | src_id | dst_type | dst_id | rel_type | provenance | confidence | -+------------------+---------------------------+----------+--------+----------+------------+------------+ -| service.instance | cart-0 | service | cart | part_of | attribute | 1.0 | -| service.instance | cart-0 | host | h1 | runs_on | attribute | 1.0 | -| process | host=h1,service_name=cart | host | h1 | runs_on | attribute | 1.0 | -+------------------+---------------------------+----------+--------+----------+------------+------------+ ++------------------+---------+----------+--------+----------+------------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | confidence | ++------------------+---------+----------+--------+----------+------------+------------+ +| service.instance | cart-0 | service | cart | part_of | attribute | 1.0 | +| process | cart,h1 | host | h1 | runs_on | attribute | 1.0 | +| service.instance | cart-0 | host | h1 | runs_on | attribute | 1.0 | ++------------------+---------+----------+--------+----------+------------+------------+ drop table graph_app_metrics; @@ -570,38 +570,38 @@ select entity_type, entity_id, source_tables from greptime_private.semantic_entities order by entity_type, entity_id, source_tables; -+------------------+-----------------------------------------------------------+------------------------------------+ -| entity_type | entity_id | source_tables | -+------------------+-----------------------------------------------------------+------------------------------------+ -| k8s.container | container=main,uid=uid-1 | ["public.kube_pod_container_info"] | -| k8s.node | node-a | ["public.kube_pod_info"] | -| k8s.pod | uid-1 | ["public.kube_pod_info"] | -| k8s.pod | uid-1 | ["public.kube_pod_owner"] | -| k8s.pod | uid-1 | ["public.kube_pod_container_info"] | -| k8s.pod | uid-2 | ["public.kube_pod_info"] | -| k8s.pod | uid-2 | ["public.kube_pod_owner"] | -| k8s.pod | uid-3 | ["public.kube_pod_info"] | -| k8s.service | svc-uid-1 | ["public.kube_service_info"] | -| k8s.workload | namespace=default,owner_kind=ReplicaSet,owner_name=api-rs | ["public.kube_pod_owner"] | -| service | shop/api | ["public.target_info"] | -| service.instance | instance=inst-1,job=shop/api | ["public.target_info"] | -+------------------+-----------------------------------------------------------+------------------------------------+ ++------------------+---------------------------+------------------------------------+ +| entity_type | entity_id | source_tables | ++------------------+---------------------------+------------------------------------+ +| k8s.container | uid-1,main | ["public.kube_pod_container_info"] | +| k8s.node | node-a | ["public.kube_pod_info"] | +| k8s.pod | uid-1 | ["public.kube_pod_info"] | +| k8s.pod | uid-1 | ["public.kube_pod_owner"] | +| k8s.pod | uid-1 | ["public.kube_pod_container_info"] | +| k8s.pod | uid-2 | ["public.kube_pod_info"] | +| k8s.pod | uid-2 | ["public.kube_pod_owner"] | +| k8s.pod | uid-3 | ["public.kube_pod_info"] | +| k8s.service | svc-uid-1 | ["public.kube_service_info"] | +| k8s.workload | default,ReplicaSet,api-rs | ["public.kube_pod_owner"] | +| service | shop/api | ["public.target_info"] | +| service.instance | shop/api,inst-1 | ["public.target_info"] | ++------------------+---------------------------+------------------------------------+ -- SQLNESS PROTOCOL MYSQL select src_type, src_id, dst_type, dst_id, rel_type, provenance from greptime_private.semantic_relationships order by rel_type, src_id, dst_id; -+------------------+------------------------------+---------------+-----------------------------------------------------------+----------+------------+ -| src_type | src_id | dst_type | dst_id | rel_type | provenance | -+------------------+------------------------------+---------------+-----------------------------------------------------------+----------+------------+ -| k8s.pod | uid-1 | k8s.container | container=main,uid=uid-1 | contains | attribute | -| service.instance | instance=inst-1,job=shop/api | service | shop/api | part_of | attribute | -| k8s.pod | uid-1 | k8s.workload | namespace=default,owner_kind=ReplicaSet,owner_name=api-rs | part_of | attribute | -| k8s.pod | uid-2 | k8s.workload | namespace=default,owner_kind=ReplicaSet,owner_name=api-rs | part_of | attribute | -| k8s.pod | uid-1 | k8s.node | node-a | runs_on | attribute | -| k8s.pod | uid-2 | k8s.node | node-a | runs_on | attribute | -+------------------+------------------------------+---------------+-----------------------------------------------------------+----------+------------+ ++------------------+-----------------+---------------+---------------------------+----------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | ++------------------+-----------------+---------------+---------------------------+----------+------------+ +| k8s.pod | uid-1 | k8s.container | uid-1,main | contains | attribute | +| service.instance | shop/api,inst-1 | service | shop/api | part_of | attribute | +| k8s.pod | uid-1 | k8s.workload | default,ReplicaSet,api-rs | part_of | attribute | +| k8s.pod | uid-2 | k8s.workload | default,ReplicaSet,api-rs | part_of | attribute | +| k8s.pod | uid-1 | k8s.node | node-a | runs_on | attribute | +| k8s.pod | uid-2 | k8s.node | node-a | runs_on | attribute | ++------------------+-----------------+---------------+---------------------------+----------+------------+ drop table kube_pod_info; @@ -627,3 +627,182 @@ drop table http_requests_total; Affected Rows: 0 +-- OTel conventions: the ingestion-synthesized greptime_otel_resource_info descriptor +-- (stamped signal_type=metric + source=opentelemetry + metric.type=info) gets implicit +-- declarations under the raw OTel column names. host/container identities are +-- the stable ids only, so rows with an empty host.id / container.id link no +-- infrastructure, and a row without an instance value declares no +-- service.instance. The same table name stamped source=prometheus is not +-- whitelisted. +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "service.name" string, + "service.namespace" string, + "host.id" string, + "host.name" string, + "container.id" string, + "container.name" string, + greptime_value double, + primary key (job, instance, "service.name", "service.namespace", + "host.id", "host.name", "container.id", "container.name") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'opentelemetry', + 'greptime.semantic.metric.type' = 'info' +); + +Affected Rows: 0 + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'api', 'shop', 'h-1', 'node-a', 'c-1', 'api-ctr', 1), + (now(), 'shop/api', 'inst-2', 'api', 'shop', '', 'laptop', '', '', 1), + (now(), 'worker', '', 'worker', '', 'h-1', 'node-a', '', '', 1); + +Affected Rows: 3 + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +order by entity_type, entity_id, source_tables; + ++------------------+-----------------+----------------------------------------+ +| entity_type | entity_id | source_tables | ++------------------+-----------------+----------------------------------------+ +| container | c-1 | ["public.greptime_otel_resource_info"] | +| host | h-1 | ["public.greptime_otel_resource_info"] | +| service | shop/api | ["public.greptime_otel_resource_info"] | +| service | worker | ["public.greptime_otel_resource_info"] | +| service.instance | shop/api,inst-1 | ["public.greptime_otel_resource_info"] | +| service.instance | shop/api,inst-2 | ["public.greptime_otel_resource_info"] | ++------------------+-----------------+----------------------------------------+ + +-- SQLNESS PROTOCOL MYSQL +select src_type, src_id, dst_type, dst_id, rel_type, provenance +from greptime_private.semantic_relationships +order by rel_type, src_id, dst_id; + ++------------------+-----------------+-----------+----------+----------+------------+ +| src_type | src_id | dst_type | dst_id | rel_type | provenance | ++------------------+-----------------+-----------+----------+----------+------------+ +| service.instance | shop/api,inst-1 | service | shop/api | part_of | attribute | +| service.instance | shop/api,inst-2 | service | shop/api | part_of | attribute | +| container | c-1 | host | h-1 | runs_on | attribute | +| service.instance | shop/api,inst-1 | container | c-1 | runs_on | attribute | +| service.instance | shop/api,inst-1 | host | h-1 | runs_on | attribute | ++------------------+-----------------+-----------+----------+----------+------------+ + +drop table greptime_otel_resource_info; + +Affected Rows: 0 + +-- The descriptor whitelist is gated on source=opentelemetry: the same table +-- shape stamped as a prometheus source contributes nothing. +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "host.id" string, + greptime_value double, + primary key (job, instance, "host.id") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +Affected Rows: 0 + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'h-1', 1); + +Affected Rows: 1 + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id +from greptime_private.semantic_entities +order by entity_type, entity_id; + +affected_rows: 0 + +drop table greptime_otel_resource_info; + +Affected Rows: 0 + +-- Cross-signal identity: the same service reaches the graph as trace spans and +-- as an ingestion-synthesized descriptor. The two sources name their identity +-- columns differently and only the metric side pre-composes the namespace into +-- job, so this asserts one entity per service and per instance, sourced from +-- both tables. +create table graph_xsignal_traces ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + "resource_attributes.service.namespace" string, + "resource_attributes.service.instance.id" string, + primary key (service_name, "resource_attributes.service.namespace", + "resource_attributes.service.instance.id") +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +Affected Rows: 0 + +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "service.name" string, + "service.namespace" string, + greptime_value double, + primary key (job, instance, "service.name", "service.namespace") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'opentelemetry', + 'greptime.semantic.metric.type' = 'info' +); + +Affected Rows: 0 + +insert into graph_xsignal_traces values + (now(), 't1', 's1', NULL, 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'api', 100, 'shop', 'inst-1'), + (now(), 't2', 's2', NULL, 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'worker', 100, '', 'inst-2'); + +Affected Rows: 2 + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'api', 'shop', 1), + (now(), 'worker', 'inst-2', 'worker', '', 1); + +Affected Rows: 2 + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +where entity_type in ('service', 'service.instance') +order by entity_type, entity_id, source_tables; + ++------------------+-----------------+----------------------------------------+ +| entity_type | entity_id | source_tables | ++------------------+-----------------+----------------------------------------+ +| service | shop/api | ["public.graph_xsignal_traces"] | +| service | shop/api | ["public.greptime_otel_resource_info"] | +| service | worker | ["public.graph_xsignal_traces"] | +| service | worker | ["public.greptime_otel_resource_info"] | +| service.instance | shop/api,inst-1 | ["public.graph_xsignal_traces"] | +| service.instance | shop/api,inst-1 | ["public.greptime_otel_resource_info"] | +| service.instance | worker,inst-2 | ["public.graph_xsignal_traces"] | +| service.instance | worker,inst-2 | ["public.greptime_otel_resource_info"] | ++------------------+-----------------+----------------------------------------+ + +drop table graph_xsignal_traces; + +Affected Rows: 0 + +drop table greptime_otel_resource_info; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/system/semantic_graph.sql b/tests/cases/standalone/common/system/semantic_graph.sql index bfc9ea8fd7b..60d95630d91 100644 --- a/tests/cases/standalone/common/system/semantic_graph.sql +++ b/tests/cases/standalone/common/system/semantic_graph.sql @@ -387,3 +387,122 @@ drop table kube_service_info; drop table target_info; drop table http_requests_total; + +-- OTel conventions: the ingestion-synthesized greptime_otel_resource_info descriptor +-- (stamped signal_type=metric + source=opentelemetry + metric.type=info) gets implicit +-- declarations under the raw OTel column names. host/container identities are +-- the stable ids only, so rows with an empty host.id / container.id link no +-- infrastructure, and a row without an instance value declares no +-- service.instance. The same table name stamped source=prometheus is not +-- whitelisted. +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "service.name" string, + "service.namespace" string, + "host.id" string, + "host.name" string, + "container.id" string, + "container.name" string, + greptime_value double, + primary key (job, instance, "service.name", "service.namespace", + "host.id", "host.name", "container.id", "container.name") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'opentelemetry', + 'greptime.semantic.metric.type' = 'info' +); + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'api', 'shop', 'h-1', 'node-a', 'c-1', 'api-ctr', 1), + (now(), 'shop/api', 'inst-2', 'api', 'shop', '', 'laptop', '', '', 1), + (now(), 'worker', '', 'worker', '', 'h-1', 'node-a', '', '', 1); + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +order by entity_type, entity_id, source_tables; + +-- SQLNESS PROTOCOL MYSQL +select src_type, src_id, dst_type, dst_id, rel_type, provenance +from greptime_private.semantic_relationships +order by rel_type, src_id, dst_id; + +drop table greptime_otel_resource_info; + +-- The descriptor whitelist is gated on source=opentelemetry: the same table +-- shape stamped as a prometheus source contributes nothing. +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "host.id" string, + greptime_value double, + primary key (job, instance, "host.id") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'prometheus' +); + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'h-1', 1); + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id +from greptime_private.semantic_entities +order by entity_type, entity_id; + +drop table greptime_otel_resource_info; + +-- Cross-signal identity: the same service reaches the graph as trace spans and +-- as an ingestion-synthesized descriptor. The two sources name their identity +-- columns differently and only the metric side pre-composes the namespace into +-- job, so this asserts one entity per service and per instance, sourced from +-- both tables. +create table graph_xsignal_traces ( + "timestamp" timestamp(9) time index, + trace_id string, + span_id string, + parent_span_id string, + span_kind string, + span_status_code string, + service_name string, + duration_nano bigint unsigned, + "resource_attributes.service.namespace" string, + "resource_attributes.service.instance.id" string, + primary key (service_name, "resource_attributes.service.namespace", + "resource_attributes.service.instance.id") +) with ('table_data_model' = 'greptime_trace_v1', 'append_mode' = 'true'); + +create table greptime_otel_resource_info ( + greptime_timestamp timestamp(3) time index, + job string, + instance string, + "service.name" string, + "service.namespace" string, + greptime_value double, + primary key (job, instance, "service.name", "service.namespace") +) with ( + 'greptime.semantic.signal_type' = 'metric', + 'greptime.semantic.source' = 'opentelemetry', + 'greptime.semantic.metric.type' = 'info' +); + +insert into graph_xsignal_traces values + (now(), 't1', 's1', NULL, 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'api', 100, 'shop', 'inst-1'), + (now(), 't2', 's2', NULL, 'SPAN_KIND_SERVER', 'STATUS_CODE_UNSET', 'worker', 100, '', 'inst-2'); + +insert into greptime_otel_resource_info values + (now(), 'shop/api', 'inst-1', 'api', 'shop', 1), + (now(), 'worker', 'inst-2', 'worker', '', 1); + +-- SQLNESS PROTOCOL MYSQL +select entity_type, entity_id, source_tables +from greptime_private.semantic_entities +where entity_type in ('service', 'service.instance') +order by entity_type, entity_id, source_tables; + +drop table graph_xsignal_traces; + +drop table greptime_otel_resource_info; diff --git a/tests/conf/frontend-test.toml.template b/tests/conf/frontend-test.toml.template index 632947a3cde..5a879e04882 100644 --- a/tests/conf/frontend-test.toml.template +++ b/tests/conf/frontend-test.toml.template @@ -7,3 +7,6 @@ parallelism = 4 [event_recorder] event_types = [] + +[otlp] +experimental_enable_resource_info = true diff --git a/tests/conf/standalone-test.toml.template b/tests/conf/standalone-test.toml.template index ebe25f82487..8d7e307d511 100644 --- a/tests/conf/standalone-test.toml.template +++ b/tests/conf/standalone-test.toml.template @@ -53,3 +53,6 @@ retry_delay = "500ms" [event_recorder] event_types = [] + +[otlp] +experimental_enable_resource_info = true From 144f83528db8648bc5607fdc65c249bbfa3c7ffc Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 26 Aug 2026 06:15:37 +0000 Subject: [PATCH 86/92] test: renew etcd TLS certificates (#8956) Signed-off-by: WenyXu --- .../fixtures/etcd-tls-certs/ca-key.pem | 52 +++++++++---------- .../fixtures/etcd-tls-certs/ca.crt | 34 ++++++------ .../fixtures/etcd-tls-certs/ca.srl | 2 +- .../fixtures/etcd-tls-certs/client-key.pem | 52 +++++++++---------- .../fixtures/etcd-tls-certs/client.crt | 34 ++++++------ .../fixtures/etcd-tls-certs/server-key.pem | 52 +++++++++---------- .../fixtures/etcd-tls-certs/server.crt | 32 ++++++------ 7 files changed, 129 insertions(+), 129 deletions(-) diff --git a/tests-integration/fixtures/etcd-tls-certs/ca-key.pem b/tests-integration/fixtures/etcd-tls-certs/ca-key.pem index 2066dfb7b47..07a0a86ee39 100644 --- a/tests-integration/fixtures/etcd-tls-certs/ca-key.pem +++ b/tests-integration/fixtures/etcd-tls-certs/ca-key.pem @@ -1,28 +1,28 @@ -----BEGIN PRIVATE KEY----- -MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCfIi0iTllUQyNY -ka+d9Zt30ILA0QBzLKHwez9QJj8bQLpfKR26d794CQn/OO6XO7MWsJ6nvqwRNYM0 -WX+UwUtuT9MrttYWpS5yEAVO79dMtF67gCNafXKjQyK8L1vz6Avq4XJN4VqZEifT -Fdk3s2MmS0twv4sgrK//jh7aZzBQJtw+ROe9RN1dqlKoa3TtpB7iwNTDxGiR0EYj -wHj+GdiNHtDMWvrnTNNX/IKeEwLP/Mjirzr7GnaOjR25ruqcLuwWrufu+EMhqFTU -r+LIQpZD0TFuPX6DPlQZOs/4UpjteoC84m5AnQXNlh0vzjbW1efEYWYQoKDesx9e -oE5jcABpAgMBAAECggEAA3Rw/r6CYaTG1tdBilA47LF/XTihu15mh6Y4BNadESDn -IEYav0p2mDW4sgH7lWyg4ecPaBHo124jfUFM1p8Zsvk9sERwbR5vSIpWKyq5hbtM -Fo2yIXakGgJM9ev19vDR7WqxRMVAkR6GnyZpcwmh0e+ENvGZpNR1n7qSAJOJRibW -Ye8qMF/zMLF9VekBRPFJFIZBvKgHxSGCLbXRzsi5LKMbwfBoJmxEZSOC9zODl43o -Z2i+LvlS0R/Rsf2KG+Hwgj1P7HH1znJeA5bVJFDewljjSwu1mk0JpQL7Jgp2xpqs -F/mRZmtoQSH3EsTm2Kjz1sLGXwUb9JaQ274aV4t7pwKBgQDXYF+qFYip7k+b9FTF -sFbfp3CjOqxXiAxnoE5hlfV/h0w4V1Su42AdejnbIVYFX8bvmezkNWPXhB6P5oHU -4nO+oAuhsRT4317mfRM5NWO6mv2eglSBnS04GLsX142A7jMPr9e+Cqh4dYfw/A1L -bSdsdILemTFQrak27pC11mHZOwKBgQC9JhDhQf5uRh6k8GvgKGu9ZEZmkYCVu2sr -7kpD60UhklnyQo6DKYsa6OcChRQMv9ss3Hqu3Y1G0/sTe1Tjz5zwrWWAZoLlbdOt -vipesvwNTxvJoUWPtrjOoFZtmyg0Q0TqxaxfEWAssDU9IfnIy+zx7sAg4zZIPUyy -6kudaY1SqwKBgQCAQ/71Fjn7qddzc4GA8lHqhJeKPokg3/8zP78uUtaQCo2UCD6A -oR0+sOn/3MyUCsQ5MZxpFHrPgPmKjabIl8yCvGHw+7sXtD+aWOa37VnlaiSc39Vg -E7E4dVIHEvJM1I9ISlrb7REEHErHc/Se9PTDnGfMFcPO3n2mH1HDWVeQvQKBgGGn -BHH3a08tXmbTRS5uT+lwmrQbjKJBJ3x/wtG75m4Fq/BaEk9/JDUZZyKy5/4JEzPf -BGvBME4P5QFS3CndJu5O5ydaRVwDzpRVqHRJvb11SShY3Zvrvw/WUai2wRPyYuM+ -eNaAFwIbWvEb2GSle8gP9htEkuLK2w1HzxAOzYqPAoGAKl5l/WsoLfQP5UIUjSQe -JtVTDVmoz8m8Mg2nn56UdUetG246EVhDx5ngj84kZDM0GDXhfVWV6MWnIZr5NQsr -bORIDNP7Czat8FnlXeAQ+R6A/87o5g5p6ydH9jYqrrLxYjaZyXJptE+ja4zW7Lbl -ogWAc76iZvXLXqKRfepG0zQ= +MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQC05De2Tk8TMZ0J +ME+7Rp+zvvLkEBPPXCCFcXfo91gYVHMMr1TGQJuVtbx7rIqC1tFhc2luAGfsBmW3 +IfTXZjzY6834ZUzmPL7uuBfoRpd0Paj5X4owpn0eHIgcsqx4juHSTYc3fdnrTLpq +lNMulZEemkNULLCE7pyxg3hhYIx9JEe9nv26fkyr1vW5z1cPlEpyOK+SUIeqHgXi +VtWPU9WGv5sdQeXnD4E1JTX2URihb7C8xali3YVbI3SUxy/gAhM6swQgZ4SLgW8k +f87a+PcP9C1N5Gujfw4fT90dJoutyiTSbcDrjZLjF03TgjcDaACvntb0qN85XD96 +dIvhg5jfAgMBAAECggEAAd3vHe6d+KTMDZ1HwPAAKS151ipnpX3W3aRgi/FC6OMm +r0VUpWC4jCJlif3kYvaID66yFeucJveD75CJw6VaF4bSnTa9GmKM5F2BQhsrhKj+ +69jY+MU8cztGfJMdgczNfi79qDc8zSW6TUFpoMhmD6Og8icRCODbkbRWDDnnQi1d +1CVe/lZ4yIapCGgzHfQxdl9I8LEZWyeAfqMCvOTI2eE4Z4ZYto/EmuRoJN3oqOx1 +MAtpS5mw6RTFVitO4b0j5gwAN2u6Cgo8jDc5UTzHGAiGc/dAn7S+mNGv+IXP1/aL +UA/ZsBwQhZbqtGCkSmCo5gGDL+anDTiIy6nRCwIWWQKBgQDtE7V4pNj5J2MciK+h +KAR/kT36NuGoZSnf585HJVB/xAs/d8F4kU5pCofmrtm2Bl5Lx4sCE3wPmv/HMYcI +QH2zt/YxPA3F/9zoBowfTNB8/iCDlzNb/pUry+qkBhKG2HgWaUkkGEtiO16PeCF2 +bidYA+h+XZBnHKexAKL74ee71wKBgQDDVHJ6HRo1uTAT/8Uz3AF5NRj+qj9QC3UX +On6dtzbtDGsJFkcqW5CV+LwUPN6zQ7MW7U6AEeH99q1UAUDkd5Q8/4bqs+0y2Ss7 +0UUWr1LDJVYo6FIs/f2ON0e9LzcJcG9TgZ6T95TL1kGEG1Ng96Je3hAZWolLmr+x +yfjTSYKqOQKBgEwRWtTO799hx+dL5C5tTKQx0hUKrvT9IKZ7FjC1xFJ6cLF4l1c8 +KFCD1H8r8yb4fCEMcYnE/rVzIkajmZQIPU0A5bl+b1zsb9Dy6NrSJsM0NvKB/TSz +RuG6mBrw59jkdAOc3J78PJPUQM7/2JzLU0xmVJ7XHpI3G4crkSAIp/YZAoGATX6m +gF4ldOUI9xZFleKWTxFK3laLEeXJybJyY34582g22v8UsvBq96UccWcI79RPLCxw +NY1ivNBuSeLJbRsoG99BFsLVu5O/fFb1cx+R1Uxt14L8f08xlofGFX+y9TK/aEEH +uirCxPA3RANXXCRDLiIp/vUVfYJixVWdO65xgbkCgYAVQhvJ3jPXZLXPl9fVC3Ro +qweEZHs0knPc4mqL/7bE+BCE1HcccU7O7+ZKbhW2q11EMrPJ7cHUqweHE3GsTgRC +C5FIvRxInZUgtPNx2UIZAVNjxRdCZJHBc/59eu8wVp9f5ETiSFmdz5xlsJECSxdO +mBakpZMqnESl1/6FMtcqeQ== -----END PRIVATE KEY----- diff --git a/tests-integration/fixtures/etcd-tls-certs/ca.crt b/tests-integration/fixtures/etcd-tls-certs/ca.crt index 8723f7ffedd..ecd404c8918 100644 --- a/tests-integration/fixtures/etcd-tls-certs/ca.crt +++ b/tests-integration/fixtures/etcd-tls-certs/ca.crt @@ -1,21 +1,21 @@ -----BEGIN CERTIFICATE----- -MIIDeTCCAmGgAwIBAgIUeRD4bMVd0XEzS3UHo2ptnknsWzEwDQYJKoZIhvcNAQEL +MIIDeTCCAmGgAwIBAgIUM4xko7n+KubCgGJnX4EaUzP7ZicwDQYJKoZIhvcNAQEL BQAwTDELMAkGA1UEBhMCVVMxCzAJBgNVBAgMAkNBMQswCQYDVQQHDAJTRjERMA8G -A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjUwODI1MTgwNjU1 -WhcNMjYwODI1MTgwNjU1WjBMMQswCQYDVQQGEwJVUzELMAkGA1UECAwCQ0ExCzAJ +A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjYwODI2MDQxMzQy +WhcNMjcwODI2MDQxMzQyWjBMMQswCQYDVQQGEwJVUzELMAkGA1UECAwCQ0ExCzAJ BgNVBAcMAlNGMREwDwYDVQQKDAhHcmVwdGltZTEQMA4GA1UEAwwHZXRjZC1jYTCC -ASIwDQYJKoZIhvcNAQEBBQADggEPADCCAQoCggEBAJ8iLSJOWVRDI1iRr531m3fQ -gsDRAHMsofB7P1AmPxtAul8pHbp3v3gJCf847pc7sxawnqe+rBE1gzRZf5TBS25P -0yu21halLnIQBU7v10y0XruAI1p9cqNDIrwvW/PoC+rhck3hWpkSJ9MV2TezYyZL -S3C/iyCsr/+OHtpnMFAm3D5E571E3V2qUqhrdO2kHuLA1MPEaJHQRiPAeP4Z2I0e -0Mxa+udM01f8gp4TAs/8yOKvOvsado6NHbmu6pwu7Bau5+74QyGoVNSv4shClkPR -MW49foM+VBk6z/hSmO16gLzibkCdBc2WHS/ONtbV58RhZhCgoN6zH16gTmNwAGkC -AwEAAaNTMFEwHQYDVR0OBBYEFOCkI3Uyx7F38LtXtSlg3ORE9ur6MB8GA1UdIwQY -MBaAFOCkI3Uyx7F38LtXtSlg3ORE9ur6MA8GA1UdEwEB/wQFMAMBAf8wDQYJKoZI -hvcNAQELBQADggEBAJXbWJJ9b+5nXjRiFXOg+wQnzn7kzMf6sWKIFk+AuKJWWt78 -O00t6vyAz6zel5Cj3ho9yAaMFNy8vYEnJCYngy5pT/2hOncnz/w7IKTeoEhzqAnf -MCEmCgHbTKDoFfMfrrRwtyoePVfx4xbiGVWMQFTPG+WNlE/ivFMRvFwsgPJ+7SUK -mR2FscH6DVo4sqF6s729lTmr6/U7bOD5l2HYGpKJ2cjCI8+HDv55aAT43tsTwBoJ -BgX+RT4w8ryGhYV+hrRTjPlMHWiHsNbzJeAi5bNA0f6vUP6k6zHn/Ur3MNb1avcf -cSsdYU2dhs7PeB6IpoJ0QCeJw9MIFK3XTnmbeVY= +ASIwDQYJKoZIhvcNAQEBBQADggEPADCCAQoCggEBALTkN7ZOTxMxnQkwT7tGn7O+ +8uQQE89cIIVxd+j3WBhUcwyvVMZAm5W1vHusioLW0WFzaW4AZ+wGZbch9NdmPNjr +zfhlTOY8vu64F+hGl3Q9qPlfijCmfR4ciByyrHiO4dJNhzd92etMumqU0y6VkR6a +Q1QssITunLGDeGFgjH0kR72e/bp+TKvW9bnPVw+USnI4r5JQh6oeBeJW1Y9T1Ya/ +mx1B5ecPgTUlNfZRGKFvsLzFqWLdhVsjdJTHL+ACEzqzBCBnhIuBbyR/ztr49w/0 +LU3ka6N/Dh9P3R0mi63KJNJtwOuNkuMXTdOCNwNoAK+e1vSo3zlcP3p0i+GDmN8C +AwEAAaNTMFEwHQYDVR0OBBYEFDjvuYaWEoViLOKyCse6mAwH66mIMB8GA1UdIwQY +MBaAFDjvuYaWEoViLOKyCse6mAwH66mIMA8GA1UdEwEB/wQFMAMBAf8wDQYJKoZI +hvcNAQELBQADggEBAJuGxtprq/GF/fwUJ5hk0cG80VY24G+iKKi7Jw7PWPEMvnlQ +/ZIKgNDupE8mSSyokicKZMrtwT8QF6wlDLVcz4LTPCtKUmYU3FiuNE4V57+yw7zt +XZn6rwiWGxeyfY0ZktLXiSYA5FydOJ9GQQWgxkpbd5QG1L9aoFW49cbz8ob6efxG +p+YRgr0Fr4hXn7ImJWWqeawf2tg4iLwpV4sem5atpI7y9qamxrGGUp2h/CnADqUz +YWYIYP1RCPp2bXzxO40Ni1NYmFsGUU/SL1vtJw5n2Zmx/Phe7u/CLe05dH1O54Qf +vCbk6X0CthFI0HzSepbEVAuniVHkrpWwDLHcQ6s= -----END CERTIFICATE----- diff --git a/tests-integration/fixtures/etcd-tls-certs/ca.srl b/tests-integration/fixtures/etcd-tls-certs/ca.srl index 69ac026da25..ad32d06e16e 100644 --- a/tests-integration/fixtures/etcd-tls-certs/ca.srl +++ b/tests-integration/fixtures/etcd-tls-certs/ca.srl @@ -1 +1 @@ -4484A9BBD4F25F32994F3C03D80D294105B119E6 +4484A9BBD4F25F32994F3C03D80D294105B119E8 diff --git a/tests-integration/fixtures/etcd-tls-certs/client-key.pem b/tests-integration/fixtures/etcd-tls-certs/client-key.pem index 13a23c53478..1d64fc66a3d 100644 --- a/tests-integration/fixtures/etcd-tls-certs/client-key.pem +++ b/tests-integration/fixtures/etcd-tls-certs/client-key.pem @@ -1,28 +1,28 @@ -----BEGIN PRIVATE KEY----- -MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQCOcrzRmE3sEGYo -pazDrpkkzbiaBhHClfmBiZAAGV2cvSXY8e9pc21GOgrGBNeyPPieS305rAE02EZl -oIAt97XRFxkZniUy8C5+PT3eicAKPRcrgGhIrYxUs39rdRocqVVpZsiZGzhHhboB -bXNW1l37SDQpuxAQJuhHu2wu9J7c5/TBXWN9ZOxzWTlrIdEQOA46VclE/N1TlaKs -CSWw8qH87YLwkLEXcFLOJfmWNZxOOhdL5MSzKKNcQKbpmkd65B/PHNwBIBZHpI+K -xphK1SybA+ESkN09qZaa1Kj3e0I62UjHNXUCuAod6RhRRmpiFDFgggXhD9T9Ttuf -1Lq38XxDAgMBAAECggEAANNhqQldwHzi9WCxXObLoufBe9lM73+1Df1IQaDHkZR3 -jHliixxF18XQW7qZ5tAJqjsDWLPm4BxSKXozsjsvBdgIUjbqkyx1dsGlO9TjeGu5 -H2/8jlH14boIMWJ9Y0JBWwvU7EnVU5IkUOEiwsxvFlzM5Q0H+841CrSD+Xd7iSZE -ZrpyZgpHNmvc/u7AmseZVccjX21dxwWtcd7cGUmFWC9o2a/hkgG6+baA8zH2ze3J -Smeowa/oOyOVkTM1EHLCBvddE+/W+BeMGPY7A6SztVw0gLqlSOL3cr/g2llUTmFg -T6MOS6FZZ7MD5c2HBL0V3OXIhxNABzD+c0O3CaRbSQKBgQC/+QuhMpDCzxBPaL/K -sqRp4sQgURDCXPA1nAfI+NoHos/CbjwcMGp6LUSiBW0tLTJ+LZoVgZXj6j/uxji9 -DWxQfULvQ7oanm/q/D+40NnupQQTzmMT04z/mV3XjFtzEhbO30YZfqwXy+tqB8vR -g1KajXVqLmCxKDtsFedAlIVxRQKBgQC99TKGcBwsKlYJJThU2eIIy8G0GkjjXleK -ROMO5i0keJtLac1oRvXo4JN5UBLJM6ZpSy0TEa2bxfzy6Uea+41aZmdCpPg0Kcsq -Pb1S3rsPpVdsfxSxsdKvwZmHXOzAdF+fom1+ZRAdo/8ALIPaNQSSaw+KkmmpHPEV -qdzzWJ4b5wKBgFubEtKcF3nuZwENoh+ueUhRvncRV+b3hGSAjTJ4lUn5hhxoj+R/ -sf+VJGAQKNXa8HJHfnRuvsDgYhulmSOViS8rZspXzjGvkwZV0m51stjvA3AUFzE5 -zNmXLLGTt3vEkP+siX3W9XXxh+ezyq2ydbNsdy/w65D9+sUL+qrVdIvlAoGAOOQy -2ajCB0g2tE59bIxE8jV0MiidI9uhhDvVdSTi6EVm3VM2vcBi7fg0suSUe8YIVQi6 -2zc0M688btQHKhek4ipBSuh1ncnWmzQae7NRewIeCNSWshF79D+bZ7sg/RLdgMX4 -3R4PkZEIUlkCtFuknuWJpgrrskaEveQ91HP6BokCgYAX42nKPANIMrTXuJ2vACWu -qKQ9atAC1ngVzYe2PI8rHjA4gCDraGfCaILewrd89rK9ukBeSDIm5Mn448q3ystw -p/QMxmMdHADCD6UwCjlt9N0Xigv+IKbDzDzhartRVIPpZdTViQXW3cSHt9id8zvO -HTu5C6W+N9Vb7j8Ak/mvzw== +MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQC5F3hVsri2FBzP +dtRe+aJFh9yqhXnJ7zN+xbUmNWsTZVnPwrbVo9H6PtGWzkbZ6YMBzMlQSYP2txaS +ROjE8wwEs38eKbyJrHO2Ccq3iR9NfDxSbHTClngsOjf9fEoYExif84VKWLO6+EKL +KqqOfj3+iuI3hbabr0MqGu4e2Nr6MlQtViBxTs5aSnahc8hmZbujBcWRTsw/GNx5 +Am0a8BrrJRVmVR+BzlYtzbf1ltfLxIhnc8dAGyT9B6GTwcMNST6m4WtTzR/uOr6S +2YEbIqwCHB1VQ+oJUGJvExqaNLoM/x5a8sUNIWx3LBN0/XC4D+MMmh/bbyb2cAOE +gyUYu+hLAgMBAAECggEAEkt+wksVP9toEkecFpEqgss65VimdkPK/WkHDpzxawGL +FDfvq4Py9tRrvLIXRbseObjSCI9jrai2Ne9Sv1jeDhQ8kzDWk8MfItX/8A6Whf8A +KrjRUlneIfz/2Mcxaa9wMSs4YxC2dIYidFqwxhSLWyluJ8Ud4hs7kndb+beUmW3i +DgkLsQ9ytQ0HDQrAERnrbRf19fmVaqQapfxpYT2gg80BMIq+zlzoiXareDz3r/1g +aoponxArWV6yCvQOYYjWcDr4WZPvIA2APEX5isxyvMBvb72Fpd2z+0JX+4LVe6N1 +FqZLNpSeJn6O+Q1vJSScSJcoJYkLWCfZbFLHPttZGQKBgQDdVr5PtgPIVM+67cvv +p3x7NNC10Gp387PgQS/FlG+QtbP87RTBnANT1uFd/FMFqL35or7Y46YKDUdk00OZ +LinpM4ARFRYFqkthM6bcUWatpF6Fg2dfsKuD4H6ARp8YHhlmCnDCQ0swZp1HCmAB +6vltUNTZjQjRdd44GlQ0icX0mQKBgQDWE50TXRszea75S+zUfXeyvxMyptIOS/rY +L50XVuhGVknH4OuBQbGF6UZulIGSpN7gmmbO29Yx1HKmqiJVaWlvB71baK4hWLDd +2grIUy0zuq1B7I1IOLmOpMcL3xau76g7ZWRUprX26hxAuxYyPb9V8gNo1P7HN8si +VLJVE8ZugwKBgQCaxoGmM91JRSVNzeOB3ljJvxEDUo5g+uWZt3u0aivpwWXvQ8nz +6Sjag7RsiHl1x52w5wEVoXsGJGr8Mk9e2k0saXrwdxJDO+YiPoA8KB/o5LvEGTM8 +UspdGarcAIZX0xRnqn1XGr+FRPxOJQ8lyC5LJu7wghLchdOy35Zqdr0aYQKBgQC/ +38yFsonS1VnS8A5RVjOW7lPSrlrPnaIzalmutaJyiJyQnjP3Il5u2+rY6hpIyaVK +QpmrBrcw6m3om80yKMzrS1CZQXXxRYEhF3Fao9J77vGjiNYIyW7nPyF4rneyS/PJ +aNNIXDP0H1k7W3RFi7qW2dfceivxezyChM9iGdtc6QKBgAIfb63+1q98MNlAB7eD +x7K8vuSxHrJ9gX9KcLrG1ScLOxnAC4Op+8xf06ZHt828CajchAOqee2nh52MsFOr +AdPNdhS/8mTq8KosTfgVuawv4yTHpZViMR2KSOcTKhKaBmhufY+k8mJ0YAAOTnzO +JJqjDSJdjWsDG+7UlNgFcvuU -----END PRIVATE KEY----- diff --git a/tests-integration/fixtures/etcd-tls-certs/client.crt b/tests-integration/fixtures/etcd-tls-certs/client.crt index 935351a48b1..0dee5915ffa 100644 --- a/tests-integration/fixtures/etcd-tls-certs/client.crt +++ b/tests-integration/fixtures/etcd-tls-certs/client.crt @@ -1,20 +1,20 @@ -----BEGIN CERTIFICATE----- -MIIDMjCCAhqgAwIBAgIURISpu9TyXzKZTzwD2A0pQQWxGeYwDQYJKoZIhvcNAQEL +MIIDMjCCAhqgAwIBAgIURISpu9TyXzKZTzwD2A0pQQWxGegwDQYJKoZIhvcNAQEL BQAwTDELMAkGA1UEBhMCVVMxCzAJBgNVBAgMAkNBMQswCQYDVQQHDAJTRjERMA8G -A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjUwODI1MTgwNjU1 -WhcNMjYwODI1MTgwNjU1WjAWMRQwEgYDVQQDDAtldGNkLWNsaWVudDCCASIwDQYJ -KoZIhvcNAQEBBQADggEPADCCAQoCggEBAI5yvNGYTewQZiilrMOumSTNuJoGEcKV -+YGJkAAZXZy9Jdjx72lzbUY6CsYE17I8+J5LfTmsATTYRmWggC33tdEXGRmeJTLw -Ln49Pd6JwAo9FyuAaEitjFSzf2t1GhypVWlmyJkbOEeFugFtc1bWXftINCm7EBAm -6Ee7bC70ntzn9MFdY31k7HNZOWsh0RA4DjpVyUT83VOVoqwJJbDyofztgvCQsRdw -Us4l+ZY1nE46F0vkxLMoo1xApumaR3rkH88c3AEgFkekj4rGmErVLJsD4RKQ3T2p -lprUqPd7QjrZSMc1dQK4Ch3pGFFGamIUMWCCBeEP1P1O25/UurfxfEMCAwEAAaNC -MEAwHQYDVR0OBBYEFINcPBk3iqCw1B2sZzcDpSF6ZNGnMB8GA1UdIwQYMBaAFOCk -I3Uyx7F38LtXtSlg3ORE9ur6MA0GCSqGSIb3DQEBCwUAA4IBAQBgyF1XSpAyiQNA -sFCqXrBfwPnv6V+R5jhO6Glmn+Nhj0gbzNBlU8EKsK/WZllDrJjmxubPBqb563Zt -J+QcNLqugZdioXPbxmmi9xj3oK25pUCRBj03nlHpGqIwuCJyn/ZxlfXMpNYE1KjM -nprFNUKgnYPk2AGffAngQdm0yrnmSpxXhNTuhQE8avgRhd5nEUEDyPIvHWwo4iU6 -pw2aP2FpWorESv9LVna+JortThYDQrIFgCtPM6BP1N7Aeqn0H846+JsXlq2ch1Iy -aBRqKNQRc1Q2Qb9/QCDEJ3dEixgFjxztIaqrIcr18ZJ797MhTFS3aVEwE8clSRR5 -9ow9jK8B +A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjYwODI2MDQxMzQy +WhcNMjcwODI2MDQxMzQyWjAWMRQwEgYDVQQDDAtldGNkLWNsaWVudDCCASIwDQYJ +KoZIhvcNAQEBBQADggEPADCCAQoCggEBALkXeFWyuLYUHM921F75okWH3KqFecnv +M37FtSY1axNlWc/CttWj0fo+0ZbORtnpgwHMyVBJg/a3FpJE6MTzDASzfx4pvIms +c7YJyreJH018PFJsdMKWeCw6N/18ShgTGJ/zhUpYs7r4Qosqqo5+Pf6K4jeFtpuv +Qyoa7h7Y2voyVC1WIHFOzlpKdqFzyGZlu6MFxZFOzD8Y3HkCbRrwGuslFWZVH4HO +Vi3Nt/WW18vEiGdzx0AbJP0HoZPBww1JPqbha1PNH+46vpLZgRsirAIcHVVD6glQ +Ym8TGpo0ugz/HlryxQ0hbHcsE3T9cLgP4wyaH9tvJvZwA4SDJRi76EsCAwEAAaNC +MEAwHQYDVR0OBBYEFD1aTqJO87ovA2usg/v8UckSsnwhMB8GA1UdIwQYMBaAFDjv +uYaWEoViLOKyCse6mAwH66mIMA0GCSqGSIb3DQEBCwUAA4IBAQArP3mtjdj5a9uT +XTys4DUjp+CJaSmsZ/4ICf/KeddaEyijx6uoiqbNsAtfwlD5A5BtyjJrxg3SpjqD +Fm5mTV0HlwvYHxcFPTUlR56//sOyeRhm+kVPp4xdVq2p6IZXFF/hDkeT/AETWOjj +l4+90gAsJsUHrKUZmzlqHBeOXoAhzo2PwSh8KLcGK0705FGql4QnS731PHo9a9qp +e8uxk1KKoc+RHm2OjJvZ7hvkCXuCm3TJeOHQEtjZ5buQg40BRrt11GVnsPePpnQR +8OtV0DvPT3UcOss88kLDQuRuZeJf8UyaevGGvItv3+Yk5cwUBSaJ0mY6sgrmkXn0 +TstUzYIT -----END CERTIFICATE----- diff --git a/tests-integration/fixtures/etcd-tls-certs/server-key.pem b/tests-integration/fixtures/etcd-tls-certs/server-key.pem index fca8a0122f9..a73ac974434 100644 --- a/tests-integration/fixtures/etcd-tls-certs/server-key.pem +++ b/tests-integration/fixtures/etcd-tls-certs/server-key.pem @@ -1,28 +1,28 @@ -----BEGIN PRIVATE KEY----- -MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCtYFxH+oet6jTD -0eF55ISO+bwm1NPTluP9zLJrQFNDkqlddyiIasK5q7qN1pPXOw9mJF/EN8Hf1lWj -42an6tr7PgE0NHLcOrqDv4dRjegRi87IWiPkfep8peUE3glI4IOMXfZmAUNp+nRO -JP/ymmR0pzHYJlx+D082OYmeBny7uFQxo+vxgVNo2MVVwy6oTv+jtjLgZ4iOrIvh -iILoXtXcdg/tjB1rS2hSABGPJMNHwKswcrt6xXv+7Llsxroq0cx0E2PrAbc5qB40 -v1Qp7PjFTJ0drgf0g51LkOb7znpB1viPD2BLMxpELTZqa4OruarcDNxShrK8WWMP -1mygLHLxAgMBAAECggEAC4TgtLaQP62VImKGLs1QQliU1+ahiUgX70Ojog0UyyNK -JefWFVQ0ilX+z9AvI+hsYj6t7zE+LBNHPtuLtUHdGT66IUAP1pJ/VGQMB07cmZfW -ngihJFv6UZxLDkr7RnCGRPP0PDw+wKKPiiaaq8F2xapbHS+VSxnUyzdA7bMkI+uj -+gLAFj+kEU0FNbIp/mBzz8sksDzcB29oy3VnZkR1211KkL1BJDvr9oKiUg8NFKaN -NF2WXRYI/m0CQMlX6zlTyv7+h/Ay7y8oNG7cfHavmCrWvYIMfaDggYxbIpNdGxb5 -MAdbrtXJ6daUBa2sezxgQUyJX8+8GGF2zpGvCfhkfQKBgQDiAwA8rlxrBnXclwRK -9s/apILKCxPEGwU6XZi/8na3ugIO14y2WD0zZTjUK/3xwoKN6j6+ObLkUBX2RRYZ -ewwxpqos2cyxi5EBKPhsydqbJQlat68I5+NOG4KOZQkqO92tBtgY3gGFWVntaqtE -UVE/Clh2gFcc/5nR/5C5Mil8LQKBgQDEYXtrl2kl4qcyXuuXKDA/bLSnDuIEMhLO -4oIiqmywXQgjPjGAWy0eztjMq2a353pqTFCqto/cfOif+rjkB7oFfi9gMm0plIKI -+F8QsBB7HpGqWos3zb1yt8X9hbxKuFFwZS/un3Li4lwPlk8EZ1YIvhC44BHAvNkD -VzxzBwoYVQKBgArRe/RroC7bS074x5LTB5X+o+gJ6bNMW860Zjhh4b7fn3OYa7ra -tGs+YB7/0BL/bYJfgQtX9bEqCDMWkX08v5Os1553+m1RMeqtTF7gtp8Qgcce3bj+ -aIn3lSM9wNeNsAm1NyjRj58TbNOJdJM7lTkARMW/VOwla/Z6VjIXLZctAoGBAJKf -xira7eMfi36MaJJ/qyZv36Ir9ozzZh+Z91gyrtwvWfgWY5dWfCXYgv6tqw/8gOYE -/OW5UUhq6rUn2gxHyJh5Up4ciGzXOW9TIoevLV7/v/rVh8SulJimpelYhPG1FPk6 -U8NywbCtGdd5fp3nGdGFN68Rfa/OUKmx5KxtwRfRAoGAWpfzmoI8EmCBZNXYk03k -hnDKmICsKi9seDclcQq2PbyNvE7LDEd8O4pDnwV6RwKVEDTer9Y4IGcjkEBZcJVm -za3l4En9+M27Xlf1I8XeB8o8RYbFjoDQyScFFIgQpvTDsOezXeDR55FGPZN7V6Si -lKmd+8BuzJ7d31KJtRjoS1U= +MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQDQLtVfdk9l1rdH +/ssuy/BK9LpmNPx7gatp/E3EqpgekBjizMXIxnCHDFNM4Dtq1v3afaxeijXnnW/b +0CVKSeGkOBmb0dO0K21xe86Y5VfNHkGotQyCZ8mPOzep/EQOGJY5+dprErL0YKM+ +2drgiRNGwidmNk/jD3OtzXSMR+8UvDWJSxUoZ1ncvgmqf3zTTiaTiABHrfLdXaGI +23+JftyLgoeU4vEI+x+RQJMAa6Gzjo9vALxX3kTtR/uH34hfm8vb3yo68UHGKVMp +i3yi6YDP4YEao5lOpm7fIyEes10fPkD224++S34KHykAk4q+eTPetWkDYFCYzdSp +AR2Ob3cxAgMBAAECggEAMDgAxO0mw7xBVGQgFJU48WuQtvaj2kl09gXxz1ECDeYr +VXC/iNrpmmYQ7zfqmzrzrkU4hOc3SA/PplamJHhLUpmJ2Oz3P35liYj3F7PbK8/L +vnM81AGNDmdVY8Jh0u//76q+29kHaRHvDbIw/5vQQq3aqVKAG2PrU8DIM2u/5Qmt +0GHJTMmyVADAOLe7j2agN9hrOrjeoWobcELwtwgKXxmU0l+g0XX1OdJpbjA3qDPa +S5a4MzcucGINeB3bohcqbcx2f1HwVUoC19mti3Cz+KvxF2YZV9bS9izW4hSEsAwU +0upNHF3MwTUQg7ftLPAZcfNoXIZMFXwYMS1HW2v7HwKBgQDuHHkFRN2lvor/Ycjr +gJD6LGb9RlQ7NMYBQHB+5g49Y/q8oDZJgyl4CfgYWdVjjOfEX8oMst1WU/Stbva5 +HlHR7yYFyebScQYOPHZaPyVKGIbcl2p7Qy9H4pxJc1ZMNxO8SxoSC68LZEan08FD +/Qsish51fZ2Sqf37b+8Uz/CGdwKBgQDf0sIp88MPGbhtzdYtF9gA4AoIOPo9UxeI +F6wTb05uCqZyK+hcdK37av1xcNpczd/yvlALxRdM9xGedJkcS2cgeExtMRcfhAzL +as0dYPzuhRkpEIfN/y9cqtHj2wigrbu35eDgRRnJDgXgL9dYl6r+Swqn3jBaasP7 +KtQdQb/RlwKBgQC7qe0n3fLi4p4iUStNkPKyebRiAb/5Ocqkyejf2ul2MQo5B/xB +TAKu/QxwBL1NzIwOFYDlKUOQ+nJpDn+dvuu1jcpl3Y7yZOnk5npQ/luhXltMGHpv +06+79DpBGYn2X6JKUNanSlYXoFyfgSFdOF5CZifjabF7Gkd2l+3SdWCYWQKBgQDB +3k4sBGZKaB7ljUscl/CTIXvPD3tBLv3M9aQo2Vp32mW9suZ7Xt1sTonkfrnFdNWr +7shqyXabRc5PD/OnHHDhIRIh6kl7FOf4MjQkZGPxPfxDI3xeI9EkVRmkYY6hjppw +eX9FAtWI3sqcGxROOmD0Do/WQ5BiYOQMZFaCWPcLVQKBgGVT+iUYkYTs+UJeD4RC +QeL6QEJDszvAI/gavUp7sD/LWn7Ig+hJiAl2O+HFYMurh7eda0JnE2KF40UyPiYA +DDl19xRAXrLPOlOj9YKDA9HkROcaVZCQCqyDalKsbw1IqMrAkRmLCIZW3EHaMrUf +UNjg+GaxwvrRj11gbd4IPgai -----END PRIVATE KEY----- diff --git a/tests-integration/fixtures/etcd-tls-certs/server.crt b/tests-integration/fixtures/etcd-tls-certs/server.crt index 6ed8b1054e7..8dcb7b13b6d 100644 --- a/tests-integration/fixtures/etcd-tls-certs/server.crt +++ b/tests-integration/fixtures/etcd-tls-certs/server.crt @@ -1,21 +1,21 @@ -----BEGIN CERTIFICATE----- -MIIDjDCCAnSgAwIBAgIURISpu9TyXzKZTzwD2A0pQQWxGeUwDQYJKoZIhvcNAQEL +MIIDjDCCAnSgAwIBAgIURISpu9TyXzKZTzwD2A0pQQWxGecwDQYJKoZIhvcNAQEL BQAwTDELMAkGA1UEBhMCVVMxCzAJBgNVBAgMAkNBMQswCQYDVQQHDAJTRjERMA8G -A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjUwODI1MTgwNjU1 -WhcNMjYwODI1MTgwNjU1WjATMREwDwYDVQQDDAhldGNkLXRsczCCASIwDQYJKoZI -hvcNAQEBBQADggEPADCCAQoCggEBAK1gXEf6h63qNMPR4XnkhI75vCbU09OW4/3M -smtAU0OSqV13KIhqwrmruo3Wk9c7D2YkX8Q3wd/WVaPjZqfq2vs+ATQ0ctw6uoO/ -h1GN6BGLzshaI+R96nyl5QTeCUjgg4xd9mYBQ2n6dE4k//KaZHSnMdgmXH4PTzY5 -iZ4GfLu4VDGj6/GBU2jYxVXDLqhO/6O2MuBniI6si+GIguhe1dx2D+2MHWtLaFIA -EY8kw0fAqzByu3rFe/7suWzGuirRzHQTY+sBtzmoHjS/VCns+MVMnR2uB/SDnUuQ -5vvOekHW+I8PYEszGkQtNmprg6u5qtwM3FKGsrxZYw/WbKAscvECAwEAAaOBnjCB +A1UECgwIR3JlcHRpbWUxEDAOBgNVBAMMB2V0Y2QtY2EwHhcNMjYwODI2MDQxMzQy +WhcNMjcwODI2MDQxMzQyWjATMREwDwYDVQQDDAhldGNkLXRsczCCASIwDQYJKoZI +hvcNAQEBBQADggEPADCCAQoCggEBANAu1V92T2XWt0f+yy7L8Er0umY0/HuBq2n8 +TcSqmB6QGOLMxcjGcIcMU0zgO2rW/dp9rF6KNeedb9vQJUpJ4aQ4GZvR07QrbXF7 +zpjlV80eQai1DIJnyY87N6n8RA4Yljn52msSsvRgoz7Z2uCJE0bCJ2Y2T+MPc63N +dIxH7xS8NYlLFShnWdy+Cap/fNNOJpOIAEet8t1doYjbf4l+3IuCh5Ti8Qj7H5FA +kwBrobOOj28AvFfeRO1H+4ffiF+by9vfKjrxQcYpUymLfKLpgM/hgRqjmU6mbt8j +IR6zXR8+QPbbj75LfgofKQCTir55M961aQNgUJjN1KkBHY5vdzECAwEAAaOBnjCB mzAJBgNVHRMEAjAAMAsGA1UdDwQEAwIEMDBBBgNVHREEOjA4gglsb2NhbGhvc3SC CGV0Y2QtdGxzggkxMjcuMC4wLjGHBH8AAAGHEAAAAAAAAAAAAAAAAAAAAAEwHQYD -VR0OBBYEFK8y9gWv8Si9AusP280BrgDF6y3jMB8GA1UdIwQYMBaAFOCkI3Uyx7F3 -8LtXtSlg3ORE9ur6MA0GCSqGSIb3DQEBCwUAA4IBAQBMBsAPIlk17I2ioN9xwxwl -yYbSjZlTs+18wSZAoCNLfxwWIYAQ08dPoEUdALfMUPEEe1Ol2IAp/qx2JoDrZWje -maA43hppBUpFFCkTKWUMYxsetN7d7BVxWL43GDoMwoD/k36nhoUKBUjlbF0+nkem -dOMr8SA4GZbEg7qk2cL93g0UHv/Z/dZgKf3epZkR9hEN4/R2jSP6OPmY9XIvQCoZ -8D9jgbQKavAAXmUcP5a81alQMRroGaBCzI1f5OlS3EuVE7ZTEBgxKK1idbouSrPt -4UHEbOGz9zXDjmfut6CTm247+lJm9jzYe2Xx+XGw29l0pzd/8tGFN9zIvW8mCv7a +VR0OBBYEFNqGaSuuVHwEIvQvYxt9H0kdWFPSMB8GA1UdIwQYMBaAFDjvuYaWEoVi +LOKyCse6mAwH66mIMA0GCSqGSIb3DQEBCwUAA4IBAQAHttiuyEZieeTsggAKzAIC +UXsjWCQMAmjRxery6mzM8P33CdCKA4cwGyy6tzoFPiZC/wh+gPq+LiUO9G+vkqNA +NQ/DSJS+fsqsOBWUt5DSYekD25o80mZVE/4H/ynFOgENBqgkVl1ikpkXu9BK8wTm +Wij+p9RWyNup8lr9OFYZRpy7wrxPp7opwp3qG2Bfu/BGY1YEr04CIrLMKnk1zRMm +73X1lcrlrM9ycCmYcfXxyU69G0iFKaPcbS7T0fngm1NciA1qxYYLuMyzDtF7/SyC +J2ynvDBBnc9+/X8piu8p9D0I9mbcbEySv6zZfYMT0kqXbsn2GijZTH5mCJ3P8cwc -----END CERTIFICATE----- From 35ea88a4ef4b06c8c5d10e57cc8bad44f4f38569 Mon Sep 17 00:00:00 2001 From: localhost Date: Wed, 26 Aug 2026 12:11:14 +0000 Subject: [PATCH 87/92] feat(ci): run query regression on ephemeral Aliyun ECS runners (#8937) * feat(ci): add aliyun ecs ephemeral runner path for query regression Signed-off-by: paomian * fix: improve condition for query-regression job execution in workflow * feat: update Docker installation to use official repository and add GPG key handling * Refactor query regression runner setup and configuration - Removed deprecated PersistentVolumeClaim for build cache. - Introduced a new bootstrap script for setting up the ECS runner host. - Deleted obsolete Helm values files for runner configuration. - Updated the Aliyun ECS runner provisioning script to reflect new cache paths. - Modified GitHub workflows to use the new Aliyun ECS runner setup. - Adjusted documentation to clarify the new runner lifecycle and provisioning process. * fix: enhance runner service management during bootstrap process * fix: update alibabacloud_tea_openapi dependency version in metadata * feat: enhance ECS runner scripts with region_id and resource_group_id support * fix: move containerd content store to data root for improved storage management * feat: rename query-regression runner to ephemeral-github runner and update related scripts * fix: update sentinel polling method to use serial console output for improved reliability * fix: add environment variable checks for Alibaba Cloud access keys in ECS client * fix: improve error handling in GitHub API requests for better diagnostics * fix: improve cache disk detection logic for Aliyun ECS instances * fix: enhance cache disk waiting logic with detailed output and error handling * fix: update dependency version for alibabacloud_tea_openapi in teardown script * fix: enhance cache disk waiting logic for better compatibility and clarity * fix: enhance console output handling and add incremental logging during instance provisioning * fix: add PATH environment variable for runner jobs in service and provision script * fix: add machine telemetry sampling and logging during query regression jobs * fix: update query regression documentation and provision script for cache disk handling * fix: update SCCACHE_CACHE_SIZE validation to 10G for improved caching efficiency * fix: remove outdated cache size checks and cleanup logic for fresh system disk runs * fix: enhance instance deletion logic with region handling and console output export * fix: add swap file setup and OOM handling for ECS runner to improve stability * fix: update OOM handling and service restart logic for ECS runner to enhance stability * fix: increase system disk size to 100 GiB for cold double nightly builds to prevent ENOSPC errors * fix: increase system disk size to 150 GiB for ECS runner to prevent ENOSPC errors * fix: add keep_instance option to preserve ECS instance for post-mortem debugging * fix: disable unattended upgrades to prevent job cancellations during library updates * fix: reduce system disk size to 40 GiB for ECS runner to prevent ENOSPC errors * feat: Refactor Aliyun ECS runner provisioning and introduce nightly regression comparison - Update `aliyun-ecs-runner-provision.py` to remove cache disk handling, simplifying the provisioning process. - Introduce `query-regression-nightly-refs.py` to resolve and compare SHAs from successful nightly builds. - Create `query-regression-nightly.yml` workflow to trigger nightly comparisons based on successful builds. - Enhance `query-regression.yml` to include a `test-tooling` job for validating Python scripts before provisioning. - Update tests for the new nightly reference selection logic and refactor existing tests to align with the new caching strategy. - Modify documentation to reflect changes in caching and nightly comparison workflows. * fix: enhance runner image tool verification with detailed checks * fix: improve error handling in runner image tool verification * fix: update tool versions in ECS image and workflow for consistency * fix: correct typo in error message for unparseable ECS creation time * fix: update README and workflow files for query regression tests and image hygiene --------- Signed-off-by: paomian --- .../query-regression/README.md | 453 ++++++---------- .../query-regression/cache-pvc.yaml | 15 - .../ecs-image/bootstrap-runner-host.sh | 195 +++++++ .../ecs-image/build-ecs-image.py | 392 ++++++++++++++ .../ecs-image/ephemeral-github-runner.service | 45 ++ .../ecs-image/start-runner.sh | 37 ++ .../query-regression/values-8-cores.yaml | 146 ------ .../query-regression/values-paused.yaml | 2 - .../scripts/aliyun-ecs-runner-provision.py | 485 ++++++++++++++++++ .github/scripts/aliyun-ecs-runner-teardown.py | 237 +++++++++ .../scripts/query-regression-nightly-refs.py | 270 ++++++++++ .github/workflows/checks.yml | 8 + .../workflows/query-regression-janitor.yml | 36 ++ .../workflows/query-regression-nightly.yml | 104 ++++ .github/workflows/query-regression.yml | 433 +++++++++++----- .github/workflows/release.yml | 2 +- tests/perf/AGENTS.md | 20 +- tests/perf/README.md | 7 +- tests/perf/test_aliyun_ecs_runner_scripts.py | 122 +++++ .../test_query_regression_nightly_refs.py | 120 +++++ 20 files changed, 2538 insertions(+), 591 deletions(-) delete mode 100644 .github/runner-scale-sets/query-regression/cache-pvc.yaml create mode 100644 .github/runner-scale-sets/query-regression/ecs-image/bootstrap-runner-host.sh create mode 100644 .github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py create mode 100644 .github/runner-scale-sets/query-regression/ecs-image/ephemeral-github-runner.service create mode 100644 .github/runner-scale-sets/query-regression/ecs-image/start-runner.sh delete mode 100644 .github/runner-scale-sets/query-regression/values-8-cores.yaml delete mode 100644 .github/runner-scale-sets/query-regression/values-paused.yaml create mode 100644 .github/scripts/aliyun-ecs-runner-provision.py create mode 100644 .github/scripts/aliyun-ecs-runner-teardown.py create mode 100644 .github/scripts/query-regression-nightly-refs.py create mode 100644 .github/workflows/query-regression-janitor.yml create mode 100644 .github/workflows/query-regression-nightly.yml create mode 100644 tests/perf/test_aliyun_ecs_runner_scripts.py create mode 100644 tests/perf/test_query_regression_nightly_refs.py diff --git a/.github/runner-scale-sets/query-regression/README.md b/.github/runner-scale-sets/query-regression/README.md index a1d58de7166..ce447dbdd0a 100644 --- a/.github/runner-scale-sets/query-regression/README.md +++ b/.github/runner-scale-sets/query-regression/README.md @@ -1,321 +1,184 @@ # Query regression self-hosted runners -The `Query Regression` workflow uses the dedicated ARC runner scale set -`perf-regression-8-cores`. ARC runner Pods run in the target Kubernetes cluster -and connect outbound to GitHub. The live scale set is currently **paused**: -`minRunners=0`, `maxRunners=0`, and no runner Pods. Do not resume it without -explicit approval. +The `Query Regression` workflow runs on **Aliyun ECS ephemeral runners** +(`aliyun-ecs`, the default and only automated path): a `provision` job on +`ubuntu-latest` creates one pay-as-you-go ECS instance per run, and the +instance registers itself as an ephemeral GitHub runner with a per-run +label. A `teardown` job (`if: always()`) deletes the instance; +`query-regression-janitor.yml` sweeps tagged leftovers older than 4 hours +daily. -## Prerequisites and trust admission +Build caches live on the instance's system disk, so **every run compiles +cold**; only the within-run reuse (base build warms the candidate build +through the shared target dir and sccache) applies. There is no retained +data disk. -Install the ARC scale set controller if it is not already installed: +Dispatching with any other `runner` value uses it as a literal self-hosted +runner label, which is how a manually prepared host (see +`ecs-image/bootstrap-runner-host.sh`) runs the workflow. For PR labels, set +the repository variable `QUERY_REGRESSION_PR_RUNNER` to such a label to +redirect PR runs away from ECS. + +The office ARC scale set `perf-regression-8-cores` that previously ran this +workflow is retired; see git history for its values files and pause/deploy +procedures. Tearing down the office cluster (helm release, runner namespace, +and the `query-regression-build-cache` PVC) is a manual operator action +outside this repository. This directory keeps its historical +`runner-scale-sets` name for path stability. + +## Nightly vs previous nightly + +`query-regression-nightly.yml` runs after a successful `GreptimeDB Nightly +Build` (`workflow_run`). It resolves that run's `head_sha` as the candidate +and the previous successful nightly (same branch, typically Friday when +Monday's nightly fires) as the base, then calls `query-regression.yml` with +those immutable SHAs. Both binaries are still compiled on the ECS runner; +this is not an artifact-download path. `workflow_dispatch` can pass explicit +`base_ref` / `candidate_ref` or a nightly run id. If there is no previous +nightly, or both nightlies built the same commit, the comparison is skipped. + +## Aliyun ECS path + +Configuration lives in repository variables/secrets: + +| Kind | Name | Purpose | +| --- | --- | --- | +| secret | `ALICLOUD_ECS_ACCESS_KEY_ID` / `ALICLOUD_ECS_ACCESS_KEY_SECRET` | RAM user scoped to ECS RunInstances/DeleteInstances/Describe*/CreateImage/RunCommand. Used only by provision/teardown jobs on `ubuntu-latest`; never reaches the ECS instance. | +| secret | `GH_PERSONAL_ACCESS_TOKEN` | Creates the short-lived runner registration token (shared with the jsonbench EC2 path). | +| vars | `ALIYUN_ECS_REGION_ID` / `ALIYUN_ECS_VSWITCH_ID` / `ALIYUN_ECS_SECURITY_GROUP_ID` | Network placement. The security group should allow egress only; no inbound rules are needed. The vSwitch pins the zone. | +| vars | `ALIYUN_ECS_INSTANCE_TYPE` | Dedicated (non-burstable, non-shared) instance family. Prefer 32 GiB (e.g. `ecs.g8i.2xlarge`); `ecs.c9i.2xlarge` is 8c16g and nightly thin-LTO of greptime peaks above that. Both base and candidate clusters run on the same machine, so noisy neighbors break thresholds. | +| vars | `QUERY_REGRESSION_ECS_IMAGE_ID` | Custom image built by `ecs-image/build-ecs-image.py`. | + +The system disk is 40 GiB, which covers the image, a 16 GiB swapfile, the +checkout, and cold build caches (target dir, cargo registry, sccache). ENOSPC +stops the runner itself from writing logs, which GitHub reports as `The +operation was canceled` with no telemetry, indistinguishable from a +platform-side cancellation. +cloud-init masks `systemd-oomd`, disables `unattended-upgrades` / +`apt-daily-upgrade`, creates `/swapfile`, and sets `OOMPolicy=continue` +on the runner unit. Ubuntu 24.04 defaults to `DefaultOOMPolicy=stop`, +which SIGTERM-s the whole unit when rustc is OOM-killed and GitHub +reports `The operation was canceled` with no telemetry. Unattended +upgrades can do the same via `systemctl restart` of the runner after a +library update; GitHub then records `UserCancelled` even though the job +was still valid. Swap is a safety net for 16 GiB types, not a substitute +for 32 GiB; linking on swap is slow. + +When a run dies with an unexplained `The operation was canceled` (no +"Canceled by" banner, healthy machine), re-dispatch with `keep_instance` +checked: the teardown job is skipped and the instance survives for +post-mortem inspection. The security group is egress-only, so inspect via +Cloud Assistant (`RunCommand`) or VNC: the runner's `_diag` logs under the +runner home record reconnects and worker crashes, `journalctl -u +ephemeral-github-runner.service` mirrors the console stream, `dmesg -T` +shows kernel OOM kills, and `machine-telemetry.log` in the job workspace +has the 30 s sampler history. The janitor sweep still deletes the instance +after its TTL, so finish the inspection within that window. + +Trust model on the ECS path: the instance receives only the one-hour runner +registration token via user data and holds no cloud credentials; the Aliyun +AK/SK exist only in the control-plane jobs. Instance tags +(`managed-by=query-regression-ci`, `query-regression-run-id`) feed the janitor +sweep. + +### Building and updating the ECS image + +The runner `Dockerfile` in the parent directory stays the single source of the +tool contract. Build a new ECS image from it: ```bash -helm upgrade --install arc \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set-controller \ - --namespace arc-systems \ - --create-namespace \ - --version 0.14.2 +ALIBABA_CLOUD_ACCESS_KEY_ID=... ALIBABA_CLOUD_ACCESS_KEY_SECRET=... \ +uv run .github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py \ + --region-id --vswitch-id --security-group-id \ + --base-image-id ``` -Create the GitHub App secret in the runner namespace. Prefer an App limited to -`GreptimeTeam/greptimedb`: +The script boots a temporary builder instance, `docker build`s the runner +image, materializes `/opt/rustup`, `/opt/cargo`, `/usr/local/bin` tools, and +`/home/runner` (actions-runner) onto the host, installs the ephemeral-runner +systemd unit from `ecs-image/`, and snapshots a custom image. It prints the +image id; set it as `QUERY_REGRESSION_ECS_IMAGE_ID`, and bump +`RUNNER_IMAGE_EPOCH` in `query-regression.yml` at the same time so the target +cache invalidates. Builder sentinel polling requires the Cloud Assistant +agent, which Aliyun public Ubuntu images include. -```bash -kubectl -n arc-runners create secret generic greptimedb-arc-github-app \ - --from-literal=github_app_id= \ - --from-literal=github_app_installation_id= \ - --from-file=github_app_private_key= -``` - -The values files here reference that secret by name. +## Trust admission for PR runs A maintainer applying the `query-regression` or `heavy-regression` label is -**trust admission for that exact PR revision**. `query-regression` runs the six -routine default cases; `heavy-regression` runs only the high-cardinality -`prom_remote_write_7913` remote-write case. The admitted job may use this scale -set's dedicated, writable persistent cache. `pull_request: labeled` is the only -PR trigger: the label event snapshots its merge, head, and base SHAs. A queued -job fetches that immutable event merge SHA directly, verifies it is a two-parent -merge whose parents include the snapshotted head exactly once, and uses its -other parent as the actual base build revision. The snapshotted event base is -retained for audit only, so a difference from the merge's non-head parent is -not a failure. The job never follows a newer mutable PR merge ref. An -unavailable event merge, or one that does not contain exactly one snapshotted -head parent, fails closed. A later PR head change does not retarget an already -queued run: it may execute only its previously trusted event revision if that -revision remains fetchable. To run the new revision, the maintainer must review -it, remove the label, and re-add the desired regression label; cancel the old -run if it is no longer wanted. An existing label does not automatically rerun -the benchmark. +**trust admission for that exact PR revision**. `query-regression` runs the +six routine default cases; `heavy-regression` runs only the high-cardinality +`prom_remote_write_7913` remote-write case. `pull_request: labeled` is the only PR +trigger: the label event snapshots its merge, head, and base SHAs. A queued +job fetches that immutable event merge SHA directly, verifies it is a +two-parent merge whose parents include the snapshotted head exactly once, and +uses its other parent as the actual base build revision. The snapshotted +event base is retained for audit only, so a difference from the merge's +non-head parent is not a failure. The job never follows a newer mutable PR +merge ref. An unavailable event merge, or one that does not contain exactly +one snapshotted head parent, fails closed. A later PR head change does not +retarget an already queued run: it may execute only its previously trusted +event revision if that revision remains fetchable. To run the new revision, +the maintainer must review it, remove the label, and re-add the desired +regression label; cancel the old run if it is no longer wanted. An existing +label does not automatically rerun the benchmark. -Admission does not relax runner hardening or GitHub permissions. Keep -service-account token mounting disabled; do not mount host paths, the Docker -socket, kubeconfig, or long-lived credentials. The runner and cache initializer -use UID/GID 1001, disallow privilege escalation, drop all capabilities, and use -the RuntimeDefault seccomp profile. Keep GitHub tokens least-privilege and -review workflow changes before admission. Where the CNI supports it, restrict -egress to required GitHub Actions, artifact/cache, Rust/crate/toolchain, DNS, -and image-registry endpoints; block unrelated cluster services, private ranges, -and metadata endpoints unless a case requires them. - -### Network routing prerequisite - -Required split routing is an **external environment-specific prerequisite**. The -responsible network operator must route GitHub Actions, GitHub content, -artifact/cache, crates.io, Rust toolchain, and image-registry traffic through -the approved path rather than the VPN where applicable. Neither this repository -nor Kubernetes configures that route. Verify it with the responsible network -operator before any canary. +Admission does not relax runner hardening or GitHub permissions. Keep the ECS +instance free of cloud credentials and long-lived tokens, keep the security +group egress-only, keep GitHub tokens least-privilege, and review workflow +changes before admission. ## Runner image and workflow tools -Build and push the derived runner image; it preserves the official -`/home/runner/run.sh` entrypoint and supplies CI tools needed at runtime. The -image builds `otelgen` from +The runner `Dockerfile` builds `otelgen` from [`WenyXu/otelgen`](https://github.com/WenyXu/otelgen) commit [`863a3f395d062c7322cc1de08a38774b7fdaa6c8`](https://github.com/WenyXu/otelgen/commit/863a3f395d062c7322cc1de08a38774b7fdaa6c8) -so trace cases do not download or compile tools during a benchmark run: +so trace cases do not download or compile tools during a benchmark run. It is +built only as a toolchain factory: `build-ecs-image.py` and +`bootstrap-runner-host.sh` both materialize its contents onto a host, and the +benchmark itself runs host-native. -```bash -docker build \ - --platform linux/amd64 \ - -f .github/runner-scale-sets/query-regression/Dockerfile \ - -t greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner:latest \ - .github/runner-scale-sets/query-regression - -docker push greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner:latest -``` - -Deploy by digest, not mutable tag, by updating both image references in -`values-8-cores.yaml` after a rebuild. Update `RUNNER_IMAGE_DIGEST` and bump -`RUNNER_IMAGE_EPOCH` in `query-regression.yml` at the same time. If the registry -is private, use a dedicated read-only pull secret only as `imagePullSecrets`; -never expose registry credentials to runner containers. -Both digest-pinned init and runner containers use `IfNotPresent`: the immutable -digest makes a cached image safe and avoids adding a registry dependency to every -runner startup. - -The runner optionally imports only the non-sensitive `HTTP_PROXY`, -`HTTPS_PROXY`, and `NO_PROXY` variables from the -`query-regression-runner-local-env` ConfigMap. Manage that ConfigMap locally in -the target namespace; private endpoint configuration must not be committed, and -credentials or secrets must never be placed in a ConfigMap. - -Before builds, the workflow asserts UID/GID 1001 and exact image tool versions: -`libprotoc 3.21.12`, `uv 0.11.26`, `mold 2.30.0`, `Python 3.12.3`, `sccache -0.16.0`, `otelgen` commit `863a3f395d062c7322cc1de08a38774b7fdaa6c8`, -root-owned `rustup 1.29.0`, and the image-baked -`nightly-2026-03-21` Rust toolchain. Rustup, Cargo, and Rustc must resolve from -`/opt/cargo/bin`; the runner cannot write `/opt/rustup` or `/opt/cargo/bin`. -Protobuf well-known includes, including `google/protobuf/any.proto` and -`google/protobuf/empty.proto`, are an image contract and must compile with -`protoc`. +Before builds, the workflow asserts the runner UID/GID (1001 in the ECS image; +overridable via `QUERY_REGRESSION_RUNNER_UID`/`QUERY_REGRESSION_RUNNER_GID`) +and exact tool versions: `libprotoc 3.21.12`, `uv 0.11.26`, `mold 2.40.4`, +`Python 3.14.4`, `sccache 0.16.0`, `otelgen` commit +`863a3f395d062c7322cc1de08a38774b7fdaa6c8`, root-owned `rustup 1.29.0`, and +the image-baked `nightly-2026-03-21` Rust toolchain. `mold` and `python3` +come from apt at image-build time (not Ubuntu 24.04's default 3.12); bump +the Verify pins together with `QUERY_REGRESSION_ECS_IMAGE_ID` when the +image is rebuilt. Rustup, Cargo, and Rustc +must resolve from `/opt/cargo/bin`; the runner cannot write `/opt/rustup` or +`/opt/cargo/bin`. Protobuf well-known includes, including +`google/protobuf/any.proto` and `google/protobuf/empty.proto`, are an image +contract and must compile with `protoc`. `actions-rust-lang/setup-rust-toolchain@v1` is intentionally removed. The -workflow sets its warning-denying mold `RUSTFLAGS` directly, disables automatic -Rustup installation, and performs no runtime toolchain downloads. +workflow sets its warning-denying mold `RUSTFLAGS` directly, disables +automatic Rustup installation, and performs no runtime toolchain downloads. -The workflow no longer uses GitHub `rust-cache`, `setup-protoc`, `setup-uv`, or -runtime Rust setup: the image establishes immutable executable state and the PVC -supplies only reusable Cargo data. Do not reintroduce those actions unless the -corresponding cache or image contract changes. +The workflow no longer uses GitHub `rust-cache`, `setup-protoc`, `setup-uv`, +or runtime Rust setup: the image establishes immutable executable state and +each instance's system disk holds only that run's Cargo data. Do not +reintroduce those actions unless the image contract changes. -## Capacity and persistent cache +## Capacity -`values-8-cores.yaml` is normal operation: `minRunners=0`, `maxRunners=1`. -`values-paused.yaml` is the mandatory pause overlay: `minRunners=0`, -`maxRunners=0`. The job uses group `query-regression-persistent-cache-v1`, -`queue: max`, and `cancel-in-progress: false`; admitted jobs queue rather than -replacing older pending jobs. During maintenance, cancel admitted queued runs as -well as pausing ARC. Runner Pods have `activeDeadlineSeconds=12600`. -The runner requests 6 CPU and limits at 8 CPU to preserve `minipc-3` -allocatable-capacity scheduling headroom; do not reset the request to 8 CPU -without revalidating scheduling capacity. +Each run provisions its own ECS instance, so overlapping dispatches proceed +in parallel. There is no workflow `concurrency` group. All Cargo state +(`CARGO_HOME` including registry/git, `CARGO_TARGET_DIR`, sccache, cache +metadata) lives on the 40 GiB system disk and is discarded with the VM. +`RUSTUP_HOME=/opt/rustup` and `/opt/cargo/bin` are image-owned. The runner +sets `RUSTC_WRAPPER=/usr/local/bin/sccache`, +`SCCACHE_DIR=/home/runner/.cache/sccache`, `SCCACHE_CACHE_SIZE=10G`, and +`CARGO_INCREMENTAL=0`. sccache uses its local disk backend and self-evicts +at 10G. Base and candidate builds share the target on that disk; Cargo +fingerprints invalidate source and dependency changes. -The cache claim `query-regression-build-cache` is a nominal 600Gi `local-path` -PVC in `arc-runners`. It is `ReadWriteOnce`; `local-path` uses -WaitForFirstConsumer binding and Delete reclaim behavior, produces a -node-affine local PV, is non-expandable, and the 600Gi request is not a hard -storage quota. The runner's `minipc-3` selector is its only consumer candidate. - -The initializer mounts the PVC root at `/cache`, creates and write-tests these -versioned subpaths as non-root UID/GID 1001, and the runner mounts them as: - -| Persistent state | PVC subpath | Runner mount | -| --- | --- | --- | -| Ephemeral Cargo home | `emptyDir` | `/home/runner/.cargo` | -| Cargo registry data | `cargo-registry-v1` | `/home/runner/.cargo/registry` | -| Cargo Git data | `cargo-git-v1` | `/home/runner/.cargo/git` | -| Cargo target | `query-regression-target-v1` | `/home/runner/query-regression-target` | -| Cache metadata | `meta-v1` | `/home/runner/query-regression-cache-meta` | -| sccache local disk cache | `sccache-v1` | `/home/runner/.cache/sccache` | -| Immutable Rust toolchain | image-owned | `/opt/rustup`, `/opt/cargo/bin` | - -The Pod security context uses UID/GID and `fsGroup` 1001 with -`fsGroupChangePolicy: OnRootMismatch`; no privileged `chown` or raw `hostPath` -is used. `CARGO_HOME` is a per-Pod `emptyDir`; only its nested `registry` and -`git` mounts are persistent. `RUSTUP_HOME=/opt/rustup` and `/opt/cargo/bin` are -image-owned immutable paths, while `CARGO_TARGET_DIR`, cache metadata, and -`SCCACHE_DIR` are persistent absolute paths. The runner sets -`RUSTC_WRAPPER=/usr/local/bin/sccache`, -`SCCACHE_DIR=/home/runner/.cache/sccache`, `SCCACHE_CACHE_SIZE=40G`, and -`CARGO_INCREMENTAL=0`. sccache uses its local PVC disk backend and self-evicts -at 40G; do not add runtime downloads, object storage, or a shared backend. - -The repository's `.cargo/config.toml` remains a trusted per-revision build input. -In contrast, `$CARGO_HOME/config*`, credentials, installed bins, and Cargo -metadata outside the persistent `registry` and `git` data mounts are ephemeral -and cannot survive to another Pod. - -The local disk backend has a one-server constraint. `maxRunners=1` and the -unchanged `query-regression-persistent-cache-v1` workflow concurrency group -serialize runs; do not increase runner capacity or relax that serialization -while this backend is in use. Base and candidate builds share the target; Cargo -fingerprints invalidate source and dependency changes. The workflow records the -sccache version and relevant environment in the target ABI marker, starts and -zeros sccache after cache and toolchain checks, shows initial/base/candidate -statistics, and resets statistics between base and candidate builds. - -### Disk preflight and cleanup contract - -Before applying or unpausing, verify the backing filesystem on `minipc-3` has -at least 900GiB free. The current local-path provisioner source is -`/opt/local-path-provisioner`; measure the filesystem containing it: - -```bash -df -PB1G /opt/local-path-provisioner -``` - -The workflow reports `du`, `df -P`, human-readable free space, and inode -availability before builds and in an always-run report. Its cleanup is narrow -and non-destructive: - -- warn at target size 400GiB; at 450GiB clear only the complete target root; -- warn at Cargo registry-plus-Git data size 60GiB; at 80GiB remove only - `registry/src` and `git/checkouts`, then abort if that persistent data remains - at least 80GiB; -- below 300GiB backing free space, clear the complete target root first, - remeasure, then remove only those Cargo extracted trees and checkouts; abort - if free space is still below 300GiB; -- never automatically remove Cargo registry cache/index, Git database, the - image-owned Cargo bin or Rustup toolchain, cache metadata, the self-evicting - sccache directory, or the PVC. - -The target clear uses fixed absolute roots and removes all entries, including -dotfiles. After migration validation, remove obsolete `cargo-home-v1` and -`rustup-home-v1` only in explicit maintenance while ARC is 0/0 and no runner Pod -exists; they are not mounted by the current configuration. - -## Deploy and pause safely - -First verify the configured external network route and the disk preflight. Apply the PVC; while -the scale set is paused, it is expected to remain `Pending` because -WaitForFirstConsumer has no scheduled runner: - -```bash -kubectl apply --dry-run=server \ - -f .github/runner-scale-sets/query-regression/cache-pvc.yaml -kubectl apply -f .github/runner-scale-sets/query-regression/cache-pvc.yaml -``` - -Render normal and paused configurations. Normal values are always first; the -pause overlay is always last: - -```bash -helm template perf-regression-8-cores \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ - --namespace arc-runners --version 0.14.2 \ - --set controllerServiceAccount.name=arc-gha-rs-controller \ - --set controllerServiceAccount.namespace=arc-systems \ - -f .github/runner-scale-sets/query-regression/values-8-cores.yaml - -helm template perf-regression-8-cores \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ - --namespace arc-runners --version 0.14.2 \ - --set controllerServiceAccount.name=arc-gha-rs-controller \ - --set controllerServiceAccount.namespace=arc-systems \ - -f .github/runner-scale-sets/query-regression/values-8-cores.yaml \ - -f .github/runner-scale-sets/query-regression/values-paused.yaml -``` - -The **first post-merge Helm deployment must reconcile the release in paused -mode**. Keep the pause overlay last: - -```bash -# First post-merge deployment and every return to paused mode: 0/0. -helm upgrade --install perf-regression-8-cores \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ - --namespace arc-runners --create-namespace --version 0.14.2 \ - --reset-values --wait \ - -f .github/runner-scale-sets/query-regression/values-8-cores.yaml \ - -f .github/runner-scale-sets/query-regression/values-paused.yaml - -# Expect 0/0 and no runner resources before considering normal mode. -kubectl -n arc-runners get autoscalingrunnerset perf-regression-8-cores \ - -o jsonpath='{.spec.minRunners}{"/"}{.spec.maxRunners}{"\n"}' -kubectl -n arc-runners get ephemeralrunners,pods \ - -l actions.github.com/scale-set-name=perf-regression-8-cores -``` - -Only after that verification and separate explicit approval, apply normal 0/1 -operation without the pause overlay: - -```bash -helm upgrade --install perf-regression-8-cores \ - oci://ghcr.io/actions/actions-runner-controller-charts/gha-runner-scale-set \ - --namespace arc-runners --create-namespace --version 0.14.2 \ - --reset-values --wait \ - -f .github/runner-scale-sets/query-regression/values-8-cores.yaml -``` - -Do not use bare `helm rollback`, `--atomic`, or `--reuse-values`: a stored -revision can restore nonzero runner capacity. Inspect rendered manifests for -capacity, the `minipc-3` selector, cache claim and mounts, initializer, -security context, and resources. After approved normal mode receives its first -canary, the PVC binds to `minipc-3`. - -For local-PV node loss, cache recovery is intentionally disposable: return to -0/0, recreate the PVC on a healthy node, cold-fill it, and run a new canary. - -## Canary and rollback - -With explicit approval, run two identical `workflow_dispatch` canaries on -`perf-regression-8-cores`, using immutable full base and candidate commit SHAs -and `cargo_profile=nightly`. The first is the cold fill; the second verifies warm -reuse. Record the workflow's base/candidate build elapsed logs and cache -ABI-marker output, initial/base/candidate sccache statistics, and cache report. -Confirm the image tool contract (root-owned Rustup/Cargo paths, baked nightly, -and non-writable `/opt` roots) and that the ephemeral Cargo home contains only -the mounted registry/Git data before Cargo creates per-Pod state. -Obtain dependency and tool network byte counters from the configured -environment counter source, filtered to `minipc-3` and the dependency/tool -destinations. - -Accept the canary only when all of the following hold: - -- exactly one runner Pod runs on `minipc-3`, and both jobs use the same bound PV; -- UID/GID 1001 cache mounts are writable; the warm run does not invalidate the - target or bulk-redownload crates or toolchains; sccache reports separate base - and candidate build statistics without server or cache-path errors; immutable - Rustup/Cargo roots remain non-writable and only registry/Git data persists; -- warm base build time is at most 50% of cold base build time; -- warm dependency/tool network bytes are at most 10% of cold fill bytes; -- cache sizes remain below soft watermarks, node free space remains at least - 300GiB, and the benchmark is correct without TLS EOFs or timeouts; -- the configured environment counter source confirms this traffic is outside VPN - accounting. - -Immediately return to 0/0 after either canary unless ongoing normal operation -has been explicitly approved; return immediately on any traffic, cache, disk, -TLS, or correctness failure. To roll back, use the paused Helm upgrade above, -or another explicit `helm upgrade --install` with known-good values followed by -`values-paused.yaml`, `--reset-values`, and `--wait`. Do not delete the PVC -automatically; preserve it for diagnosis unless intentionally discarding cache. +The workflow reports `du` / `df` in telemetry. It does not try to reclaim +space across runs: a cold 40 GiB disk that fills up fails the build. ## Future optional phases -The current phase uses a digest-pinned image with sccache 0.16.0 and no shared -cache service. Optional follow-ups are an image additionally seeded with the -exact Rust toolchain and `cargo fetch --locked`; or an internal read/write -sccache backend or Cargo/Git mirror. Evaluate them only if persistent PVC reuse -is insufficient. +The current phase uses a materialized runner toolchain with sccache 0.16.0 and +no shared cache service. Optional follow-ups are an image additionally seeded +with `cargo fetch --locked` results, or an internal read/write sccache +backend or Cargo/Git mirror. Evaluate them only if cold compile time on the +system disk becomes the bottleneck. diff --git a/.github/runner-scale-sets/query-regression/cache-pvc.yaml b/.github/runner-scale-sets/query-regression/cache-pvc.yaml deleted file mode 100644 index 692f3a2d6e2..00000000000 --- a/.github/runner-scale-sets/query-regression/cache-pvc.yaml +++ /dev/null @@ -1,15 +0,0 @@ -apiVersion: v1 -kind: PersistentVolumeClaim -metadata: - name: query-regression-build-cache - namespace: arc-runners - labels: - app.kubernetes.io/part-of: perf-regression-8-cores -spec: - volumeMode: Filesystem - accessModes: - - ReadWriteOnce - storageClassName: local-path - resources: - requests: - storage: 600Gi diff --git a/.github/runner-scale-sets/query-regression/ecs-image/bootstrap-runner-host.sh b/.github/runner-scale-sets/query-regression/ecs-image/bootstrap-runner-host.sh new file mode 100644 index 00000000000..883bbede092 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/ecs-image/bootstrap-runner-host.sh @@ -0,0 +1,195 @@ +#!/usr/bin/env bash +# Copyright 2023 Greptime Team +# +# 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. + +# Bootstrap an Ubuntu 24.04 host as a query-regression self-hosted runner. +# +# This is the host-native equivalent of the ECS image contract: it +# installs docker-ce from Docker's official repository, builds the runner +# image from the sibling Dockerfile (the single source of the tool contract), +# materializes the toolchain onto the host, and optionally registers the +# runner and installs its systemd service. +# +# Large state lives under --data-root (a data disk): docker's data-root, the +# runner home, and the rustup/cargo toolchain roots, symlinked back to the +# hard-coded contract paths (/home/runner, /opt/rustup, /opt/cargo). +# +# The script is idempotent: re-running refreshes the toolchain in place, and +# the runner registration (.runner/.credentials) and _work survive. A live +# runner service is stopped first and restarted at the end. +# +# Usage: +# sudo bash bootstrap-runner-host.sh --repo-dir /path/to/greptimedb +# sudo RUNNER_TOKEN= bash bootstrap-runner-host.sh --repo-dir . \ +# --register --repo-url https://github.com// +# +# Options: +# --repo-dir PATH greptimedb checkout containing the runner Dockerfile (required) +# --data-root PATH data disk mount point (default: /data) +# --uid / --gid N runner user id (default: 3141; the ECS image contract is 1001) +# --register also register the runner and install the systemd service +# --repo-url URL repository URL for registration (required with --register) +# --runner-name N runner name (default: qreg-host) +# --labels L runner labels (default: perf-regression-8-cores) +# +# With --register, provide a fresh registration token via RUNNER_TOKEN. + +set -euo pipefail + +REPO_DIR="" +DATA_ROOT="/data" +RUNNER_UID="3141" +RUNNER_GID="3141" +REGISTER="false" +REPO_URL="" +RUNNER_NAME="qreg-host" +RUNNER_LABELS="perf-regression-8-cores" +IMAGE_TAG="qreg-runner:manual" + +while [[ $# -gt 0 ]]; do + case "$1" in + --repo-dir) REPO_DIR="$2"; shift 2 ;; + --data-root) DATA_ROOT="$2"; shift 2 ;; + --uid) RUNNER_UID="$2"; shift 2 ;; + --gid) RUNNER_GID="$2"; shift 2 ;; + --register) REGISTER="true"; shift ;; + --repo-url) REPO_URL="$2"; shift 2 ;; + --runner-name) RUNNER_NAME="$2"; shift 2 ;; + --labels) RUNNER_LABELS="$2"; shift 2 ;; + -h | --help) sed -n '17,40p' "$0"; exit 0 ;; + *) echo "Unknown argument: $1" >&2; exit 2 ;; + esac +done + +[[ $EUID -eq 0 ]] || { echo "Run as root, e.g. sudo bash $0 ..." >&2; exit 1; } +[[ -n "${REPO_DIR}" ]] || { echo "--repo-dir is required" >&2; exit 2; } +DOCKERFILE="${REPO_DIR}/.github/runner-scale-sets/query-regression/Dockerfile" +[[ -f "${DOCKERFILE}" ]] || { echo "Dockerfile not found at ${DOCKERFILE}" >&2; exit 1; } +if [[ "${REGISTER}" == "true" ]]; then + [[ -n "${REPO_URL}" ]] || { echo "--repo-url is required with --register" >&2; exit 2; } + [[ -n "${RUNNER_TOKEN:-}" ]] || { echo "RUNNER_TOKEN env is required with --register" >&2; exit 2; } +fi + +step() { printf '\n==> %s\n' "$*"; } + +# A re-run refreshes files under the (possibly live) runner service, so stop +# it first and restart it at the end (unless --register re-creates it). +RUNNER_SERVICE="$(systemctl list-units --type=service --all --no-legend 'actions.runner.*' 2>/dev/null | awk '{print $1}' | head -n 1 || true)" +RUNNER_WAS_ACTIVE="false" +if [[ -n "${RUNNER_SERVICE}" ]] && systemctl is-active --quiet "${RUNNER_SERVICE}"; then + step "Stop running runner service ${RUNNER_SERVICE}" + systemctl stop "${RUNNER_SERVICE}" + RUNNER_WAS_ACTIVE="true" +fi + +step "Install docker-ce from Docker's official repository" +for pkg in docker.io docker-doc docker-compose docker-compose-v2 podman-docker containerd runc; do + apt-get remove -y "${pkg}" 2>/dev/null || true +done +apt-get update +apt-get install -y ca-certificates curl +install -m 0755 -d /etc/apt/keyrings +curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc +chmod a+r /etc/apt/keyrings/docker.asc +# shellcheck disable=SC1091 +. /etc/os-release +tee /etc/apt/sources.list.d/docker.sources > /dev/null <&2 + exit 1 +fi +echo "{ \"data-root\": \"${DATA_ROOT}/docker\" }" > /etc/docker/daemon.json +# containerd keeps its own content store (where build-time image data lands); +# move it off the system disk too. +mkdir -p /etc/containerd +containerd config default | sed "s|root = \"/var/lib/containerd\"|root = \"${DATA_ROOT}/containerd\"|" > /etc/containerd/config.toml +systemctl start containerd docker +docker info | grep "Docker Root Dir" + +step "Build the runner image" +docker build --platform linux/amd64 -f "${DOCKERFILE}" -t "${IMAGE_TAG}" "${REPO_DIR}" + +step "Materialize the toolchain onto the host" +mkdir -p "${DATA_ROOT}/opt" "${DATA_ROOT}/runner" /opt /home +container="$(docker create "${IMAGE_TAG}")" +trap 'docker rm -f "${container}" >/dev/null 2>&1 || true' EXIT +# The `/.` suffix copies directory *contents*, so re-running over an existing +# materialization refreshes it in place instead of nesting runner/runner. +# Files absent from the image (runner registration, _work) are preserved. +docker cp "${container}:/home/runner/." "${DATA_ROOT}/runner" +docker cp "${container}:/opt/rustup/." "${DATA_ROOT}/opt/rustup" +docker cp "${container}:/opt/cargo/." "${DATA_ROOT}/opt/cargo" +for tool in uv uvx otelgen sccache; do + docker cp "${container}:/usr/local/bin/${tool}" "/usr/local/bin/${tool}" +done +docker rm "${container}" > /dev/null +trap - EXIT +ln -sfn "${DATA_ROOT}/runner" /home/runner +ln -sfn "${DATA_ROOT}/opt/rustup" /opt/rustup +ln -sfn "${DATA_ROOT}/opt/cargo" /opt/cargo +docker image rm "${IMAGE_TAG}" > /dev/null + +step "Install system packages" +apt-get install -y --no-install-recommends \ + build-essential clang cmake git gzip jq libprotobuf-dev libssl-dev mold \ + openssh-client pkg-config protobuf-compiler python3 sudo tar unzip wget xz-utils zip zstd + +step "Create runner user (${RUNNER_UID}:${RUNNER_GID}) and environment" +getent group "${RUNNER_GID}" > /dev/null || groupadd -g "${RUNNER_GID}" runner +id -u runner > /dev/null 2>&1 || useradd -u "${RUNNER_UID}" -g "${RUNNER_GID}" -d /home/runner -s /bin/bash runner +chown -R "${RUNNER_UID}:${RUNNER_GID}" "${DATA_ROOT}/runner" +echo 'PATH=/opt/cargo/bin:/usr/local/bin:/usr/bin:/bin' > "${DATA_ROOT}/runner/.env" +chown "${RUNNER_UID}:${RUNNER_GID}" "${DATA_ROOT}/runner/.env" + +if [[ "${REGISTER}" == "true" ]]; then + step "Register runner ${RUNNER_NAME} (labels: ${RUNNER_LABELS})" + if [[ -n "${RUNNER_SERVICE}" ]]; then + (cd /home/runner && ./svc.sh uninstall) || true + fi + runuser -u runner -- bash -c "cd /home/runner && HOME=/home/runner ./config.sh \ + --url '${REPO_URL}' --token '${RUNNER_TOKEN}' --name '${RUNNER_NAME}' \ + --labels '${RUNNER_LABELS}' --unattended --replace --disableupdate" + + step "Install and start the runner service" + cd /home/runner + ./svc.sh install runner + ./svc.sh start + ./svc.sh status +else + step "Skipping registration (pass --register --repo-url ... with RUNNER_TOKEN to enable)" + if [[ "${RUNNER_WAS_ACTIVE}" == "true" ]]; then + step "Restart runner service ${RUNNER_SERVICE}" + systemctl start "${RUNNER_SERVICE}" + systemctl --no-pager status "${RUNNER_SERVICE}" || true + fi +fi + +step "Done" +echo "Runner home: /home/runner -> ${DATA_ROOT}/runner" +echo "Before each workflow run on this persistent host, clean transient cargo state:" +echo " sudo rm -f /home/runner/.cargo/.package-cache" +echo " sudo rm -rf /home/runner/.cargo/.global-cache" diff --git a/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py b/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py new file mode 100644 index 00000000000..7031917f910 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py @@ -0,0 +1,392 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# 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. + +# PEP 723 inline metadata (see .github/scripts/aliyun-ecs-runner-provision.py +# for the convention). +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "alibabacloud_ecs20140526>=4.1.0,<6", +# "alibabacloud_tea_openapi>=0.3.12,<1", +# ] +# /// + +"""Build the query-regression ECS custom image (manual ops tool). + +Boots a temporary pay-as-you-go ECS instance from a public Ubuntu 24.04 image, +builds the existing runner container image (the Dockerfile in the parent +directory remains the single source of the tool contract), materializes the +tool directories onto the host filesystem so the workflow's "Verify runner +image tools" step holds unchanged, installs the ephemeral-runner systemd unit, +and snapshots the result as a custom image. The temporary instance is deleted +afterwards. + +Sentinel polling reads the instance's serial console output +(GetInstanceConsoleOutput) and looks for marker lines the user data writes to +/dev/console. This has no in-guest agent dependency. + +Usage: + uv run .github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py \ + --region-id cn-hangzhou --vswitch-id vsw-... --security-group-id sg-... \ + --base-image-id ubuntu_24_04_x64_20G_alibase_*.vhd +""" + +from __future__ import annotations + +import argparse +import base64 +import os +import time +from pathlib import Path + +ASSETS_DIR = Path(__file__).resolve().parent +DONE_MARKER = "QREG_IMAGE_BUILD_DONE" +FAILED_MARKER = "QREG_IMAGE_BUILD_FAILED" +POLL_INTERVAL_SECONDS = 15 +CONSOLE_POLL_INTERVAL_SECONDS = 30 +BUILD_TIMEOUT_SECONDS = 60 * 60 + +# Same apt package contract as the runner Dockerfile; the base +# actions-runner image is Ubuntu 24.04, so an Ubuntu 24.04 host resolves the +# same tool versions (protoc 3.21.12, mold 2.40.4, Python 3.14.4). +# Docker itself comes from Docker's official repository (docker-ce), not the +# distribution-packaged docker.io. +APT_PACKAGES = [ + "build-essential", + "ca-certificates", + "clang", + "cmake", + "curl", + "git", + "gpg", + "gzip", + "jq", + "libprotobuf-dev", + "libssl-dev", + "mold", + "openssh-client", + "pkg-config", + "protobuf-compiler", + "python3", + "sudo", + "tar", + "unzip", + "wget", + "xz-utils", + "zip", + "zstd", +] + +DOCKER_CE_PACKAGES = "docker-ce docker-ce-cli containerd.io docker-buildx-plugin" + + +def render_user_data(dockerfile: str, start_runner: str, unit: str) -> str: + dockerfile_b64 = base64.b64encode(dockerfile.encode()).decode() + start_runner_b64 = base64.b64encode(start_runner.encode()).decode() + unit_b64 = base64.b64encode(unit.encode()).decode() + packages = " ".join(APT_PACKAGES) + return f"""#!/bin/bash +set -euo pipefail +trap 'echo "{FAILED_MARKER} at line $LINENO" > /dev/console' ERR +# Stream the full build log to the serial console so the poller (and anyone +# watching GetInstanceConsoleOutput) sees real progress, not a silent login +# prompt for the whole build. +exec > >(tee -a /dev/console) 2>&1 + +apt-get update +DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends {packages} + +# Docker from the official repository, not the distribution-packaged docker.io. +install -m 0755 -d /etc/apt/keyrings +curl -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /etc/apt/keyrings/docker.gpg +chmod a+r /etc/apt/keyrings/docker.gpg +. /etc/os-release +echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu ${{VERSION_CODENAME}} stable" \ + > /etc/apt/sources.list.d/docker.list +apt-get update +DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends {DOCKER_CE_PACKAGES} + +base64 -d > /tmp/Dockerfile <<'EOF' +{dockerfile_b64} +EOF +mkdir -p /tmp/image-context +docker build --platform linux/amd64 -f /tmp/Dockerfile -t qreg-runner:local /tmp/image-context + +# Materialize the tool contract onto the host filesystem. +container="$(docker create qreg-runner:local)" +trap 'docker rm -f "${{container}}" >/dev/null 2>&1 || true' EXIT +docker cp "${{container}}:/opt/rustup" /opt/rustup +docker cp "${{container}}:/opt/cargo" /opt/cargo +for tool in uv uvx otelgen sccache; do + docker cp "${{container}}:/usr/local/bin/${{tool}}" "/usr/local/bin/${{tool}}" +done +docker cp "${{container}}:/home/runner" /home/runner + +# Runner identity mirrors the container image (UID/GID 1001). +groupadd --gid 1001 runner 2>/dev/null || true +useradd --uid 1001 --gid 1001 --home-dir /home/runner --shell /bin/bash runner 2>/dev/null || true +chown -R 1001:1001 /home/runner + +mkdir -p /opt/ephemeral-github-runner +base64 -d > /opt/ephemeral-github-runner/start-runner.sh <<'EOF' +{start_runner_b64} +EOF +chmod 0755 /opt/ephemeral-github-runner/start-runner.sh +base64 -d > /etc/systemd/system/ephemeral-github-runner.service <<'EOF' +{unit_b64} +EOF +systemctl daemon-reload +systemctl enable ephemeral-github-runner.service + +# Keep the image free of the build-time docker state (also shrinks the +# snapshot: buildkit cache is several GB). +docker rm -f "${{container}}" >/dev/null +docker system prune -af >/dev/null +trap - EXIT + +echo "{DONE_MARKER}" > /dev/console +""" + + +def make_ecs_client(region_id: str): + from alibabacloud_ecs20140526.client import Client as EcsClient + from alibabacloud_tea_openapi.models import Config as OpenApiConfig + + return EcsClient( + OpenApiConfig( + access_key_id=os.environ["ALIBABA_CLOUD_ACCESS_KEY_ID"], + access_key_secret=os.environ["ALIBABA_CLOUD_ACCESS_KEY_SECRET"], + region_id=region_id, + endpoint=f"ecs.{region_id}.aliyuncs.com", + ) + ) + + +def wait_for_instance_status(client, region_id: str, instance_id: str, wanted: str, deadline: float) -> None: + from alibabacloud_ecs20140526 import models as ecs_models + + import json + + while time.monotonic() < deadline: + response = call_api_with_retry( + lambda: client.describe_instances( + ecs_models.DescribeInstancesRequest( + region_id=region_id, instance_ids=json.dumps([instance_id]) + ) + ), + "DescribeInstances", + ) + instances = response.body.instances.instance + if instances and instances[0].status == wanted: + return + time.sleep(POLL_INTERVAL_SECONDS) + raise TimeoutError(f"Instance {instance_id} did not reach status {wanted} in time") + + +# Aliyun error codes worth retrying: throttling and server-side faults. Client +# errors (4xx: permissions, bad parameters) are configuration problems and must +# fail fast instead of being retried. +TRANSIENT_ERROR_CODES = {"Throttling", "Throttling.User", "InternalError", "ServiceUnavailable"} + + +def is_transient(error: Exception) -> bool: + # UnretryableException comes from the darabonba network layer: connection + # reset, read timeout, DNS blip. + if type(error).__name__ == "UnretryableException": + return True + code = getattr(error, "code", "") or "" + status = getattr(error, "statusCode", None) + return code in TRANSIENT_ERROR_CODES or (isinstance(status, int) and status >= 500) + + +def call_api_with_retry(fn, description: str, attempts: int = 5): + """Retry transient API/network failures; the build is too long to die on a blip. + + Only call this with idempotent or safely-repeatable requests (reads, + StopInstance, CreateImage). Never with RunInstances: if the request + succeeded but the response was lost, a retry double-creates instances. + """ + for attempt in range(1, attempts + 1): + try: + return fn() + except Exception as error: # noqa: BLE001 + if attempt == attempts or not is_transient(error): + raise + print(f"{description} failed (attempt {attempt}/{attempts}): {error}", flush=True) + time.sleep(POLL_INTERVAL_SECONDS) + raise RuntimeError("unreachable: retry loop exited without returning") + + +def read_console_output(client, region_id: str, instance_id: str) -> str: + """Fetch the instance's serial console output; no in-guest agent needed.""" + from alibabacloud_ecs20140526 import models as ecs_models + + response = call_api_with_retry( + lambda: client.get_instance_console_output( + ecs_models.GetInstanceConsoleOutputRequest(region_id=region_id, instance_id=instance_id) + ), + "GetInstanceConsoleOutput", + ) + return base64.b64decode(response.body.console_output or "").decode("utf-8", "replace") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--region-id", default=os.environ.get("ALIYUN_ECS_REGION_ID")) + parser.add_argument("--vswitch-id", default=os.environ.get("ALIYUN_ECS_VSWITCH_ID")) + parser.add_argument("--security-group-id", default=os.environ.get("ALIYUN_ECS_SECURITY_GROUP_ID")) + parser.add_argument("--base-image-id", default=os.environ.get("ALIYUN_ECS_BASE_IMAGE_ID")) + parser.add_argument("--resource-group-id", default=os.environ.get("ALIYUN_ECS_RESOURCE_GROUP_ID")) + parser.add_argument("--instance-type", default="ecs.g7.xlarge") + parser.add_argument("--image-name", default=None, help="Defaults to a timestamped name.") + args = parser.parse_args() + + for name in ("region_id", "vswitch_id", "security_group_id", "base_image_id"): + if not getattr(args, name): + raise SystemExit(f"Missing required configuration: --{name.replace('_', '-')}") + + from alibabacloud_ecs20140526 import models as ecs_models + + client = make_ecs_client(args.region_id) + image_name = args.image_name or time.strftime( + "greptimedb-query-regression-runner-%Y%m%d%H%M%S", time.gmtime() + ) + + user_data = base64.b64encode( + render_user_data( + (ASSETS_DIR.parent / "Dockerfile").read_text(), + (ASSETS_DIR / "start-runner.sh").read_text(), + (ASSETS_DIR / "ephemeral-github-runner.service").read_text(), + ).encode() + ).decode() + + instance_id = None + try: + response = client.run_instances( + ecs_models.RunInstancesRequest( + region_id=args.region_id, + image_id=args.base_image_id, + resource_group_id=args.resource_group_id, + instance_type=args.instance_type, + v_switch_id=args.vswitch_id, + security_group_id=args.security_group_id, + instance_name=f"build-{image_name}", + description="Temporary builder for the query-regression ECS image", + amount=1, + instance_charge_type="PostPaid", + spot_strategy="NoSpot", + internet_charge_type="PayByTraffic", + internet_max_bandwidth_out=100, + # Peak usage is ~18G (OS+apt, docker image, and the materialized + # toolchain coexist briefly): deliberately tight, and a smaller + # disk makes the image snapshot faster. Instances created from + # the image get a larger system disk from the provision side. + system_disk=ecs_models.RunInstancesRequestSystemDisk( + category="cloud_essd", size="20" + ), + user_data=user_data, + tag=[ + ecs_models.RunInstancesRequestTag(key="managed-by", value="query-regression-ci"), + ecs_models.RunInstancesRequestTag(key="role", value="image-builder"), + ], + ) + ) + instance_id = response.body.instance_id_sets.instance_id_set[0] + print(f"Builder instance: {instance_id}", flush=True) + wait_for_instance_status(client, args.region_id, instance_id, "Running", time.monotonic() + 10 * 60) + + deadline = time.monotonic() + BUILD_TIMEOUT_SECONDS + while time.monotonic() < deadline: + console = read_console_output(client, args.region_id, instance_id) + if DONE_MARKER in console: + break + if FAILED_MARKER in console: + tail = "\n".join(console.splitlines()[-20:]) + raise RuntimeError(f"Image build failed on the builder; console tail:\n{tail}") + lines = console.splitlines() + print(f"Build in progress; last console line: {lines[-1] if lines else '(none yet)'}", flush=True) + time.sleep(CONSOLE_POLL_INTERVAL_SECONDS) + else: + raise TimeoutError("Image build did not finish in time") + + print("Stopping builder before image creation", flush=True) + try: + call_api_with_retry( + lambda: client.stop_instance( + ecs_models.StopInstanceRequest( + instance_id=instance_id, + # The instance is deleted right after the snapshot, but + # there is no reason to keep billing vCPU during the stop. + stopped_mode="StopCharging", + ) + ), + "StopInstance", + ) + except Exception as error: # noqa: BLE001 + # Instance families with local disks do not support StopCharging. + print(f"StopCharging unavailable ({error}); stopping with default mode", flush=True) + call_api_with_retry( + lambda: client.stop_instance(ecs_models.StopInstanceRequest(instance_id=instance_id)), + "StopInstance", + ) + wait_for_instance_status(client, args.region_id, instance_id, "Stopped", time.monotonic() + 10 * 60) + + image = call_api_with_retry( + lambda: client.create_image( + ecs_models.CreateImageRequest( + region_id=args.region_id, instance_id=instance_id, image_name=image_name + ) + ), + "CreateImage", + ) + image_id = image.body.image_id + deadline = time.monotonic() + 60 * 60 + while time.monotonic() < deadline: + description = call_api_with_retry( + lambda: client.describe_images( + ecs_models.DescribeImagesRequest(region_id=args.region_id, image_id=image_id) + ), + "DescribeImages", + ) + images = description.body.images.image + if images: + status = images[0].status + print(f"Image {image_id} status: {status} (progress {images[0].progress})", flush=True) + if status == "Available": + break + time.sleep(POLL_INTERVAL_SECONDS) + else: + raise TimeoutError( + f"Image {image_id} did not become Available in time. The snapshot " + f"continues server-side: check its status in the console and reuse it " + f"once Available instead of rebuilding." + ) + + print(f"Custom image ready: {image_id} ({image_name})", flush=True) + print(f"Set the repo variable QUERY_REGRESSION_ECS_IMAGE_ID={image_id}", flush=True) + return 0 + finally: + if instance_id: + try: + client.delete_instance( + ecs_models.DeleteInstanceRequest(instance_id=instance_id, force=True) + ) + print(f"Deleted builder instance {instance_id}", flush=True) + except Exception as error: # noqa: BLE001 + print(f"Failed to delete builder instance {instance_id}: {error}", flush=True) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/runner-scale-sets/query-regression/ecs-image/ephemeral-github-runner.service b/.github/runner-scale-sets/query-regression/ecs-image/ephemeral-github-runner.service new file mode 100644 index 00000000000..a1cbe6a4503 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/ecs-image/ephemeral-github-runner.service @@ -0,0 +1,45 @@ +# Copyright 2023 Greptime Team +# +# 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. + +[Unit] +Description=Ephemeral GitHub Actions runner +After=network-online.target +Wants=network-online.target +# cloud-final runs the provision user data, which writes +# /etc/ephemeral-github-runner.env and the drop-ins below. Ordering after it +# makes the runner come up configured on first boot, no restart needed. +# (cloud-init's `systemctl restart` in user data remains as the bridge for +# images built before this ordering existed.) +After=cloud-final.service + +[Service] +Type=simple +# The toolchain lives in the image at /opt/cargo/bin (the Dockerfile's ENV +# PATH does not survive materialization onto the host). Publish it here so +# every runner job inherits it. /etc/ephemeral-github-runner.env currently +# carries the same PATH line as a bridge for images built before this +# directive existed; EnvironmentFile is applied after Environment=, and the +# two values are kept identical, so there is no conflict. +Environment=PATH=/opt/cargo/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin +EnvironmentFile=/etc/ephemeral-github-runner.env +ExecStart=/opt/ephemeral-github-runner/start-runner.sh +Restart=no +# If the kernel OOM killer kills a job child (rustc/cargo), keep the runner +# service alive so the job is reported as a normal failure with logs instead +# of the runner vanishing and GitHub reporting a cancellation. cloud-init's +# oom.conf drop-in carries the same directive for images built before this. +OOMPolicy=continue + +[Install] +WantedBy=multi-user.target diff --git a/.github/runner-scale-sets/query-regression/ecs-image/start-runner.sh b/.github/runner-scale-sets/query-regression/ecs-image/start-runner.sh new file mode 100644 index 00000000000..a6aff87ca22 --- /dev/null +++ b/.github/runner-scale-sets/query-regression/ecs-image/start-runner.sh @@ -0,0 +1,37 @@ +#!/bin/bash +# Copyright 2023 Greptime Team +# +# 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. + +# Configure and start the ephemeral GitHub Actions runner on an ECS instance. +# Baked into the custom image at /opt/ephemeral-github-runner/ by +# build-ecs-image.py; the per-run environment is written by cloud-init to +# /etc/ephemeral-github-runner.env. +set -euo pipefail + +# shellcheck disable=SC1091 +source /etc/ephemeral-github-runner.env + +export HOME=/home/runner +cd /home/runner + +if [[ ! -f /home/runner/.runner ]]; then + runuser -u runner -- ./config.sh \ + --url "${REPO_URL}" \ + --token "${RUNNER_TOKEN}" \ + --name "${RUNNER_NAME}" \ + --labels "${RUNNER_LABELS}" \ + --ephemeral --unattended --replace --disableupdate +fi + +exec runuser -u runner -- ./run.sh diff --git a/.github/runner-scale-sets/query-regression/values-8-cores.yaml b/.github/runner-scale-sets/query-regression/values-8-cores.yaml deleted file mode 100644 index cff1322680b..00000000000 --- a/.github/runner-scale-sets/query-regression/values-8-cores.yaml +++ /dev/null @@ -1,146 +0,0 @@ -githubConfigUrl: "https://github.com/GreptimeTeam/greptimedb" -githubConfigSecret: greptimedb-arc-github-app - -runnerScaleSetName: "perf-regression-8-cores" -minRunners: 0 -maxRunners: 1 - -template: - spec: - automountServiceAccountToken: false - activeDeadlineSeconds: 12600 - nodeSelector: - kubernetes.io/hostname: minipc-3 - securityContext: - runAsNonRoot: true - runAsUser: 1001 - runAsGroup: 1001 - fsGroup: 1001 - fsGroupChangePolicy: OnRootMismatch - seccompProfile: - type: RuntimeDefault - volumes: - - name: cargo-home - emptyDir: {} - - name: build-cache - persistentVolumeClaim: - claimName: query-regression-build-cache - initContainers: - - name: initialize-build-cache - image: greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner@sha256:e713b294e23b7e15184e558866c90025e59930033e72c97650dbc7f1ca022d11 - imagePullPolicy: IfNotPresent - command: - - /bin/sh - - -ec - - | - umask 0002 - for directory in cargo-registry-v1 cargo-git-v1 query-regression-target-v1 meta-v1 sccache-v1; do - cache_directory="/cache/${directory}" - mkdir -p "${cache_directory}" - test -w "${cache_directory}" - probe_file="${cache_directory}/.write-probe" - : > "${probe_file}" - test -f "${probe_file}" - rm "${probe_file}" - done - securityContext: - runAsNonRoot: true - runAsUser: 1001 - runAsGroup: 1001 - allowPrivilegeEscalation: false - capabilities: - drop: - - ALL - resources: - requests: - cpu: 100m - memory: 128Mi - ephemeral-storage: 1Gi - limits: - cpu: 500m - memory: 512Mi - ephemeral-storage: 1Gi - volumeMounts: - - name: build-cache - mountPath: /cache - containers: - - name: runner - image: greptime-registry.cn-hangzhou.cr.aliyuncs.com/greptime/greptimedb-query-regression-runner@sha256:e713b294e23b7e15184e558866c90025e59930033e72c97650dbc7f1ca022d11 - imagePullPolicy: IfNotPresent - command: ["/home/runner/run.sh"] - env: - - name: CARGO_HOME - value: /home/runner/.cargo - - name: HTTP_PROXY - valueFrom: - configMapKeyRef: - name: query-regression-runner-local-env - key: HTTP_PROXY - optional: true - - name: HTTPS_PROXY - valueFrom: - configMapKeyRef: - name: query-regression-runner-local-env - key: HTTPS_PROXY - optional: true - - name: NO_PROXY - valueFrom: - configMapKeyRef: - name: query-regression-runner-local-env - key: NO_PROXY - optional: true - - name: UV_CACHE_DIR - value: /home/runner/.cargo/uv-cache - - name: RUSTUP_HOME - value: /opt/rustup - - name: RUSTUP_TOOLCHAIN - value: nightly-2026-03-21 - - name: RUSTUP_AUTO_INSTALL - value: "0" - - name: CARGO_TARGET_DIR - value: /home/runner/query-regression-target - - name: QUERY_REGRESSION_CACHE_META - value: /home/runner/query-regression-cache-meta - - name: RUSTC_WRAPPER - value: /usr/local/bin/sccache - - name: SCCACHE_DIR - value: /home/runner/.cache/sccache - - name: SCCACHE_CACHE_SIZE - value: 40G - - name: CARGO_INCREMENTAL - value: "0" - securityContext: - runAsNonRoot: true - runAsUser: 1001 - runAsGroup: 1001 - allowPrivilegeEscalation: false - capabilities: - drop: - - ALL - resources: - requests: - cpu: "6" - memory: 24Gi - ephemeral-storage: 80Gi - limits: - cpu: "8" - memory: 24Gi - ephemeral-storage: 80Gi - volumeMounts: - - name: cargo-home - mountPath: /home/runner/.cargo - - name: build-cache - mountPath: /home/runner/.cargo/registry - subPath: cargo-registry-v1 - - name: build-cache - mountPath: /home/runner/.cargo/git - subPath: cargo-git-v1 - - name: build-cache - mountPath: /home/runner/query-regression-target - subPath: query-regression-target-v1 - - name: build-cache - mountPath: /home/runner/query-regression-cache-meta - subPath: meta-v1 - - name: build-cache - mountPath: /home/runner/.cache/sccache - subPath: sccache-v1 diff --git a/.github/runner-scale-sets/query-regression/values-paused.yaml b/.github/runner-scale-sets/query-regression/values-paused.yaml deleted file mode 100644 index fe734ef7615..00000000000 --- a/.github/runner-scale-sets/query-regression/values-paused.yaml +++ /dev/null @@ -1,2 +0,0 @@ -minRunners: 0 -maxRunners: 0 diff --git a/.github/scripts/aliyun-ecs-runner-provision.py b/.github/scripts/aliyun-ecs-runner-provision.py new file mode 100644 index 00000000000..74a69681f3d --- /dev/null +++ b/.github/scripts/aliyun-ecs-runner-provision.py @@ -0,0 +1,485 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# 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. + +# This is the first `.github/scripts/` user of PEP 723 inline script metadata: +# `uv run` installs the pinned Aliyun SDK before executing. The SDK import stays +# lazy (inside `make_ecs_client`) so unit tests can import this module with a +# plain stdlib interpreter. +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# # Pin exactly after the first live run against the Aliyun API. +# "alibabacloud_ecs20140526>=4.1.0,<6", +# "alibabacloud_tea_openapi>=0.3.12,<1", +# ] +# /// + +"""Provision an ephemeral Aliyun ECS query-regression runner. + +Creates one pay-as-you-go ECS instance from the query-regression custom image +and waits until the instance registers itself as an ephemeral GitHub Actions +runner with a per-run label. Build caches live on the instance system disk +and disappear with the VM. The instance id is written to job outputs as soon +as the VM exists so the teardown job is the single DeleteInstance caller (a +timeout here used to delete and then teardown deleted again). On runner +online timeout the console is dumped and the job fails; teardown releases +the instance. + +The Aliyun credentials are used only in this control-plane job (ubuntu-latest) +and are never passed to the ECS instance; the instance receives only a +short-lived GitHub runner registration token via user data. +""" + +from __future__ import annotations + +import argparse +import base64 +import json +import os +import time +import urllib.error +import urllib.request +from dataclasses import dataclass + +RUNNER_NAME_PREFIX = "qreg-ecs" +RUNNER_LABEL_PREFIX = "query-regression-ecs" +MANAGED_BY_TAG_KEY = "managed-by" +MANAGED_BY_TAG_VALUE = "query-regression-ci" +RUN_TAG_KEY = "query-regression-run-id" + +# Runner cache paths on the instance system disk. They are created empty +# every provision and deleted with the VM. +CACHE_PATHS = ( + "/home/runner/.cargo/registry", + "/home/runner/.cargo/git", + "/home/runner/query-regression-target", + "/home/runner/query-regression-cache-meta", + "/home/runner/.cache/sccache", +) + +RUNNER_ONLINE_TIMEOUT_SECONDS = 10 * 60 +POLL_INTERVAL_SECONDS = 5 + +# Nightly thin-LTO of greptime peaks above ecs.c9i.2xlarge's 16 GiB. A swap +# file plus masking systemd-oomd lets the kernel reclaim rustc pages instead +# of SIGTERM-ing the runner service cgroup (GitHub then reports the step as +# "The operation was canceled"). Sized to match that 16 GiB instance; a +# 32 GiB type still benefits as a safety net. Lives on the system disk next +# to the cold build caches. +SWAP_FILE = "/swapfile" +SWAP_SIZE_GIB = 16 + + +@dataclass(frozen=True) +class ProvisionConfig: + region_id: str + vswitch_id: str + security_group_id: str + image_id: str + instance_type: str + repo: str + run_id: str + github_token: str + # Optional resource group; required when the RAM grant is scoped to one. + resource_group_id: str | None = None + # Runner identity inside the image; the workflow asserts the same values. + runner_uid: str = "1001" + runner_gid: str = "1001" + + +def runner_name_for_run(run_id: str) -> str: + return f"{RUNNER_NAME_PREFIX}-{run_id}" + + +def runner_label_for_run(run_id: str) -> str: + return f"{RUNNER_LABEL_PREFIX}-{run_id}" + + +def render_user_data( + runner_name: str, + runner_label: str, + runner_token: str, + repo: str, + runner_uid: str = "1001", + runner_gid: str = "1001", +) -> str: + """Render the cloud-init shell script for the runner instance.""" + destinations = " ".join(f'"{dst}"' for dst in CACHE_PATHS) + cache_setup = f"""# Caches live on the system disk, are deleted with the instance, and every +# run compiles cold. Within-run reuse (base warming candidate via the shared +# target dir and sccache) still applies. +mkdir -p {destinations} +chown -R {runner_uid}:{runner_gid} /home/runner""" + swap_setup = f"""# Mask systemd-oomd before enabling swap: Ubuntu 24.04 kills the whole +# service cgroup on PSI pressure, and swap thrashing looks like pressure. +systemctl disable --now systemd-oomd.socket systemd-oomd.service || true +systemctl mask systemd-oomd.socket systemd-oomd.service || true +# apt-daily-upgrade / unattended-upgrades can `systemctl restart` services +# whose libraries were updated. That SIGTERM-s the runner mid-job; GitHub +# records UserCancelled (job still valid, child exit 143). These VMs live +# ~1h; security updates belong in the image, not at job time. +systemctl disable --now unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service || true +systemctl mask unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service || true +systemctl stop unattended-upgrades.service || true +killall -9 unattended-upgr || true +cat > /etc/apt/apt.conf.d/99disable-auto-updates <<'APTEOF' +APT::Periodic::Update-Package-Lists "0"; +APT::Periodic::Unattended-Upgrade "0"; +APT::Periodic::Download-Upgradeable-Packages "0"; +APTEOF +if [[ ! -f "{SWAP_FILE}" ]]; then + fallocate --length {SWAP_SIZE_GIB}G "{SWAP_FILE}" + chmod 600 "{SWAP_FILE}" + mkswap "{SWAP_FILE}" +fi +swapon "{SWAP_FILE}" +sysctl --write vm.swappiness=10 +swapon --show +free --human""" + return f"""#!/bin/bash +set -euo pipefail + +{cache_setup} + +{swap_setup} + +cat > /etc/ephemeral-github-runner.env <<'ENVEOF' +RUNNER_NAME={runner_name} +RUNNER_LABELS={runner_label} +RUNNER_TOKEN={runner_token} +REPO_URL=https://github.com/{repo} +# The container image baked /opt/cargo/bin into PATH via Dockerfile ENV; on +# the host nothing inherits that, so publish it through the unit's +# EnvironmentFile: runner job processes inherit the runner's environment. +PATH=/opt/cargo/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin +ENVEOF + +# Stream the runner service's own logs (config output, job lifecycle, crash +# messages) to the serial console, next to the kernel's OOM-killer records. +# Teardown dumps the console tail before deleting the instance, so the +# witness of what killed the runner survives the machine. +mkdir -p /etc/systemd/system/ephemeral-github-runner.service.d +cat > /etc/systemd/system/ephemeral-github-runner.service.d/console.conf <<'CONFEOF' +[Service] +StandardOutput=journal+console +StandardError=journal+console +CONFEOF +# systemd's DefaultOOMPolicy=stop SIGTERM-s the whole unit when *any* +# process in the cgroup is OOM-killed (typically rustc during thin LTO). +# That is the "Session terminated, killing shell..." / job Canceled path: +# runuser dies, GitHub deletes the ephemeral registration, always() steps +# never run. continue keeps the runner up so cargo can report signal 9. +cat > /etc/systemd/system/ephemeral-github-runner.service.d/oom.conf <<'CONFEOF' +[Service] +OOMPolicy=continue +CONFEOF +systemctl daemon-reload +# restart (not start): the image enables this unit, so it may already +# have come up without the drop-ins if it raced cloud-init. restart +# applies OOMPolicy=continue to the running cgroup. --no-block matters: +# newer image units carry After=cloud-final.service, and this script IS +# cloud-final — a blocking restart would deadlock until TimeoutStartSec. +systemctl restart --no-block ephemeral-github-runner.service +""" + + +def encode_user_data(script: str) -> str: + return base64.b64encode(script.encode("utf-8")).decode("ascii") + + +def github_api(token: str, method: str, path: str, body: dict | None = None) -> dict: + data = json.dumps(body).encode("utf-8") if body is not None else None + request = urllib.request.Request( + f"https://api.github.com{path}", + data=data, + method=method, + headers={ + "Accept": "application/vnd.github+json", + "Authorization": f"Bearer {token}", + "X-GitHub-Api-Version": "2022-11-28", + }, + ) + try: + with urllib.request.urlopen(request, timeout=30) as response: + return json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as error: + # GitHub's error body says exactly why (e.g. "Must have admin rights to + # Repository" for a PAT without the required scope); surface it instead + # of a bare "HTTP Error 403". + body = error.read().decode("utf-8", "replace") + raise SystemExit( + f"GitHub API {method} {path} failed: HTTP {error.code}: {body}\n" + "The token comes from the GH_PERSONAL_ACCESS_TOKEN secret; it needs " + "'repo' scope (classic PAT) or 'Administration: write' on the " + "repository (fine-grained PAT)." + ) from error + + +def create_registration_token(github_token: str, repo: str) -> str: + response = github_api( + github_token, + "POST", + f"/repos/{repo}/actions/runners/registration-token", + body={}, + ) + return response["token"] + + +def find_runner_by_name(github_token: str, repo: str, name: str) -> dict | None: + page = 1 + while True: + response = github_api( + github_token, + "GET", + f"/repos/{repo}/actions/runners?per_page=100&page={page}", + ) + runners = response.get("runners", []) + for runner in runners: + if runner.get("name") == name: + return runner + if len(runners) < 100: + return None + page += 1 + + +def make_ecs_client(config: ProvisionConfig): + from alibabacloud_ecs20140526.client import Client as EcsClient + from alibabacloud_tea_openapi.models import Config as OpenApiConfig + + access_key_id = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_ID", "") + access_key_secret = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_SECRET", "") + if not access_key_id or not access_key_secret: + raise SystemExit( + "ALIBABA_CLOUD_ACCESS_KEY_ID/SECRET are empty; in CI they come from the " + "ALICLOUD_ECS_ACCESS_KEY_ID/SECRET repository secrets (a missing or " + "misnamed secret expands to an empty string)." + ) + return EcsClient( + OpenApiConfig( + access_key_id=access_key_id, + access_key_secret=access_key_secret, + region_id=config.region_id, + endpoint=f"ecs.{config.region_id}.aliyuncs.com", + ) + ) + + +def wait_for_instance_status(client, region_id: str, instance_id: str, wanted: str, deadline: float) -> None: + from alibabacloud_ecs20140526 import models as ecs_models + + while time.monotonic() < deadline: + response = client.describe_instances( + ecs_models.DescribeInstancesRequest( + region_id=region_id, instance_ids=json.dumps([instance_id]) + ) + ) + instances = response.body.instances.instance + if instances and instances[0].status == wanted: + return + time.sleep(POLL_INTERVAL_SECONDS) + raise TimeoutError(f"Instance {instance_id} did not reach status {wanted} in time") + + +def fetch_console_output(client, region_id: str, instance_id: str) -> str: + from alibabacloud_ecs20140526 import models as ecs_models + + response = client.get_instance_console_output( + ecs_models.GetInstanceConsoleOutputRequest(region_id=region_id, instance_id=instance_id) + ) + return base64.b64decode(response.body.console_output or "").decode("utf-8", "replace") + + +class ConsoleTailer: + """Incrementally prints new serial-console lines so cloud-init progress is + visible in the CI log while waiting, not only after a failure dump.""" + + def __init__(self, client, region_id: str, instance_id: str) -> None: + self.client = client + self.region_id = region_id + self.instance_id = instance_id + self.printed_lines = 0 + + def poll(self) -> None: + try: + lines = fetch_console_output(self.client, self.region_id, self.instance_id).splitlines() + except Exception as error: # noqa: BLE001 + print(f"[console] unable to fetch console output: {error}", flush=True) + return + # The API returns a bounded tail; if it ever shrinks, restart from the + # beginning of the new buffer rather than skipping lines. + if len(lines) < self.printed_lines: + self.printed_lines = 0 + for line in lines[self.printed_lines :]: + print(f"[console] {line}", flush=True) + self.printed_lines = len(lines) + + +def dump_console_output(client, region_id: str, instance_id: str) -> None: + try: + output = fetch_console_output(client, region_id, instance_id) + except Exception as error: # noqa: BLE001 + print(f"Unable to fetch console output for {instance_id}: {error}", flush=True) + return + tail = "\n".join(output.splitlines()[-80:]) + print(f"::group::Console output tail for {instance_id}\n{tail}\n::endgroup::", flush=True) + summary_path = os.environ.get("GITHUB_STEP_SUMMARY") + if summary_path: + with open(summary_path, "a", encoding="utf-8") as summary: + summary.write(f"
ECS console output ({instance_id})\n\n```\n") + summary.write(tail) + summary.write("\n```\n
\n") + + +def append_github_output(name: str, value: str) -> None: + output_path = os.environ.get("GITHUB_OUTPUT") + if output_path: + with open(output_path, "a", encoding="utf-8") as output: + output.write(f"{name}={value}\n") + + +def run_instance(client, config: ProvisionConfig, user_data: str) -> str: + from alibabacloud_ecs20140526 import models as ecs_models + + request = ecs_models.RunInstancesRequest( + region_id=config.region_id, + image_id=config.image_id, + resource_group_id=config.resource_group_id, + instance_type=config.instance_type, + v_switch_id=config.vswitch_id, + security_group_id=config.security_group_id, + instance_name=runner_name_for_run(config.run_id), + host_name=runner_name_for_run(config.run_id), + description=f"Ephemeral query-regression runner for run {config.run_id}", + amount=1, + instance_charge_type="PostPaid", + spot_strategy="NoSpot", + internet_charge_type="PayByTraffic", + internet_max_bandwidth_out=100, + # Holds the image (~12G), the 16G swapfile, the source checkout, the + # base+candidate cluster data homes, and cold build caches. Deleted + # with the instance. + system_disk=ecs_models.RunInstancesRequestSystemDisk(category="cloud_essd", size="40"), + user_data=user_data, + tag=[ + ecs_models.RunInstancesRequestTag(key=MANAGED_BY_TAG_KEY, value=MANAGED_BY_TAG_VALUE), + ecs_models.RunInstancesRequestTag(key=RUN_TAG_KEY, value=config.run_id), + ], + ) + response = client.run_instances(request) + instance_id = response.body.instance_id_sets.instance_id_set[0] + return instance_id + + +def provision(config: ProvisionConfig) -> int: + client = make_ecs_client(config) + runner_name = runner_name_for_run(config.run_id) + runner_label = runner_label_for_run(config.run_id) + + registration_token = create_registration_token(config.github_token, config.repo) + user_data = encode_user_data( + render_user_data( + runner_name, + runner_label, + registration_token, + config.repo, + config.runner_uid, + config.runner_gid, + ) + ) + + print(f"::group::Run instance {runner_name}", flush=True) + instance_id = run_instance(client, config, user_data) + print(f"Instance id: {instance_id}", flush=True) + # Teardown is the only DeleteInstance caller. Write outputs immediately + # so a later status/online failure still releases this VM. + append_github_output("label", runner_label) + append_github_output("instance_id", instance_id) + append_github_output("runner_name", runner_name) + wait_for_instance_status(client, config.region_id, instance_id, "Running", time.monotonic() + 5 * 60) + print("::endgroup::", flush=True) + + print("::group::Wait for runner online", flush=True) + deadline = time.monotonic() + RUNNER_ONLINE_TIMEOUT_SECONDS + tailer = ConsoleTailer(client, config.region_id, instance_id) + last_console_poll = 0.0 + while time.monotonic() < deadline: + runner = find_runner_by_name(config.github_token, config.repo, runner_name) + if runner is not None: + status = runner.get("status") + print(f"Runner {runner_name} status: {status}", flush=True) + if status == "online": + print("::endgroup::", flush=True) + print(f"Runner {runner_name} is online with label {runner_label}", flush=True) + return 0 + # The serial console lags the guest by a minute or so; 30s polling is + # enough to follow cloud-init without tripping API throttling. + if time.monotonic() - last_console_poll >= 30: + tailer.poll() + last_console_poll = time.monotonic() + time.sleep(POLL_INTERVAL_SECONDS) + + print("::endgroup::", flush=True) + print( + f"Runner {runner_name} did not come online in time; " + "leaving instance for the teardown job to delete", + flush=True, + ) + dump_console_output(client, config.region_id, instance_id) + return 1 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--region-id", default=os.environ.get("ALIYUN_ECS_REGION_ID")) + parser.add_argument("--vswitch-id", default=os.environ.get("ALIYUN_ECS_VSWITCH_ID")) + parser.add_argument("--security-group-id", default=os.environ.get("ALIYUN_ECS_SECURITY_GROUP_ID")) + parser.add_argument("--image-id", default=os.environ.get("QUERY_REGRESSION_ECS_IMAGE_ID")) + parser.add_argument("--instance-type", default=os.environ.get("ALIYUN_ECS_INSTANCE_TYPE")) + parser.add_argument("--repo", default=os.environ.get("GITHUB_REPOSITORY")) + parser.add_argument("--run-id", default=os.environ.get("GITHUB_RUN_ID")) + parser.add_argument("--github-token", default=os.environ.get("GH_PERSONAL_ACCESS_TOKEN")) + parser.add_argument("--runner-uid", default=os.environ.get("QUERY_REGRESSION_RUNNER_UID", "1001")) + parser.add_argument("--runner-gid", default=os.environ.get("QUERY_REGRESSION_RUNNER_GID", "1001")) + parser.add_argument("--resource-group-id", default=os.environ.get("ALIYUN_ECS_RESOURCE_GROUP_ID")) + args = parser.parse_args() + + missing = [ + name + for name, value in vars(args).items() + if name not in ("resource_group_id",) + and (value is None or (isinstance(value, str) and not value)) + ] + if missing: + flags = ", ".join(f"--{name.replace('_', '-')}" for name in missing) + raise SystemExit(f"Missing required configuration: {flags} (or their env defaults)") + + config = ProvisionConfig( + region_id=args.region_id, + vswitch_id=args.vswitch_id, + security_group_id=args.security_group_id, + image_id=args.image_id, + instance_type=args.instance_type, + repo=args.repo, + run_id=args.run_id, + github_token=args.github_token, + resource_group_id=args.resource_group_id or None, + runner_uid=args.runner_uid, + runner_gid=args.runner_gid, + ) + return provision(config) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/scripts/aliyun-ecs-runner-teardown.py b/.github/scripts/aliyun-ecs-runner-teardown.py new file mode 100644 index 00000000000..36c80253b95 --- /dev/null +++ b/.github/scripts/aliyun-ecs-runner-teardown.py @@ -0,0 +1,237 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# 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. + +# PEP 723 inline metadata (see aliyun-ecs-runner-provision.py for the +# convention). The SDK import stays lazy so unit tests run on a plain +# stdlib interpreter. +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "alibabacloud_ecs20140526>=4.1.0,<6", +# "alibabacloud_tea_openapi>=0.3.12,<1", +# ] +# /// + +"""Tear down ephemeral Aliyun ECS query-regression runners. + +Two modes: + +- Targeted (per workflow run): delete one instance by id and deregister its + runner by name. Both steps are idempotent and best-effort; a missing + instance or runner is not an error. +- Sweep (scheduled janitor): delete every instance tagged as managed by + query-regression CI whose creation time is older than the given TTL, and + deregister the matching runners. This is the safety net for runs whose + teardown job never executed. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import os +import sys +from datetime import datetime, timedelta, timezone +from pathlib import Path + +# Reuse the GitHub API client, runner lookup, and tag constants from the +# provision script. +_PROVISION_SPEC = importlib.util.spec_from_file_location( + "aliyun_ecs_runner_provision", + Path(__file__).resolve().parent / "aliyun-ecs-runner-provision.py", +) +assert _PROVISION_SPEC is not None and _PROVISION_SPEC.loader is not None +provision = importlib.util.module_from_spec(_PROVISION_SPEC) +sys.modules[_PROVISION_SPEC.name] = provision +_PROVISION_SPEC.loader.exec_module(provision) + +SWEEP_TTL = timedelta(hours=4) +# ECS DescribeInstances creation_time is ISO8601 UTC, e.g. "2026-08-17T02:13Z". +CREATION_TIME_FORMATS = ("%Y-%m-%dT%H:%M:%SZ", "%Y-%m-%dT%H:%MZ") + + +def parse_creation_time(value: str) -> datetime: + for fmt in CREATION_TIME_FORMATS: + try: + return datetime.strptime(value, fmt).replace(tzinfo=timezone.utc) + except ValueError: + continue + raise ValueError(f"Unparsable ECS creation time: {value}") + + +def expired_instance_names( + instances: list[tuple[str, str, str]], now: datetime, ttl: timedelta +) -> list[tuple[str, str]]: + """Pick (instance_id, instance_name) pairs whose creation is older than ttl. + + `instances` items are (instance_id, instance_name, creation_time). + """ + expired = [] + for instance_id, instance_name, creation_time in instances: + age = now - parse_creation_time(creation_time) + if age >= ttl: + expired.append((instance_id, instance_name)) + return expired + + +def make_ecs_client(region_id: str): + from alibabacloud_ecs20140526.client import Client as EcsClient + from alibabacloud_tea_openapi.models import Config as OpenApiConfig + + access_key_id = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_ID", "") + access_key_secret = os.environ.get("ALIBABA_CLOUD_ACCESS_KEY_SECRET", "") + if not access_key_id or not access_key_secret: + raise SystemExit( + "ALIBABA_CLOUD_ACCESS_KEY_ID/SECRET are empty; in CI they come from the " + "ALICLOUD_ECS_ACCESS_KEY_ID/SECRET repository secrets (a missing or " + "misnamed secret expands to an empty string)." + ) + return EcsClient( + OpenApiConfig( + access_key_id=access_key_id, + access_key_secret=access_key_secret, + region_id=region_id, + endpoint=f"ecs.{region_id}.aliyuncs.com", + ) + ) + + +def delete_instance(client, instance_id: str, region_id: str | None = None) -> bool: + import time + + from alibabacloud_ecs20140526 import models as ecs_models + + # Export the serial console before destroying the evidence: a cancelled + # run still runs this teardown job, and the console tail (cloud-init, + # runner service logs, kernel OOM records) is the only witness of what + # happened on the machine. + if region_id is not None: + provision.dump_console_output(client, region_id, instance_id) + + # DeleteInstance rejects instances that are still Initializing (or in + # another transitional status). Teardown runs right after a cancelled + # run, when the instance may be only seconds old, so retry the transient + # status errors for a few minutes instead of leaking the instance. + deadline = time.monotonic() + 5 * 60 + while True: + try: + client.delete_instance(ecs_models.DeleteInstanceRequest(instance_id=instance_id, force=True)) + print(f"Deleted instance {instance_id}", flush=True) + return True + except Exception as error: # noqa: BLE001 + message = str(error) + if "InvalidInstanceId.NotFound" in message: + print(f"Instance {instance_id} already gone", flush=True) + return True + if "IncorrectInstanceStatus" in message and time.monotonic() < deadline: + print(f"Instance {instance_id} is in a transitional status; retrying delete", flush=True) + time.sleep(15) + continue + print(f"Failed to delete instance {instance_id}: {error}", flush=True) + return False + + +def deregister_runner(token: str, repo: str, runner_name: str) -> bool: + runner = provision.find_runner_by_name(token, repo, runner_name) + if runner is None: + print(f"Runner {runner_name} is not registered", flush=True) + return True + try: + provision.github_api(token, "DELETE", f"/repos/{repo}/actions/runners/{runner['id']}") + print(f"Deregistered runner {runner_name} (id {runner['id']})", flush=True) + return True + except Exception as error: # noqa: BLE001 + print(f"Failed to deregister runner {runner_name}: {error}", flush=True) + return False + + +def list_managed_instances(client, region_id: str) -> list[tuple[str, str, str]]: + from alibabacloud_ecs20140526 import models as ecs_models + + result: list[tuple[str, str, str]] = [] + next_token = None + while True: + request = ecs_models.DescribeInstancesRequest( + region_id=region_id, + tag=[ + ecs_models.DescribeInstancesRequestTag( + key=provision.MANAGED_BY_TAG_KEY, value=provision.MANAGED_BY_TAG_VALUE + ) + ], + max_results=100, + next_token=next_token, + ) + response = client.describe_instances(request) + for instance in response.body.instances.instance: + result.append((instance.instance_id, instance.instance_name, instance.creation_time)) + next_token = response.body.next_token + if not next_token: + return result + + +def sweep(client, region_id: str, repo: str, github_token: str, ttl: timedelta) -> int: + instances = list_managed_instances(client, region_id) + print(f"Found {len(instances)} managed instance(s) in {region_id}", flush=True) + expired = expired_instance_names(instances, datetime.now(timezone.utc), ttl) + ok = True + for instance_id, instance_name in expired: + print(f"Instance {instance_id} ({instance_name}) exceeds TTL {ttl}; deleting", flush=True) + # Runner names mirror instance names by construction in the provision + # script (both are qreg-ecs-). + ok &= delete_instance(client, instance_id, region_id) + ok &= deregister_runner(github_token, repo, instance_name) + return 0 if ok else 1 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--region-id", default=os.environ.get("ALIYUN_ECS_REGION_ID")) + parser.add_argument("--repo", default=os.environ.get("GITHUB_REPOSITORY")) + parser.add_argument("--github-token", default=os.environ.get("GH_PERSONAL_ACCESS_TOKEN")) + parser.add_argument("--instance-id", default=os.environ.get("QUERY_REGRESSION_ECS_INSTANCE_ID")) + parser.add_argument("--runner-name", default=os.environ.get("QUERY_REGRESSION_ECS_RUNNER_NAME")) + parser.add_argument( + "--sweep", + action="store_true", + help="Janitor mode: delete all managed instances older than --sweep-ttl-hours.", + ) + parser.add_argument("--sweep-ttl-hours", type=float, default=SWEEP_TTL.total_seconds() / 3600) + args = parser.parse_args() + + for name in ("region_id", "repo", "github_token"): + if not getattr(args, name): + raise SystemExit(f"Missing required configuration: --{name.replace('_', '-')}") + + client = make_ecs_client(args.region_id) + + if args.sweep: + return sweep( + client, args.region_id, args.repo, args.github_token, timedelta(hours=args.sweep_ttl_hours) + ) + + ok = True + if args.instance_id: + ok &= delete_instance(client, args.instance_id, args.region_id) + else: + print("No instance id given; skipping instance deletion", flush=True) + if args.runner_name: + ok &= deregister_runner(args.github_token, args.repo, args.runner_name) + else: + print("No runner name given; skipping runner deregistration", flush=True) + return 0 if ok else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/scripts/query-regression-nightly-refs.py b/.github/scripts/query-regression-nightly-refs.py new file mode 100644 index 00000000000..6e9f54682d4 --- /dev/null +++ b/.github/scripts/query-regression-nightly-refs.py @@ -0,0 +1,270 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# 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. + +"""Resolve consecutive successful Nightly Build SHAs for query-regression. + +Query regression builds both binaries from git; this script only picks the +head SHAs of two Nightly Build workflow runs (today vs the previous success). +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import urllib.error +import urllib.parse +import urllib.request +from dataclasses import dataclass +from typing import Any + + +NIGHTLY_WORKFLOW = "nightly-build.yml" +NIGHTLY_EVENTS = frozenset({"schedule", "workflow_dispatch"}) + + +@dataclass(frozen=True) +class WorkflowRun: + id: int + head_sha: str + head_branch: str + html_url: str + created_at: str + conclusion: str + event: str + + +@dataclass(frozen=True) +class RefPair: + base: WorkflowRun | None + candidate: WorkflowRun | None + skip: bool + reason: str + + +def parse_run(payload: dict[str, Any]) -> WorkflowRun: + return WorkflowRun( + id=int(payload["id"]), + head_sha=str(payload["head_sha"]), + head_branch=str(payload.get("head_branch") or ""), + html_url=str(payload.get("html_url") or ""), + created_at=str(payload.get("created_at") or ""), + conclusion=str(payload.get("conclusion") or ""), + event=str(payload.get("event") or ""), + ) + + +def github_get(token: str, path: str, query: dict[str, str] | None = None) -> dict[str, Any]: + encoded = f"?{urllib.parse.urlencode(query)}" if query else "" + request = urllib.request.Request( + f"https://api.github.com{path}{encoded}", + method="GET", + headers={ + "Accept": "application/vnd.github+json", + "Authorization": f"Bearer {token}", + "X-GitHub-Api-Version": "2022-11-28", + }, + ) + try: + with urllib.request.urlopen(request, timeout=30) as response: + return json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as error: + body = error.read().decode("utf-8", "replace") + raise SystemExit( + f"GitHub API GET {path} failed: HTTP {error.code}: {body}" + ) from error + + +def list_successful_nightly_runs( + token: str, + repo: str, + *, + branch: str | None = None, + per_page: int = 30, +) -> list[WorkflowRun]: + query = { + "status": "success", + "per_page": str(per_page), + } + if branch: + query["branch"] = branch + payload = github_get( + token, + f"/repos/{repo}/actions/workflows/{NIGHTLY_WORKFLOW}/runs", + query, + ) + runs = [parse_run(item) for item in payload.get("workflow_runs") or []] + return [ + run + for run in runs + if run.conclusion == "success" and run.event in NIGHTLY_EVENTS and run.head_sha + ] + + +def fetch_run(token: str, repo: str, run_id: int) -> WorkflowRun: + payload = github_get(token, f"/repos/{repo}/actions/runs/{run_id}") + return parse_run(payload) + + +def select_base_and_candidate( + runs: list[WorkflowRun], + *, + candidate_run_id: int | None = None, + candidate: WorkflowRun | None = None, +) -> RefPair: + """Pick candidate (newest / requested) and the previous successful run. + + `runs` must be newest-first. Same-SHA consecutive nightlies are skipped + (no commits that day). + """ + if candidate is None: + if candidate_run_id is not None: + candidate = next((run for run in runs if run.id == candidate_run_id), None) + if candidate is None: + return RefPair( + None, + None, + True, + f"candidate nightly run {candidate_run_id} was not a successful " + f"{NIGHTLY_WORKFLOW} run", + ) + elif runs: + candidate = runs[0] + else: + return RefPair(None, None, True, "no successful Nightly Build runs") + + branch = candidate.head_branch + older = [ + run + for run in runs + if run.id != candidate.id + and (not branch or not run.head_branch or run.head_branch == branch) + and (not candidate.created_at or run.created_at <= candidate.created_at) + ] + if not older: + return RefPair( + None, + candidate, + True, + "no previous successful Nightly Build to compare against", + ) + base = older[0] + if base.head_sha.lower() == candidate.head_sha.lower(): + return RefPair( + base, + candidate, + True, + f"previous nightly SHA {base.head_sha} matches candidate; nothing to compare", + ) + return RefPair(base, candidate, False, "") + + +def _override_run(sha: str) -> WorkflowRun: + return WorkflowRun( + id=0, + head_sha=sha, + head_branch="", + html_url="", + created_at="", + conclusion="success", + event="workflow_dispatch", + ) + + +def override_pair(base_ref: str, candidate_ref: str) -> RefPair: + return RefPair(_override_run(base_ref), _override_run(candidate_ref), False, "explicit refs") + + +def write_outputs(pair: RefPair) -> None: + values = { + "skip": "true" if pair.skip else "false", + "reason": pair.reason, + "base_sha": pair.base.head_sha if pair.base else "", + "candidate_sha": pair.candidate.head_sha if pair.candidate else "", + "base_run_url": pair.base.html_url if pair.base else "", + "candidate_run_url": pair.candidate.html_url if pair.candidate else "", + "base_run_id": str(pair.base.id) if pair.base and pair.base.id else "", + "candidate_run_id": str(pair.candidate.id) if pair.candidate and pair.candidate.id else "", + } + output_path = os.environ.get("GITHUB_OUTPUT") + if output_path: + with open(output_path, "a", encoding="utf-8") as handle: + for key, value in values.items(): + handle.write(f"{key}={value}\n") + for key, value in values.items(): + print(f"{key}={value}") + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--repo", default=os.environ.get("GITHUB_REPOSITORY")) + parser.add_argument( + "--token", + default=os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN"), + ) + parser.add_argument("--branch", default=os.environ.get("NIGHTLY_BRANCH") or "") + parser.add_argument( + "--candidate-run-id", + default=os.environ.get("CANDIDATE_RUN_ID") or "", + ) + parser.add_argument("--base-ref", default=os.environ.get("BASE_REF") or "") + parser.add_argument("--candidate-ref", default=os.environ.get("CANDIDATE_REF") or "") + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + base_ref = args.base_ref.strip() + candidate_ref = args.candidate_ref.strip() + if base_ref and candidate_ref: + write_outputs(override_pair(base_ref, candidate_ref)) + return 0 + if bool(base_ref) != bool(candidate_ref): + print("Provide both --base-ref and --candidate-ref, or neither.", file=sys.stderr) + return 2 + if not args.repo or not args.token: + print("--repo and --token (or GITHUB_REPOSITORY / GITHUB_TOKEN) are required.", file=sys.stderr) + return 2 + + candidate_run_id = int(args.candidate_run_id) if str(args.candidate_run_id).strip() else None + candidate: WorkflowRun | None = None + branch = args.branch.strip() or None + if candidate_run_id is not None: + candidate = fetch_run(args.token, args.repo, candidate_run_id) + if candidate.conclusion != "success": + write_outputs( + RefPair( + None, + candidate, + True, + f"nightly run {candidate_run_id} conclusion is {candidate.conclusion}", + ) + ) + return 0 + branch = branch or candidate.head_branch or None + + runs = list_successful_nightly_runs(args.token, args.repo, branch=branch) + pair = select_base_and_candidate( + runs, + candidate_run_id=candidate_run_id if candidate is None else None, + candidate=candidate, + ) + write_outputs(pair) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 7f064f52eb5..8d7a37080b3 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -80,6 +80,14 @@ jobs: run: python3 scripts/check-enterprise-license-test.py - name: Run fuzz orchestration script tests run: .github/scripts/run-fuzz-targets-test.sh + - name: Run query-regression Python tooling tests + run: | + python3 tests/perf/test_query_regression_runner_compaction_toctou.py + python3 tests/perf/test_query_regression_runner_otlp_trace_load.py + python3 tests/perf/test_query_regression_summary_otlp.py + python3 tests/perf/test_query_regression_case_selection.py + python3 tests/perf/test_query_regression_nightly_refs.py + python3 tests/perf/test_aliyun_ecs_runner_scripts.py check: if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} diff --git a/.github/workflows/query-regression-janitor.yml b/.github/workflows/query-regression-janitor.yml new file mode 100644 index 00000000000..ce277ce2f02 --- /dev/null +++ b/.github/workflows/query-regression-janitor.yml @@ -0,0 +1,36 @@ +name: Query Regression ECS Janitor + +# Safety net for the aliyun-ecs query-regression path: deletes managed ECS +# instances (and their runner registrations) that outlived their workflow run, +# so a failed teardown never leaks pay-as-you-go spend. + +on: + schedule: + - cron: "23 3 * * *" + workflow_dispatch: + +permissions: + contents: read + +jobs: + sweep: + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - name: Checkout teardown script + uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Install uv + uses: astral-sh/setup-uv@v6 + + - name: Sweep expired ECS runners + env: + ALIBABA_CLOUD_ACCESS_KEY_ID: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + ALIBABA_CLOUD_ACCESS_KEY_SECRET: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + ALIYUN_ECS_REGION_ID: ${{ vars.ALIYUN_ECS_REGION_ID }} + run: >- + uv run .github/scripts/aliyun-ecs-runner-teardown.py + --sweep --sweep-ttl-hours 4 diff --git a/.github/workflows/query-regression-nightly.yml b/.github/workflows/query-regression-nightly.yml new file mode 100644 index 00000000000..027fcd16987 --- /dev/null +++ b/.github/workflows/query-regression-nightly.yml @@ -0,0 +1,104 @@ +name: Query Regression Nightly + +# After a successful GreptimeDB Nightly Build, compare that commit against +# the previous successful nightly. The reusable Query Regression workflow +# still compiles both SHAs; this wrapper only resolves which two SHAs. + +on: + workflow_run: + workflows: + - GreptimeDB Nightly Build + types: + - completed + workflow_dispatch: + inputs: + base_ref: + description: Base ref/SHA (empty = previous successful nightly) + required: false + type: string + default: "" + candidate_ref: + description: Candidate ref/SHA (empty = latest successful nightly) + required: false + type: string + default: "" + candidate_run_id: + description: Nightly Build run id to treat as candidate (empty = latest) + required: false + type: string + default: "" + case: + description: Query perf case path(s), all, or heavy + required: false + type: string + default: all + +permissions: + contents: read + actions: read + +jobs: + resolve-refs: + name: Resolve previous vs current nightly SHAs + if: ${{ github.event_name == 'workflow_dispatch' || github.event.workflow_run.conclusion == 'success' }} + runs-on: ubuntu-latest + timeout-minutes: 5 + outputs: + skip: ${{ steps.refs.outputs.skip }} + reason: ${{ steps.refs.outputs.reason }} + base_sha: ${{ steps.refs.outputs.base_sha }} + candidate_sha: ${{ steps.refs.outputs.candidate_sha }} + steps: + - name: Checkout ref resolver + uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Resolve nightly SHAs + id: refs + env: + GITHUB_TOKEN: ${{ github.token }} + GITHUB_REPOSITORY: ${{ github.repository }} + CANDIDATE_RUN_ID: ${{ github.event_name == 'workflow_run' && github.event.workflow_run.id || inputs.candidate_run_id || '' }} + BASE_REF: ${{ github.event_name == 'workflow_dispatch' && inputs.base_ref || '' }} + CANDIDATE_REF: ${{ github.event_name == 'workflow_dispatch' && inputs.candidate_ref || '' }} + run: python3 .github/scripts/query-regression-nightly-refs.py + + - name: Summarize comparison + if: always() + env: + SKIP: ${{ steps.refs.outputs.skip }} + REASON: ${{ steps.refs.outputs.reason }} + BASE_SHA: ${{ steps.refs.outputs.base_sha }} + CANDIDATE_SHA: ${{ steps.refs.outputs.candidate_sha }} + BASE_RUN_URL: ${{ steps.refs.outputs.base_run_url }} + CANDIDATE_RUN_URL: ${{ steps.refs.outputs.candidate_run_url }} + run: | + set -euo pipefail + { + if [[ "${SKIP}" == "true" ]]; then + printf 'Skipping query-regression nightly: %s\n' "${REASON}" + else + printf 'Comparing previous nightly `%s` -> current nightly `%s`\n' \ + "${BASE_SHA}" "${CANDIDATE_SHA}" + if [[ -n "${BASE_RUN_URL}" ]]; then + printf -- '- Previous Nightly Build: %s\n' "${BASE_RUN_URL}" + fi + if [[ -n "${CANDIDATE_RUN_URL}" ]]; then + printf -- '- Current Nightly Build: %s\n' "${CANDIDATE_RUN_URL}" + fi + fi + } | tee -a "${GITHUB_STEP_SUMMARY}" + + query-regression: + name: Query regression nightly + needs: [resolve-refs] + if: ${{ needs.resolve-refs.outputs.skip != 'true' && needs.resolve-refs.outputs.base_sha != '' && needs.resolve-refs.outputs.candidate_sha != '' }} + uses: ./.github/workflows/query-regression.yml + secrets: inherit + with: + case: ${{ github.event_name == 'workflow_dispatch' && inputs.case || 'all' }} + base_ref: ${{ needs.resolve-refs.outputs.base_sha }} + candidate_ref: ${{ needs.resolve-refs.outputs.candidate_sha }} + cargo_profile: nightly + runner: aliyun-ecs diff --git a/.github/workflows/query-regression.yml b/.github/workflows/query-regression.yml index 0806d4d3e38..209c8d76fff 100644 --- a/.github/workflows/query-regression.yml +++ b/.github/workflows/query-regression.yml @@ -32,10 +32,10 @@ on: type: string default: nightly runner: - description: Self-hosted runner label or ARC runner scale set for this query regression run + description: Self-hosted runner label; aliyun-ecs provisions a fresh ECS instance per run required: false type: string - default: perf-regression-8-cores + default: aliyun-ecs workflow_dispatch: inputs: case: @@ -68,12 +68,20 @@ on: - release - dev runner: - description: Self-hosted runner label or ARC runner scale set for this query regression run + description: >- + Self-hosted runner label; aliyun-ecs provisions a fresh ECS instance + per run, any other value is used as a literal runner label required: true - type: choice - default: perf-regression-8-cores - options: - - perf-regression-8-cores + type: string + default: aliyun-ecs + keep_instance: + description: >- + Debug: keep the ECS instance after the run (skip teardown) so its + runner _diag logs, journal, and telemetry can be inspected; the + janitor still sweeps it after the TTL + required: false + type: boolean + default: false pull_request: types: [labeled] @@ -81,19 +89,99 @@ permissions: contents: read jobs: - query-regression: + test-tooling: + # Stdlib unittests for the workflow Python (case selection, report + # helpers, nightly SHA picking, and rendered ECS user-data). They do + # not talk to Aliyun or the Actions runner process; running them on + # ubuntu-latest fails fast before any ECS spend. Not gated on the + # regression labels: those labels boot a VM, these tests should not. + # Ordinary PRs also run the same tests from checks.yml, because this + # workflow only starts on `labeled` (or dispatch / workflow_call). + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + persist-credentials: false + + - name: Test query regression tooling + run: | + python3 tests/perf/test_query_regression_runner_compaction_toctou.py + python3 tests/perf/test_query_regression_runner_otlp_trace_load.py + python3 tests/perf/test_query_regression_summary_otlp.py + python3 tests/perf/test_query_regression_case_selection.py + python3 tests/perf/test_query_regression_nightly_refs.py + python3 tests/perf/test_aliyun_ecs_runner_scripts.py + + provision: + # Runs when the aliyun-ecs path is selected: explicitly via the runner + # input, or for PR labels by default (a QUERY_REGRESSION_PR_RUNNER + # repository variable set to another value redirects PRs to that literal + # runner label instead). Uses trusted scripts from the PR base (or the + # dispatched ref), never from candidate code. Waits for test-tooling so + # a broken user-data template does not still create a VM. + needs: [test-tooling] if: >- - ${{ github.event_name != 'pull_request' || - (github.event_name == 'pull_request' && - !github.event.pull_request.draft && - (github.event.label.name == 'query-regression' || - github.event.label.name == 'heavy-regression')) }} - runs-on: ${{ github.event_name != 'pull_request' && inputs.runner || 'perf-regression-8-cores' }} + ${{ !failure() && !cancelled() && + ((github.event_name != 'pull_request' && inputs.runner == 'aliyun-ecs') || + (github.event_name == 'pull_request' && + !github.event.pull_request.draft && + (github.event.label.name == 'query-regression' || + github.event.label.name == 'heavy-regression') && + (vars.QUERY_REGRESSION_PR_RUNNER || 'aliyun-ecs') == 'aliyun-ecs')) }} + runs-on: ubuntu-latest + timeout-minutes: 45 + outputs: + label: ${{ steps.provision.outputs.label }} + instance_id: ${{ steps.provision.outputs.instance_id }} + runner_name: ${{ steps.provision.outputs.runner_name }} + steps: + - name: Checkout trusted provisioning scripts + uses: actions/checkout@v4 + with: + repository: ${{ github.repository }} + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || github.sha }} + persist-credentials: false + + - name: Install uv + uses: astral-sh/setup-uv@v6 + + - name: Provision ECS runner + id: provision + env: + ALIBABA_CLOUD_ACCESS_KEY_ID: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + ALIBABA_CLOUD_ACCESS_KEY_SECRET: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + ALIYUN_ECS_REGION_ID: ${{ vars.ALIYUN_ECS_REGION_ID }} + ALIYUN_ECS_VSWITCH_ID: ${{ vars.ALIYUN_ECS_VSWITCH_ID }} + ALIYUN_ECS_SECURITY_GROUP_ID: ${{ vars.ALIYUN_ECS_SECURITY_GROUP_ID }} + ALIYUN_ECS_INSTANCE_TYPE: ${{ vars.ALIYUN_ECS_INSTANCE_TYPE }} + ALIYUN_ECS_RESOURCE_GROUP_ID: ${{ vars.ALIYUN_ECS_RESOURCE_GROUP_ID }} + QUERY_REGRESSION_ECS_IMAGE_ID: ${{ vars.QUERY_REGRESSION_ECS_IMAGE_ID }} + QUERY_REGRESSION_RUNNER_UID: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }} + QUERY_REGRESSION_RUNNER_GID: ${{ vars.QUERY_REGRESSION_RUNNER_GID || '1001' }} + run: >- + uv run .github/scripts/aliyun-ecs-runner-provision.py + + query-regression: + needs: [provision, test-tooling] + # `!failure() && !cancelled()` both suppresses the implicit success() and + # lets the job run when provision was intentionally skipped because a + # literal runner label was selected; a failed or cancelled provision still + # blocks the run because no ECS runner would be waiting. + if: >- + ${{ !failure() && !cancelled() && + (github.event_name != 'pull_request' || + (github.event_name == 'pull_request' && + !github.event.pull_request.draft && + (github.event.label.name == 'query-regression' || + github.event.label.name == 'heavy-regression'))) }} + runs-on: >- + ${{ needs.provision.outputs.label || + (github.event_name != 'pull_request' && inputs.runner || + (vars.QUERY_REGRESSION_PR_RUNNER || 'aliyun-ecs')) }} timeout-minutes: 180 - concurrency: - group: query-regression-persistent-cache-v1 - queue: max - cancel-in-progress: false env: CARGO_PROFILE: ${{ github.event_name == 'pull_request' && 'nightly' || inputs.cargo_profile }} CARGO_HOME: /home/runner/.cargo @@ -105,7 +193,8 @@ jobs: QUERY_REGRESSION_CACHE_META: /home/runner/query-regression-cache-meta RUSTC_WRAPPER: /usr/local/bin/sccache SCCACHE_DIR: /home/runner/.cache/sccache - SCCACHE_CACHE_SIZE: 40G + # Caps the local sccache on the system disk next to target and cargo. + SCCACHE_CACHE_SIZE: 10G CARGO_INCREMENTAL: "0" RUSTFLAGS: -D warnings -C link-arg=-fuse-ld=mold QUERY_REGRESSION_CACHE_EPOCH: "1" @@ -114,7 +203,55 @@ jobs: EVENT_MERGE_SHA: ${{ github.event_name == 'pull_request' && github.sha || '' }} EVENT_HEAD_SHA: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || '' }} EVENT_BASE_SHA: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || '' }} + # Runner identity contract. The ECS image uses 1001; a manually prepared + # host may override via repo variables when 1001 is already taken. + EXPECTED_RUNNER_UID: ${{ vars.QUERY_REGRESSION_RUNNER_UID || '1001' }} + EXPECTED_RUNNER_GID: ${{ vars.QUERY_REGRESSION_RUNNER_GID || '1001' }} steps: + - name: Report provisioned ECS runner + if: ${{ needs.provision.outputs.instance_id != '' }} + shell: bash + run: | + set -euo pipefail + { + printf -- '- ECS instance: `%s` (type `%s`, image `%s`, region `%s`)\n' \ + "${{ needs.provision.outputs.instance_id }}" \ + "${{ vars.ALIYUN_ECS_INSTANCE_TYPE }}" \ + "${{ vars.QUERY_REGRESSION_ECS_IMAGE_ID }}" \ + "${{ vars.ALIYUN_ECS_REGION_ID }}" + } >> "${GITHUB_STEP_SUMMARY}" + + - name: Start machine telemetry sampler + shell: bash + # The ECS instance is deleted when the run ends, so machine state must + # be captured while the job runs. A background sampler appends load, + # memory, disk, and top-memory processes to a log every 30s; the + # "Dump machine telemetry" step (if: always()) prints and uploads it. + # Note: when the job is *cancelled*, even always() steps are killed — + # telemetry covers failures, not cancellations. + run: | + set -euo pipefail + log="${GITHUB_WORKSPACE}/machine-telemetry.log" + { + echo "== baseline $(date -u +%Y-%m-%dT%H:%M:%SZ) ==" + nproc + free -h + df -h / /home/runner + } >> "${log}" + nohup bash -c " + while true; do + { + date -u '+== %Y-%m-%dT%H:%M:%SZ ==' + uptime + free -m + df -h / /home/runner + ps -eo pid,comm,%mem,%cpu --sort=-%mem | head -6 + } >> '${log}' 2>&1 + sleep 30 + done + " >/dev/null 2>&1 & + echo "Telemetry sampler started (pid $!), logging to ${log}" + - name: Checkout base source uses: actions/checkout@v4 with: @@ -263,14 +400,64 @@ jobs: shell: bash run: | set -euo pipefail - [[ "$(id -u)" == "1001" ]] - [[ "$(id -g)" == "1001" ]] - [[ "${UV_CACHE_DIR}" == "/home/runner/.cargo/uv-cache" ]] - [[ "$(protoc --version)" == "libprotoc 3.21.12" ]] + errors=() + record() { + printf 'FAIL %s\n' "$*" + errors+=("$*") + } + require_eq() { + local name="$1" actual="$2" expected="$3" + printf 'check %s: %s\n' "${name}" "${actual}" + [[ "${actual}" == "${expected}" ]] || record "${name}: expected '${expected}', got '${actual}'" + } + require_match() { + local name="$1" actual="$2" pattern="$3" + printf 'check %s: %s\n' "${name}" "${actual}" + [[ "${actual}" =~ $pattern ]] || record "${name}: expected to match ${pattern}, got '${actual}'" + } + require() { + local name="$1" + shift + printf 'check %s\n' "${name}" + "$@" || record "${name}" + } + capture() { + local out + if out="$("$@" 2>&1)"; then + printf '%s' "${out}" + else + printf '' "$*" + fi + } + + printf 'uid=%s gid=%s PATH=%s CARGO_HOME=%s\n' "$(id -u)" "$(id -g)" "${PATH}" "${CARGO_HOME}" + ls -la "${CARGO_HOME}" 2>&1 || printf '(CARGO_HOME missing)\n' + + # Image/host hygiene: CARGO_HOME must be empty of config and + # install-state before cargo/rustup run. Those tools create + # .package-cache / bin on first use, so this cannot come after. + cargo_home_dirty=() + for entry in config config.toml credentials credentials.toml bin .crates.toml .crates2.json .global-cache .package-cache; do + if [[ -e "${CARGO_HOME}/${entry}" ]]; then + cargo_home_dirty+=("${entry}") + record "CARGO_HOME must not contain ${entry}" + fi + done + if (( ${#cargo_home_dirty[@]} > 0 )); then + printf 'CARGO_HOME contents:\n' >&2 + ls -la "${CARGO_HOME}" >&2 || true + fi + + require_eq uid "$(id -u)" "${EXPECTED_RUNNER_UID}" + require_eq gid "$(id -g)" "${EXPECTED_RUNNER_GID}" + require_eq UV_CACHE_DIR "${UV_CACHE_DIR}" "/home/runner/.cargo/uv-cache" + require_eq protoc "$(capture protoc --version)" "libprotoc 3.21.12" temporary_proto_dir="$(mktemp --directory)" trap 'rm -rf "${temporary_proto_dir}"' EXIT - test -r /usr/include/google/protobuf/any.proto - test -r /usr/include/google/protobuf/empty.proto + require "readable /usr/include/google/protobuf/any.proto" \ + test -r /usr/include/google/protobuf/any.proto + require "readable /usr/include/google/protobuf/empty.proto" \ + test -r /usr/include/google/protobuf/empty.proto printf '%s\n' \ 'syntax = "proto3";' \ 'package smoke;' \ @@ -278,47 +465,51 @@ jobs: 'import "google/protobuf/empty.proto";' \ 'message Smoke { google.protobuf.Any any = 1; google.protobuf.Empty empty = 2; }' \ > "${temporary_proto_dir}/smoke.proto" - protoc --proto_path="${temporary_proto_dir}" --proto_path=/usr/include \ + if protoc --proto_path="${temporary_proto_dir}" --proto_path=/usr/include \ --descriptor_set_out="${temporary_proto_dir}/smoke.pb" \ - "${temporary_proto_dir}/smoke.proto" - test -s "${temporary_proto_dir}/smoke.pb" - [[ "$(uv --version)" =~ ^uv[[:space:]]0\.11\.26([[:space:]]|$) ]] - mold_version="$(mold --version)" - [[ "${mold_version}" =~ ^mold[[:space:]]2\.30\.0([[:space:]]|$) ]] - [[ "$(python3 --version)" == "Python 3.12.3" ]] + "${temporary_proto_dir}/smoke.proto"; then + require "protoc smoke descriptor is non-empty" test -s "${temporary_proto_dir}/smoke.pb" + else + record "protoc smoke compile failed" + fi + require_match uv "$(capture uv --version)" '^uv[[:space:]]0\.11\.26([[:space:]]|$)' + require_match mold "$(capture mold --version)" '^mold[[:space:]]2\.40\.4([[:space:]]|$)' + require_eq python3 "$(capture python3 --version)" "Python 3.14.4" otelgen_path="$(command -v otelgen || true)" otelgen_version="" if [[ -n "${otelgen_path}" ]]; then - otelgen_version="$("${otelgen_path}" --version 2>&1 || true)" + otelgen_version="$(capture "${otelgen_path}" --version)" fi - printf 'otelgen path: %s\n' "${otelgen_path:-}" - printf 'otelgen version: %s\n' "${otelgen_version}" - [[ "${otelgen_path}" == "/usr/local/bin/otelgen" ]] - [[ "${otelgen_version}" == *"863a3f395d062c7322cc1de08a38774b7fdaa6c8"* ]] - sccache_version="$(sccache --version)" - [[ "${sccache_version}" =~ ^sccache[[:space:]]0\.16\.0([[:space:]]|$) ]] - [[ "$(command -v rustup)" == "/opt/cargo/bin/rustup" ]] - [[ "$(command -v cargo)" == "/opt/cargo/bin/cargo" ]] - [[ "$(command -v rustc)" == "/opt/cargo/bin/rustc" ]] - [[ "$(rustup --version)" =~ ^rustup[[:space:]]1\.29\.0([[:space:]]|$) ]] - cargo_version="$(cargo --version)" - [[ "${cargo_version}" =~ ^cargo[[:space:]]1\.96\.0-nightly[[:space:]]\(cbb9bb8bd[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$ ]] - rustc_version="$(rustc --version)" - [[ "${rustc_version}" =~ ^rustc[[:space:]]1\.96\.0-nightly[[:space:]]\(ac7f9ec7d[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$ ]] - active_toolchain="$(rustup show active-toolchain)" - [[ "${active_toolchain}" =~ ^nightly-2026-03-21-x86_64-unknown-linux-gnu([[:space:]]|$) ]] - [[ "${RUSTUP_HOME}" == "/opt/rustup" ]] - [[ "${RUSTUP_TOOLCHAIN}" == "nightly-2026-03-21" ]] - [[ "${RUSTUP_AUTO_INSTALL}" == "0" ]] - test -r /opt/rustup && test -x /opt/rustup - test ! -w /opt/rustup - test ! -w /opt/cargo/bin - for entry in config config.toml credentials credentials.toml bin .crates.toml .crates2.json .global-cache .package-cache; do - test ! -e "${CARGO_HOME}/${entry}" - done + require_eq otelgen_path "${otelgen_path}" "/usr/local/bin/otelgen" + printf 'check otelgen_version: %s\n' "${otelgen_version}" + [[ "${otelgen_version}" == *"863a3f395d062c7322cc1de08a38774b7fdaa6c8"* ]] \ + || record "otelgen_version: expected commit 863a3f395d062c7322cc1de08a38774b7fdaa6c8, got '${otelgen_version}'" + require_match sccache "$(capture sccache --version)" '^sccache[[:space:]]0\.16\.0([[:space:]]|$)' + require_eq rustup_path "$(command -v rustup || true)" "/opt/cargo/bin/rustup" + require_eq cargo_path "$(command -v cargo || true)" "/opt/cargo/bin/cargo" + require_eq rustc_path "$(command -v rustc || true)" "/opt/cargo/bin/rustc" + require_match rustup "$(capture rustup --version)" '^rustup[[:space:]]1\.29\.0([[:space:]]|$)' + require_match cargo "$(capture cargo --version)" \ + '^cargo[[:space:]]1\.96\.0-nightly[[:space:]]\(cbb9bb8bd[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$' + require_match rustc "$(capture rustc --version)" \ + '^rustc[[:space:]]1\.96\.0-nightly[[:space:]]\(ac7f9ec7d[[:space:]][0-9]{4}-[0-9]{2}-[0-9]{2}\)$' + require_match active_toolchain "$(capture rustup show active-toolchain)" \ + '^nightly-2026-03-21-x86_64-unknown-linux-gnu([[:space:]]|$)' + require_eq RUSTUP_HOME "${RUSTUP_HOME}" "/opt/rustup" + require_eq RUSTUP_TOOLCHAIN "${RUSTUP_TOOLCHAIN}" "nightly-2026-03-21" + require_eq RUSTUP_AUTO_INSTALL "${RUSTUP_AUTO_INSTALL}" "0" + require "readable /opt/rustup" test -r /opt/rustup + require "executable /opt/rustup" test -x /opt/rustup + require "runner cannot write /opt/rustup" test ! -w /opt/rustup + require "runner cannot write /opt/cargo/bin" test ! -w /opt/cargo/bin mkdir -p "${CARGO_HOME}/registry" "${CARGO_HOME}/git" + if (( ${#errors[@]} > 0 )); then + printf '\nVerify runner image tools failed (%d checks):\n' "${#errors[@]}" >&2 + printf ' - %s\n' "${errors[@]}" >&2 + exit 1 + fi - - name: Prepare persistent query regression cache + - name: Prepare query regression cache shell: bash working-directory: src run: | @@ -362,24 +553,6 @@ jobs: find "${root}" -mindepth 1 -maxdepth 1 -exec rm -rf -- {} + } - clear_cargo_extracted_trees() { - require_expected_root CARGO_REGISTRY "${CARGO_HOME}/registry" "${EXPECTED_CARGO_REGISTRY}" - require_expected_root CARGO_GIT "${CARGO_HOME}/git" "${EXPECTED_CARGO_GIT}" - rm -rf -- "${CARGO_HOME}/registry/src" "${CARGO_HOME}/git/checkouts" - } - - cargo_size_kib() { - du -sk -- "${CARGO_HOME}/registry" "${CARGO_HOME}/git" | awk '{ total += $1 } END { print total }' - } - - target_size_kib() { - du -sk -- "${CARGO_TARGET_DIR}" | cut -f1 - } - - free_kib() { - df -Pk "${CARGO_TARGET_DIR}" | awk 'NR == 2 { print $4 }' - } - report_cache_usage() { du -sh -- "${CARGO_HOME}" "${CARGO_HOME}/registry" "${CARGO_HOME}/git" \ "${RUSTUP_HOME}" "${CARGO_TARGET_DIR}" "${QUERY_REGRESSION_CACHE_META}" "${SCCACHE_DIR}" @@ -400,7 +573,7 @@ jobs: printf 'Refusing unexpected RUSTC_WRAPPER: %s\n' "${RUSTC_WRAPPER}" >&2 exit 1 } - [[ "${SCCACHE_CACHE_SIZE}" == "40G" ]] || { + [[ "${SCCACHE_CACHE_SIZE}" == "10G" ]] || { printf 'Refusing unexpected SCCACHE_CACHE_SIZE: %s\n' "${SCCACHE_CACHE_SIZE}" >&2 exit 1 } @@ -487,46 +660,10 @@ jobs: report_cache_usage - target_size="$(target_size_kib)" - if (( target_size >= 400 * 1024 * 1024 )); then - printf 'Warning: target cache is at least 400 GiB (%s KiB)\n' "${target_size}" >&2 - fi - if (( target_size >= 450 * 1024 * 1024 )); then - printf 'Target cache is at least 450 GiB; clearing complete target root\n' >&2 - clear_directory "${CARGO_TARGET_DIR}" - fi - - cargo_size="$(cargo_size_kib)" - if (( cargo_size >= 60 * 1024 * 1024 )); then - printf 'Warning: Cargo cache is at least 60 GiB (%s KiB)\n' "${cargo_size}" >&2 - fi - if (( cargo_size >= 80 * 1024 * 1024 )); then - printf 'Cargo cache is at least 80 GiB; removing extracted sources and checkouts\n' >&2 - clear_cargo_extracted_trees - cargo_size="$(cargo_size_kib)" - if (( cargo_size >= 80 * 1024 * 1024 )); then - printf 'Cargo cache remains at least 80 GiB after cleanup (%s KiB)\n' "${cargo_size}" >&2 - exit 1 - fi - fi - - free_space="$(free_kib)" - if (( free_space < 300 * 1024 * 1024 )); then - printf 'Free space is below 300 GiB; clearing complete target root\n' >&2 - clear_directory "${CARGO_TARGET_DIR}" - free_space="$(free_kib)" - if (( free_space < 300 * 1024 * 1024 )); then - printf 'Free space remains below 300 GiB; removing Cargo extracted sources and checkouts\n' >&2 - clear_cargo_extracted_trees - free_space="$(free_kib)" - if (( free_space < 300 * 1024 * 1024 )); then - printf 'Free space remains below 300 GiB after cleanup (%s KiB)\n' "${free_space}" >&2 - exit 1 - fi - fi - fi - - report_cache_usage + # Every run starts from a fresh system disk, so size/free-space + # watermarks from the retired retained-disk era are not restored. + # A run that overflows the disk fails the build outright, which the + # telemetry step makes diagnosable. sccache --start-server sccache --zero-stats @@ -567,14 +704,6 @@ jobs: git reset --hard "${VERIFIED_CANDIDATE_SHA}" git clean -ffdx - - name: Test query regression tooling - working-directory: src - run: | - uv run --no-project python tests/perf/test_query_regression_runner_compaction_toctou.py - uv run --no-project python tests/perf/test_query_regression_runner_otlp_trace_load.py - uv run --no-project python tests/perf/test_query_regression_summary_otlp.py - uv run --no-project python tests/perf/test_query_regression_case_selection.py - - name: Build candidate greptime and query regression helpers working-directory: src run: | @@ -643,6 +772,7 @@ jobs: query-regression-work/**/logs/** query-regression-work/**/otelgen/** query-regression-summary.md + machine-telemetry.log if-no-files-found: warn retention-days: 7 @@ -657,7 +787,7 @@ jobs: if-no-files-found: warn retention-days: 7 - - name: Report persistent cache usage + - name: Report cache usage if: ${{ always() }} shell: bash run: | @@ -697,6 +827,55 @@ jobs: printf 'sccache is unavailable (report only)\n' >&2 fi + - name: Dump machine telemetry + if: ${{ always() }} + shell: bash + run: | + log="${GITHUB_WORKSPACE}/machine-telemetry.log" + if [[ -f "${log}" ]]; then + echo "::group::Machine telemetry (last 200 lines)" + tail -n 200 "${log}" + echo "::endgroup::" + else + echo "No telemetry log found (sampler never started?)" + fi + echo "::group::dmesg tail (OOM killer records)" + sudo dmesg -T 2>/dev/null | tail -n 50 || dmesg -T 2>/dev/null | tail -n 50 || \ + echo "dmesg unavailable without root" + echo "::endgroup::" + - name: Fail on regression failure if: ${{ steps.run.outputs.status != '0' }} run: exit 1 + + teardown: + # Releases the dynamically provisioned ECS runner. Runs even when the + # benchmark job fails or is cancelled; skipped when a literal runner label + # was selected because the provision outputs are empty, or when the + # dispatch set keep_instance to preserve the machine for post-mortem + # debugging (the janitor sweep still reclaims it after the TTL). + if: ${{ always() && needs.provision.outputs.instance_id != '' && !inputs.keep_instance }} + needs: [provision, query-regression] + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - name: Checkout trusted teardown scripts + uses: actions/checkout@v4 + with: + repository: ${{ github.repository }} + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || github.sha }} + persist-credentials: false + + - name: Install uv + uses: astral-sh/setup-uv@v6 + + - name: Teardown ECS runner + env: + ALIBABA_CLOUD_ACCESS_KEY_ID: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_ID }} + ALIBABA_CLOUD_ACCESS_KEY_SECRET: ${{ secrets.ALICLOUD_ECS_ACCESS_KEY_SECRET }} + GH_PERSONAL_ACCESS_TOKEN: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }} + ALIYUN_ECS_REGION_ID: ${{ vars.ALIYUN_ECS_REGION_ID }} + QUERY_REGRESSION_ECS_INSTANCE_ID: ${{ needs.provision.outputs.instance_id }} + QUERY_REGRESSION_ECS_RUNNER_NAME: ${{ needs.provision.outputs.runner_name }} + run: >- + uv run .github/scripts/aliyun-ecs-runner-teardown.py diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 52d62299a07..0a9c39c9034 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -510,7 +510,7 @@ jobs: base_ref: ${{ needs.prepare-release-validation.outputs.previous-release-tag }} candidate_ref: ${{ needs.prepare-release-validation.outputs.candidate-ref }} cargo_profile: nightly - runner: perf-regression-8-cores + runner: aliyun-ecs release-images-to-dockerhub: name: Build and push images to DockerHub diff --git a/tests/perf/AGENTS.md b/tests/perf/AGENTS.md index 1f13dea58dc..822b99080f9 100644 --- a/tests/perf/AGENTS.md +++ b/tests/perf/AGENTS.md @@ -3,9 +3,23 @@ - Keep GitHub Actions YAML thin. Put non-trivial control flow, case expansion, report generation, and metadata writing in scripts under `.github/scripts/`; workflow steps should mostly invoke those scripts. +- Runner lifecycle: the default path provisions one ephemeral Aliyun ECS + instance per run via `.github/scripts/aliyun-ecs-runner-provision.py` and + always releases it via `aliyun-ecs-runner-teardown.py`; a scheduled janitor + workflow sweeps leftovers. Build caches live on that instance's system disk + and are discarded with the VM. Runs do not share a workflow concurrency + group. The ECS custom image is built from the runner Dockerfile by + `.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py`; + keep the Dockerfile the single source of the tool contract. Dispatching with + any other `runner` value treats it as a literal self-hosted runner label + (see `ecs-image/bootstrap-runner-host.sh` for preparing such a host). - Query regression PR runs should build base/candidate binaries once, then run the default case set. Do not hard-code a single case such as `promql_pushdown_7913` into the workflow path. +- Scheduled nightly comparison lives in `query-regression-nightly.yml`: it + waits for a successful Nightly Build, then calls `query-regression.yml` + with the previous vs current nightly SHAs. Keep SHA selection in + `.github/scripts/query-regression-nightly-refs.py`. - The case DSL is not required to keep compatibility inside this PR. When the DSL changes, update TOML cases, the outer lifecycle script, Rust helpers, and docs together. @@ -19,8 +33,10 @@ - Keep the direct-SST generator generic. Issue-specific behavior belongs in case files and thresholds, not in Rust generator logic. - Before pushing perf harness changes, run at least: - - the Python tests in the `Test query regression tooling` step of - `.github/workflows/query-regression.yml` + - the Python tests in the `test-tooling` job of + `.github/workflows/query-regression.yml` (ubuntu-latest, not the ECS runner). + The Checks workflow runs the same tests on ordinary PRs so they are not + gated on the `query-regression` / `heavy-regression` labels. - `cargo fmt --all -- --check` - `cargo build -p cmd --bin query_perf_fixture --features dev-tools` - `cargo build -p cmd --bin query_regression_runner --features dev-tools` diff --git a/tests/perf/README.md b/tests/perf/README.md index cbe03b51794..47495ca531f 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -255,11 +255,12 @@ gh workflow run query-regression.yml \ -f candidate_ref= \ -f cargo_profile=nightly \ -f http_timeout=300 \ - -f runner=perf-regression-8-cores + -f runner=aliyun-ecs ``` -The selected ARC scale set must already be deployed with the runner-image -digest built from the current query-regression Dockerfile. +The `aliyun-ecs` path provisions a fresh ECS instance per run from the custom +image built from the current query-regression Dockerfile; see +`.github/runner-scale-sets/query-regression/README.md` for its configuration. ## Generator contract diff --git a/tests/perf/test_aliyun_ecs_runner_scripts.py b/tests/perf/test_aliyun_ecs_runner_scripts.py new file mode 100644 index 00000000000..581fc88abd0 --- /dev/null +++ b/tests/perf/test_aliyun_ecs_runner_scripts.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# 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. + +"""Coverage for the pure parts of the Aliyun ECS runner provision/teardown scripts.""" + +import base64 +import importlib.util +import sys +import unittest +from datetime import datetime, timedelta, timezone +from pathlib import Path + + +SCRIPTS_DIR = Path(__file__).parents[2] / ".github/scripts" + + +def load_module(name: str, filename: str): + spec = importlib.util.spec_from_file_location(name, SCRIPTS_DIR / filename) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +provision = load_module("aliyun_ecs_runner_provision_under_test", "aliyun-ecs-runner-provision.py") +teardown = load_module("aliyun_ecs_runner_teardown_under_test", "aliyun-ecs-runner-teardown.py") + + +class ProvisionNamingTest(unittest.TestCase): + def test_runner_name_and_label_derive_from_run_id(self) -> None: + self.assertEqual(provision.runner_name_for_run("12345"), "qreg-ecs-12345") + self.assertEqual(provision.runner_label_for_run("12345"), "query-regression-ecs-12345") + + +class ProvisionUserDataTest(unittest.TestCase): + def render(self) -> str: + return provision.render_user_data( + runner_name="qreg-ecs-12345", + runner_label="query-regression-ecs-12345", + runner_token="TOKEN", + repo="GreptimeTeam/greptimedb", + ) + + def test_user_data_creates_cache_paths_on_the_system_disk(self) -> None: + script = self.render() + self.assertNotIn("DISK_SERIAL", script) + self.assertNotIn("mount --bind", script) + self.assertNotIn("mkfs.ext4", script) + for destination in provision.CACHE_PATHS: + self.assertIn(f'"{destination}"', script) + + def test_user_data_wires_runner_registration(self) -> None: + script = self.render() + self.assertIn("RUNNER_NAME=qreg-ecs-12345", script) + self.assertIn("RUNNER_LABELS=query-regression-ecs-12345", script) + self.assertIn("RUNNER_TOKEN=TOKEN", script) + self.assertIn("REPO_URL=https://github.com/GreptimeTeam/greptimedb", script) + self.assertIn("PATH=/opt/cargo/bin:", script) + self.assertIn("systemctl restart --no-block ephemeral-github-runner.service", script) + + def test_encode_user_data_round_trips(self) -> None: + script = self.render() + self.assertEqual( + base64.b64decode(provision.encode_user_data(script)).decode("utf-8"), script + ) + + def test_user_data_enables_swap_and_masks_oomd(self) -> None: + script = self.render() + self.assertIn("systemctl mask systemd-oomd.socket systemd-oomd.service", script) + self.assertIn("systemctl mask unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer", script) + self.assertIn('APT::Periodic::Unattended-Upgrade "0"', script) + self.assertIn(f'fallocate --length {provision.SWAP_SIZE_GIB}G "{provision.SWAP_FILE}"', script) + self.assertIn(f'swapon "{provision.SWAP_FILE}"', script) + self.assertIn("sysctl --write vm.swappiness=10", script) + self.assertIn("OOMPolicy=continue", script) + self.assertNotIn("OOMScoreAdjust", script) + self.assertLess(script.index("swapon"), script.index("systemctl restart --no-block ephemeral-github-runner.service")) + + +class TeardownExpiryTest(unittest.TestCase): + NOW = datetime(2026, 8, 17, 6, 0, tzinfo=timezone.utc) + TTL = timedelta(hours=4) + + def test_parse_creation_time_formats(self) -> None: + self.assertEqual( + teardown.parse_creation_time("2026-08-17T01:02:03Z"), + datetime(2026, 8, 17, 1, 2, 3, tzinfo=timezone.utc), + ) + self.assertEqual( + teardown.parse_creation_time("2026-08-17T01:02Z"), + datetime(2026, 8, 17, 1, 2, tzinfo=timezone.utc), + ) + with self.assertRaises(ValueError): + teardown.parse_creation_time("not-a-time") + + def test_expired_instance_names_selects_only_old_instances(self) -> None: + instances = [ + ("i-old", "qreg-ecs-1", "2026-08-17T01:00Z"), # 5h old: expired + ("i-edge", "qreg-ecs-2", "2026-08-17T02:00Z"), # exactly TTL: expired + ("i-fresh", "qreg-ecs-3", "2026-08-17T05:30Z"), # 30m old: kept + ] + self.assertEqual( + teardown.expired_instance_names(instances, self.NOW, self.TTL), + [("i-old", "qreg-ecs-1"), ("i-edge", "qreg-ecs-2")], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/perf/test_query_regression_nightly_refs.py b/tests/perf/test_query_regression_nightly_refs.py new file mode 100644 index 00000000000..74b168bb464 --- /dev/null +++ b/tests/perf/test_query_regression_nightly_refs.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 +# Copyright 2023 Greptime Team +# +# 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. + +"""Unit tests for consecutive Nightly Build SHA selection.""" + +import importlib.util +import sys +import unittest +from pathlib import Path + + +SCRIPTS_DIR = Path(__file__).parents[2] / ".github/scripts" + + +def load_module(): + spec = importlib.util.spec_from_file_location( + "query_regression_nightly_refs_under_test", + SCRIPTS_DIR / "query-regression-nightly-refs.py", + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +refs = load_module() + + +def run(*, run_id: int, sha: str, branch: str = "main", created_at: str = "2026-08-24T00:00:00Z"): + return refs.WorkflowRun( + id=run_id, + head_sha=sha, + head_branch=branch, + html_url=f"https://github.com/example/run/{run_id}", + created_at=created_at, + conclusion="success", + event="schedule", + ) + + +class SelectNightlyRefsTest(unittest.TestCase): + def test_picks_newest_and_previous_on_same_branch(self) -> None: + pair = refs.select_base_and_candidate( + [ + run(run_id=3, sha="ccc", created_at="2026-08-22T00:00:00Z"), + run(run_id=2, sha="bbb", created_at="2026-08-21T00:00:00Z"), + run(run_id=1, sha="aaa", created_at="2026-08-20T00:00:00Z"), + ] + ) + self.assertFalse(pair.skip) + assert pair.candidate is not None and pair.base is not None + self.assertEqual(pair.candidate.head_sha, "ccc") + self.assertEqual(pair.base.head_sha, "bbb") + + def test_candidate_run_id_uses_that_run_and_the_next_older(self) -> None: + pair = refs.select_base_and_candidate( + [ + run(run_id=3, sha="ccc", created_at="2026-08-22T00:00:00Z"), + run(run_id=2, sha="bbb", created_at="2026-08-21T00:00:00Z"), + run(run_id=1, sha="aaa", created_at="2026-08-20T00:00:00Z"), + ], + candidate_run_id=2, + ) + self.assertFalse(pair.skip) + assert pair.candidate is not None and pair.base is not None + self.assertEqual(pair.candidate.id, 2) + self.assertEqual(pair.base.id, 1) + + def test_skips_when_consecutive_nightlies_share_a_sha(self) -> None: + pair = refs.select_base_and_candidate( + [ + run(run_id=2, sha="same", created_at="2026-08-22T00:00:00Z"), + run(run_id=1, sha="same", created_at="2026-08-21T00:00:00Z"), + ] + ) + self.assertTrue(pair.skip) + self.assertIn("matches candidate", pair.reason) + + def test_skips_other_branches_when_picking_previous(self) -> None: + pair = refs.select_base_and_candidate( + [ + run(run_id=3, sha="ccc", branch="main", created_at="2026-08-22T00:00:00Z"), + run(run_id=2, sha="other", branch="feat", created_at="2026-08-21T12:00:00Z"), + run(run_id=1, sha="aaa", branch="main", created_at="2026-08-21T00:00:00Z"), + ] + ) + self.assertFalse(pair.skip) + assert pair.base is not None + self.assertEqual(pair.base.head_sha, "aaa") + + def test_skips_when_only_one_nightly_exists(self) -> None: + pair = refs.select_base_and_candidate( + [run(run_id=1, sha="aaa")], + ) + self.assertTrue(pair.skip) + self.assertIn("no previous", pair.reason) + + def test_explicit_refs_bypass_github(self) -> None: + pair = refs.override_pair("base-sha", "cand-sha") + self.assertFalse(pair.skip) + assert pair.base is not None and pair.candidate is not None + self.assertEqual(pair.base.head_sha, "base-sha") + self.assertEqual(pair.candidate.head_sha, "cand-sha") + + +if __name__ == "__main__": + unittest.main() From 1409e66837d89ff7f9d7d0604dfdbba3ca9fc7a7 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 26 Aug 2026 14:00:39 +0000 Subject: [PATCH 88/92] refactor(flight): add request builder and defer DoGet execution (#8953) * refactor: add flight request builder Signed-off-by: WenyXu * refactor: use flight request builder in flow Signed-off-by: WenyXu * refactor: defer frontend flight query execution Signed-off-by: WenyXu * fix(grpc): avoid cloning requests during auth Signed-off-by: WenyXu * test(grpc): cover flight request timeout Signed-off-by: WenyXu * refactor(client): share Flight message reader Signed-off-by: WenyXu * feat(flow): add Flight DoGet timeout Signed-off-by: WenyXu * refactor(client): gate Flight DDL helpers for testing Signed-off-by: WenyXu * fix(client): restore Flight stream error semantics Signed-off-by: WenyXu * docs(grpc): document Flight stream input constructors Signed-off-by: WenyXu * fix(flight): preserve deferred stream context Signed-off-by: WenyXu * fix(client): use Flight stream SNAFU context Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/client/src/database.rs | 490 ++++++++++++------ src/client/src/error.rs | 18 + src/client/src/flight.rs | 74 ++- src/client/src/region.rs | 189 +++++-- src/datanode/src/region_server.rs | 9 +- src/flow/src/batching_mode.rs | 3 + src/flow/src/batching_mode/frontend_client.rs | 61 +-- src/servers/src/grpc/flight.rs | 75 +-- src/servers/src/grpc/flight/stream.rs | 150 ++++-- src/servers/src/grpc/greptime_handler.rs | 72 +-- tests-integration/src/grpc/flight.rs | 180 ++++++- tests-integration/tests/http.rs | 1 + 12 files changed, 923 insertions(+), 399 deletions(-) diff --git a/src/client/src/database.rs b/src/client/src/database.rs index 4fff3636b63..7c33794f477 100644 --- a/src/client/src/database.rs +++ b/src/client/src/database.rs @@ -12,20 +12,25 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; use std::pin::Pin; use std::str::FromStr; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, RwLock}; use std::task::{Context, Poll}; +use std::time::Duration; use api::v1::auth_header::AuthScheme; +#[cfg(feature = "testing")] use api::v1::ddl_request::Expr as DdlExpr; use api::v1::greptime_database_client::GreptimeDatabaseClient; use api::v1::greptime_request::Request; use api::v1::query_request::Query; +#[cfg(feature = "testing")] +use api::v1::{AlterTableExpr, CreateTableExpr, DdlRequest}; use api::v1::{ - AlterTableExpr, AuthHeader, Basic, CreateTableExpr, DdlRequest, GreptimeRequest, - InsertRequests, QueryRequest, RequestHeader, RowInsertRequests, + AuthHeader, Basic, GreptimeRequest, InsertRequests, QueryRequest, RequestHeader, + RowInsertRequests, }; use arc_swap::ArcSwapOption; use arrow_flight::{FlightData, Ticket}; @@ -34,7 +39,7 @@ use base64::Engine; use base64::prelude::BASE64_STANDARD; use common_catalog::build_db_string; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; -use common_error::ext::BoxedError; +use common_error::ext::{BoxedError, ErrorExt}; use common_grpc::flight::do_put::DoPutResponse; use common_grpc::flight::{ FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightMessage, SNAPSHOT_SEQS_METADATA_KEY, @@ -49,21 +54,23 @@ use common_telemetry::{error, warn}; use futures::future; use futures_util::{Stream, StreamExt, TryStreamExt}; use prost::Message; -use snafu::ResultExt; +use snafu::{IntoError, ResultExt}; use tonic::metadata::{AsciiMetadataKey, AsciiMetadataValue, MetadataMap, MetadataValue}; use tonic::transport::Channel; use crate::error::{ - ConvertFlightDataSnafu, Error, FlightGetSnafu, IllegalFlightMessagesSnafu, + ConvertFlightDataSnafu, Error, FlightGetSnafu, FlightStreamSnafu, IllegalFlightMessagesSnafu, InvalidTonicMetadataValueSnafu, }; -use crate::flight::decode_flight_data; +use crate::flight::{FlightMessageReader, decode_flight_data}; use crate::{Client, Result, error, from_grpc_response}; type FlightDataStream = Pin + Send>>; type DoPutResponseStream = Pin>>>; +const HINTS_METADATA_KEY: &str = "x-greptime-hints"; + /// Terminal metrics associated with a query output. /// /// For streaming outputs, metrics are only final after the stream is fully @@ -257,19 +264,17 @@ fn attach_terminal_metrics(output: Output, terminal_metrics: &OutputMetrics) -> } async fn output_from_flight_message_stream( - mut flight_message_stream: S, + remote_addr: String, + flight_message_stream: S, ) -> Result where S: Stream> + Send + Unpin + 'static, { - let Some(first_flight_message) = flight_message_stream.next().await else { - return IllegalFlightMessagesSnafu { - reason: "Expect the response not to be empty", - } - .fail(); - }; - - let first_flight_message = first_flight_message?; + let mut reader = FlightMessageReader::new(remote_addr, flight_message_stream); + let first_flight_message = reader + .read_first() + .await + .map_err(|error| flight_stream_error(reader.remote_addr(), error))?; match first_flight_message { FlightMessage::AffectedRows { rows, metrics } => { @@ -277,7 +282,10 @@ where if let Some(metrics) = metrics { terminal_metrics.update(Some(parse_terminal_metrics(&metrics)?)); } - let next_message = flight_message_stream.next().await.transpose()?; + let next_message = reader + .read_next() + .await + .map_err(|error| flight_stream_error(reader.remote_addr(), error))?; match next_message { None => terminal_metrics.mark_ready(), Some(FlightMessage::Metrics(s)) if terminal_metrics.get().is_none() => { @@ -316,15 +324,19 @@ where ); let schema_cloned = schema.clone(); let stream = Box::pin(stream!({ - while let Some(flight_message_item) = flight_message_stream.next().await { - let flight_message = match flight_message_item { - Ok(message) => message, - Err(e) => { - yield Err(BoxedError::new(e)).context(ExternalSnafu); + loop { + let flight_message = match reader.read_next().await { + Ok(Some(message)) => message, + Ok(None) => break, + Err(error) => { + yield Err(BoxedError::new(flight_stream_error( + reader.remote_addr(), + error, + ))) + .context(ExternalSnafu); break; } }; - match flight_message { FlightMessage::RecordBatch(arrow_batch) => { yield Ok(RecordBatch::from_df_record_batch( @@ -371,6 +383,25 @@ where } } +fn flight_stream_error(addr: &str, error: Error) -> Error { + let tonic_code = error.tonic_code().unwrap_or(tonic::Code::Unknown); + let message = error.to_string(); + if error.status_code().should_log_error() { + error!( + error; "Failed to receive Flight data, addr: {}, code: {}", + addr, + tonic_code + ); + } + + FlightStreamSnafu { + addr: addr.to_string(), + tonic_code, + message, + } + .into_error(BoxedError::new(error)) +} + #[derive(Clone, Debug, Default)] pub struct Database { // The "catalog" and "schema" to be used in processing the requests at the server side. @@ -389,6 +420,42 @@ pub struct Database { ctx: FlightContext, } +#[derive(Default)] +struct FlightRequestOptions { + hints: Option, + flow_extensions: Option, + snapshot_seqs: Option, + timeout: Option, +} + +impl FlightRequestOptions { + fn apply_to(self, request: &mut tonic::Request) -> Result<()> { + let metadata = request.metadata_mut(); + if let Some(hints) = self.hints { + Database::put_metadata_value(metadata, HINTS_METADATA_KEY, hints)?; + } + if let Some(flow_extensions) = self.flow_extensions { + Database::put_metadata_value(metadata, FLOW_EXTENSIONS_METADATA_KEY, flow_extensions)?; + } + if let Some(snapshot_seqs) = self.snapshot_seqs { + Database::put_metadata_value(metadata, SNAPSHOT_SEQS_METADATA_KEY, snapshot_seqs)?; + } + if let Some(timeout) = self.timeout { + request.set_timeout(timeout); + } + Ok(()) + } +} + +/// A single Flight DoGet request to a [`Database`]. +/// +/// The builder carries request-scoped metadata and options. It does not modify +/// the underlying [`Database`], so its configuration cannot affect later RPCs. +pub struct DatabaseFlightRequest<'a> { + database: &'a Database, + options: FlightRequestOptions, +} + pub struct DatabaseClient { pub addr: String, pub inner: GreptimeDatabaseClient, @@ -483,6 +550,14 @@ impl Database { }); } + /// Creates a builder for a single Flight DoGet request. + pub fn flight_request(&self) -> DatabaseFlightRequest<'_> { + DatabaseFlightRequest { + database: self, + options: FlightRequestOptions::default(), + } + } + /// Make an InsertRequests request to the database. pub async fn insert(&self, requests: InsertRequests) -> Result { self.handle(Request::Inserts(requests)).await @@ -538,44 +613,31 @@ impl Database { } fn put_hints(metadata: &mut MetadataMap, hints: &[(&str, &str)]) -> Result<()> { - let Some(value) = hints - .iter() - .map(|(k, v)| format!("{}={}", k, v)) - .reduce(|a, b| format!("{},{}", a, b)) - else { + let Some(value) = Self::encode_hints(hints) else { return Ok(()); }; - let key = AsciiMetadataKey::from_static("x-greptime-hints"); - let value = AsciiMetadataValue::from_str(&value).context(InvalidTonicMetadataValueSnafu)?; - metadata.insert(key, value); - Ok(()) + Self::put_metadata_value(metadata, HINTS_METADATA_KEY, value) } - fn put_flow_extensions( - metadata: &mut MetadataMap, - flow_extensions: &[(&str, &str)], - ) -> Result<()> { - if flow_extensions.is_empty() { - return Ok(()); - } - - let value = serde_json::to_string(&flow_extensions.to_vec()) - .expect("flow extension pairs should serialize"); - Self::put_metadata_value(metadata, FLOW_EXTENSIONS_METADATA_KEY, value) + fn encode_hints(hints: &[(&str, &str)]) -> Option { + hints + .iter() + .map(|(k, v)| format!("{}={}", k, v)) + .reduce(|a, b| format!("{},{}", a, b)) } - fn put_snapshot_seqs( - metadata: &mut MetadataMap, - snapshot_seqs: &std::collections::HashMap, - ) -> Result<()> { - if snapshot_seqs.is_empty() { - return Ok(()); - } + fn encode_flow_extensions(flow_extensions: &[(&str, &str)]) -> Option { + (!flow_extensions.is_empty()).then(|| { + serde_json::to_string(&flow_extensions.to_vec()) + .expect("flow extension pairs should serialize") + }) + } - let value = - serde_json::to_string(snapshot_seqs).expect("snapshot sequence map should serialize"); - Self::put_metadata_value(metadata, SNAPSHOT_SEQS_METADATA_KEY, value) + fn encode_snapshot_seqs(snapshot_seqs: &HashMap) -> Option { + (!snapshot_seqs.is_empty()).then(|| { + serde_json::to_string(snapshot_seqs).expect("snapshot sequence map should serialize") + }) } fn put_metadata_value( @@ -673,7 +735,7 @@ impl Database { where S: AsRef, { - self.sql_with_hint(sql, &[]).await + self.flight_request().sql(sql).await } /// Executes a SQL query with optional hints for query optimization. @@ -681,12 +743,7 @@ impl Database { where S: AsRef, { - let request = Request::Query(QueryRequest { - query: Some(Query::Sql(sql.as_ref().to_string())), - }); - self.do_get(request, hints, &[], &Default::default()) - .await - .map(OutputWithMetrics::into_output) + self.flight_request().with_hints(hints).sql(sql).await } /// Executes a SQL query and returns the output with terminal metrics. @@ -701,77 +758,33 @@ impl Database { where S: AsRef, { - self.query_with_terminal_metrics_and_flow_extensions( - QueryRequest { - query: Some(Query::Sql(sql.as_ref().to_string())), - }, - hints, - &[], - &Default::default(), - ) - .await + self.flight_request() + .with_hints(hints) + .sql_with_terminal_metrics(sql) + .await } /// Executes a logical plan directly without SQL parsing. pub async fn logical_plan(&self, logical_plan: Vec) -> Result { - self.query_with_terminal_metrics_and_flow_extensions( - QueryRequest { - query: Some(Query::LogicalPlan(logical_plan)), - }, - &[], - &[], - &Default::default(), - ) - .await - .map(OutputWithMetrics::into_output) - } - - /// Executes a query and carries flow extensions through Flight metadata. - /// - /// This is the lower-level terminal-metrics API for Flow callers that need - /// to pass JSON-bearing flow extensions without going through hint metadata. - pub async fn query_with_terminal_metrics_and_flow_extensions( - &self, - request: QueryRequest, - hints: &[(&str, &str)], - flow_extensions: &[(&str, &str)], - snapshot_seqs: &std::collections::HashMap, - ) -> Result { - self.do_get( - Request::Query(request), - hints, - flow_extensions, - snapshot_seqs, - ) - .await + self.flight_request().logical_plan(logical_plan).await } /// Creates a new table using the provided table expression. + #[cfg(feature = "testing")] pub async fn create(&self, expr: CreateTableExpr) -> Result { - let request = Request::Ddl(DdlRequest { - expr: Some(DdlExpr::CreateTable(expr)), - }); - self.do_get(request, &[], &[], &Default::default()) - .await - .map(OutputWithMetrics::into_output) + self.flight_request().create(expr).await } /// Alters an existing table using the provided alter expression. + #[cfg(feature = "testing")] pub async fn alter(&self, expr: AlterTableExpr) -> Result { - let request = Request::Ddl(DdlRequest { - expr: Some(DdlExpr::AlterTable(expr)), - }); - self.do_get(request, &[], &[], &Default::default()) - .await - .map(OutputWithMetrics::into_output) + self.flight_request().alter(expr).await } async fn do_get( &self, request: Request, - hints: &[(&str, &str)], - flow_extensions: &[(&str, &str)], - snapshot_seqs: &std::collections::HashMap, + options: FlightRequestOptions, ) -> Result { let request = self.to_rpc_request(request); let request = Ticket { @@ -779,12 +792,10 @@ impl Database { }; let mut request = tonic::Request::new(request); - let metadata = request.metadata_mut(); - Self::put_hints(metadata, hints)?; - Self::put_flow_extensions(metadata, flow_extensions)?; - Self::put_snapshot_seqs(metadata, snapshot_seqs)?; + options.apply_to(&mut request)?; let mut client = self.client.make_flight_client(false, false)?; + let remote_addr = client.addr().to_string(); let response = client.mut_inner().do_get(request).await.or_else(|e| { let tonic_code = e.code(); @@ -796,7 +807,7 @@ impl Database { e ); Err(BoxedError::new(e)).with_context(|_| FlightGetSnafu { - addr: client.addr().to_string(), + addr: remote_addr.clone(), tonic_code, }) })?; @@ -808,7 +819,7 @@ impl Database { future::ready(decode_flight_data(&mut decoder, flight_data)) }); - output_from_flight_message_stream(flight_message_stream).await + output_from_flight_message_stream(remote_addr, flight_message_stream).await } /// Ingest a stream of [RecordBatch]es that belong to a table, using Arrow Flight's "`DoPut`" @@ -847,6 +858,98 @@ impl Database { } } +impl<'a> DatabaseFlightRequest<'a> { + /// Adds query optimization hints to this Flight request. + pub fn with_hints(mut self, hints: &[(&str, &str)]) -> Self { + self.options.hints = Database::encode_hints(hints); + self + } + + /// Adds Flow extensions to this Flight request. + pub fn with_flow_extensions(mut self, flow_extensions: &[(&str, &str)]) -> Self { + self.options.flow_extensions = Database::encode_flow_extensions(flow_extensions); + self + } + + /// Adds snapshot sequence fences to this Flight request. + pub fn with_snapshot_seqs(mut self, snapshot_seqs: &HashMap) -> Self { + self.options.snapshot_seqs = Database::encode_snapshot_seqs(snapshot_seqs); + self + } + + /// Sets a timeout for this Flight request only. + pub fn with_timeout(mut self, timeout: Duration) -> Self { + self.options.timeout = Some(timeout); + self + } + + /// Executes a SQL query. + pub async fn sql(self, sql: S) -> Result + where + S: AsRef, + { + let request = Request::Query(QueryRequest { + query: Some(Query::Sql(sql.as_ref().to_string())), + }); + self.do_get(request) + .await + .map(OutputWithMetrics::into_output) + } + + /// Executes a SQL query and returns terminal metrics. + pub async fn sql_with_terminal_metrics(self, sql: S) -> Result + where + S: AsRef, + { + self.query_with_terminal_metrics(QueryRequest { + query: Some(Query::Sql(sql.as_ref().to_string())), + }) + .await + } + + /// Executes a logical plan directly without SQL parsing. + pub async fn logical_plan(self, logical_plan: Vec) -> Result { + self.query_with_terminal_metrics(QueryRequest { + query: Some(Query::LogicalPlan(logical_plan)), + }) + .await + .map(OutputWithMetrics::into_output) + } + + /// Executes a query and returns terminal metrics. + pub async fn query_with_terminal_metrics( + self, + request: QueryRequest, + ) -> Result { + self.do_get(Request::Query(request)).await + } + + /// Creates a new table using the provided table expression. + #[cfg(feature = "testing")] + pub async fn create(self, expr: CreateTableExpr) -> Result { + self.do_get(Request::Ddl(DdlRequest { + expr: Some(DdlExpr::CreateTable(expr)), + })) + .await + .map(OutputWithMetrics::into_output) + } + + /// Alters an existing table using the provided alter expression. + #[cfg(feature = "testing")] + pub async fn alter(self, expr: AlterTableExpr) -> Result { + self.do_get(Request::Ddl(DdlRequest { + expr: Some(DdlExpr::AlterTable(expr)), + })) + .await + .map(OutputWithMetrics::into_output) + } + + async fn do_get(self, request: Request) -> Result { + let Self { database, options } = self; + database.do_get(request, options).await + } +} + /// by grpc standard, only `Unavailable` is retryable, see: https://github.com/grpc/grpc/blob/master/doc/statuscodes.md#status-codes-and-their-use-in-grpc pub fn is_grpc_retryable(err: &tonic::Status) -> bool { matches!(err.code(), tonic::Code::Unavailable) @@ -934,12 +1037,14 @@ mod tests { #[test] fn test_put_flow_extensions_preserves_comma_bearing_values() { let mut metadata = MetadataMap::new(); - Database::put_flow_extensions( + Database::put_metadata_value( &mut metadata, - &[ + FLOW_EXTENSIONS_METADATA_KEY, + Database::encode_flow_extensions(&[ ("flow.return_region_seq", "true"), ("flow.incremental_after_seqs", r#"{"1":10,"2":20}"#), - ], + ]) + .unwrap(), ) .unwrap(); @@ -969,7 +1074,12 @@ mod tests { (9_007_199_254_740_993_u64, 9_007_199_254_740_995_u64), ]); - Database::put_snapshot_seqs(&mut metadata, &snapshot_seqs).unwrap(); + Database::put_metadata_value( + &mut metadata, + SNAPSHOT_SEQS_METADATA_KEY, + Database::encode_snapshot_seqs(&snapshot_seqs).unwrap(), + ) + .unwrap(); let value = metadata .get(SNAPSHOT_SEQS_METADATA_KEY) @@ -980,6 +1090,50 @@ mod tests { assert_eq!(decoded, snapshot_seqs); } + #[test] + fn test_flight_request_builder_applies_request_options() { + let database = Database::new("greptime", "public", Client::default()); + let snapshot_seqs = HashMap::from([(42, 99)]); + let request = database + .flight_request() + .with_hints(&[("query_parallelism", "1")]) + .with_flow_extensions(&[("flow.return_region_seq", "true")]) + .with_snapshot_seqs(&snapshot_seqs) + .with_timeout(Duration::from_millis(50)); + let mut tonic_request = tonic::Request::new(()); + + request.options.apply_to(&mut tonic_request).unwrap(); + + let metadata = tonic_request.metadata(); + assert_eq!( + metadata.get(HINTS_METADATA_KEY).unwrap(), + "query_parallelism=1" + ); + assert_eq!( + serde_json::from_str::>( + metadata + .get(FLOW_EXTENSIONS_METADATA_KEY) + .unwrap() + .to_str() + .unwrap(), + ) + .unwrap(), + vec![("flow.return_region_seq".to_string(), "true".to_string())] + ); + assert_eq!( + serde_json::from_str::>( + metadata + .get(SNAPSHOT_SEQS_METADATA_KEY) + .unwrap() + .to_str() + .unwrap(), + ) + .unwrap(), + snapshot_seqs + ); + assert!(metadata.get("grpc-timeout").is_some()); + } + #[test] fn test_flight_ctx() { let mut ctx = FlightContext::default(); @@ -1020,6 +1174,28 @@ mod tests { assert_eq!(expected.should_retry(), actual.should_retry()); } + #[test] + fn test_flight_stream_error_preserves_addr_and_message() { + let error = flight_stream_error( + "127.0.0.1:4001", + Status::out_of_range("message length too large").into(), + ); + + assert!(matches!( + &error, + Error::FlightStream { + addr, + tonic_code: Code::OutOfRange, + message, + .. + } if addr == "127.0.0.1:4001" && message == "message length too large" + )); + assert_eq!( + "Failed to receive Flight data from 127.0.0.1:4001, code: Operation was attempted past the valid range: message length too large", + error.to_string(), + ); + } + #[test] fn test_from_tonic_status_with_retry_hint() { let mut headers = HeaderMap::new(); @@ -1117,13 +1293,13 @@ mod tests { #[tokio::test] async fn test_affected_rows_inline_metrics_are_parsed() { - let output = output_from_flight_message_stream(futures_util::stream::iter(vec![Ok( - FlightMessage::AffectedRows { + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![Ok(FlightMessage::AffectedRows { rows: 3, metrics: Some(terminal_metrics_json()), - }, - )] - as Vec>)) + })] as Vec>), + ) .await .unwrap(); @@ -1138,14 +1314,16 @@ mod tests { #[tokio::test] async fn test_affected_rows_inline_metrics_rejects_trailing_metrics() { let metrics_json = terminal_metrics_json(); - let err = output_from_flight_message_stream(futures_util::stream::iter(vec![ - Ok(FlightMessage::AffectedRows { - rows: 3, - metrics: Some(metrics_json.clone()), - }), - Ok(FlightMessage::Metrics(metrics_json)), - ] - as Vec>)) + let err = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: Some(metrics_json.clone()), + }), + Ok(FlightMessage::Metrics(metrics_json)), + ] as Vec>), + ) .await .unwrap_err(); @@ -1167,12 +1345,14 @@ mod tests { vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], ) .unwrap(); - let output = output_from_flight_message_stream(futures_util::stream::iter(vec![ - Ok(FlightMessage::Schema(schema.arrow_schema().clone())), - Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), - Ok(FlightMessage::Metrics("{not-json}".to_string())), - ] - as Vec>)) + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + Ok(FlightMessage::Metrics("{not-json}".to_string())), + ] as Vec>), + ) .await .unwrap(); let terminal_metrics = output.metrics.clone(); @@ -1211,18 +1391,20 @@ mod tests { vec![Arc::new(Int32Vector::from_slice([2])) as VectorRef], ) .unwrap(); - let output = output_from_flight_message_stream(futures_util::stream::iter(vec![ - Ok(FlightMessage::Schema(schema.arrow_schema().clone())), - Ok(FlightMessage::RecordBatch( - first_batch.into_df_record_batch(), - )), - Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(1))), - Ok(FlightMessage::RecordBatch( - second_batch.into_df_record_batch(), - )), - Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(2))), - ] - as Vec>)) + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch( + first_batch.into_df_record_batch(), + )), + Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(1))), + Ok(FlightMessage::RecordBatch( + second_batch.into_df_record_batch(), + )), + Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(2))), + ] as Vec>), + ) .await .unwrap(); let terminal_metrics = output.metrics.clone(); diff --git a/src/client/src/error.rs b/src/client/src/error.rs index 1db0858f087..168ba422715 100644 --- a/src/client/src/error.rs +++ b/src/client/src/error.rs @@ -40,6 +40,21 @@ pub enum Error { source: BoxedError, }, + #[snafu(display( + "Failed to receive Flight data from {}, code: {}: {}", + addr, + tonic_code, + message + ))] + FlightStream { + addr: String, + tonic_code: Code, + message: String, + source: BoxedError, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to convert FlightData"))] ConvertFlightData { #[snafu(implicit)] @@ -154,6 +169,7 @@ impl ErrorExt for Error { Error::Server { code, .. } | Error::Tonic { code, .. } => *code, Error::FlightGet { source, .. } + | Error::FlightStream { source, .. } | Error::RegionServer { source, .. } | Error::FlowServer { source, .. } => source.status_code(), Error::CreateChannel { source, .. } @@ -174,6 +190,7 @@ impl ErrorExt for Error { match self { Error::Tonic { retry_hint, .. } => *retry_hint, Error::FlightGet { source, .. } + | Error::FlightStream { source, .. } | Error::RegionServer { source, .. } | Error::FlowServer { source, .. } | Error::External { source, .. } => source.retry_hint(), @@ -193,6 +210,7 @@ impl Error { pub fn tonic_code(&self) -> Option { match self { Self::FlightGet { tonic_code, .. } + | Self::FlightStream { tonic_code, .. } | Self::RegionServer { code: tonic_code, .. } diff --git a/src/client/src/flight.rs b/src/client/src/flight.rs index 4d00cba962e..1753fdd5246 100644 --- a/src/client/src/flight.rs +++ b/src/client/src/flight.rs @@ -12,12 +12,82 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::pin::Pin; + use arrow_flight::FlightData; use common_grpc::flight::{FlightDecoder, FlightMessage}; -use snafu::ResultExt; +use futures_util::stream::Peekable; +use futures_util::{Stream, StreamExt}; +use snafu::{OptionExt, ResultExt}; use crate::Result; -use crate::error::{ConvertFlightDataSnafu, Error}; +use crate::error::{ConvertFlightDataSnafu, Error, IllegalFlightMessagesSnafu}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum FlightMessageKind { + Schema, + RecordBatch, + AffectedRows, + Metrics, +} + +impl From<&FlightMessage> for FlightMessageKind { + fn from(message: &FlightMessage) -> Self { + match message { + FlightMessage::Schema(_) => Self::Schema, + FlightMessage::RecordBatch(_) => Self::RecordBatch, + FlightMessage::AffectedRows { .. } => Self::AffectedRows, + FlightMessage::Metrics(_) => Self::Metrics, + } + } +} + +pub(crate) struct FlightMessageReader { + /// Remote Flight peer associated with this response stream. + remote_addr: String, + messages: Peekable, +} + +impl FlightMessageReader +where + S: Stream> + Unpin, +{ + pub(crate) fn new(remote_addr: impl Into, messages: S) -> Self { + Self { + remote_addr: remote_addr.into(), + messages: messages.peekable(), + } + } + + pub(crate) fn remote_addr(&self) -> &str { + &self.remote_addr + } + + pub(crate) async fn read_first(&mut self) -> Result { + self.read_next().await?.context(IllegalFlightMessagesSnafu { + reason: "Expect the response not to be empty", + }) + } + + pub(crate) async fn read_next(&mut self) -> Result> { + self.messages.next().await.transpose() + } + + pub(crate) async fn peek_next_message_kind(&mut self) -> Result> { + match Pin::new(&mut self.messages).peek().await { + Some(Ok(message)) => Ok(Some(message.into())), + None => Ok(None), + Some(Err(_)) => match self.read_next().await { + // `peek` only borrows the error; consume it to preserve the source error. + Err(error) => Err(error), + Ok(_) => IllegalFlightMessagesSnafu { + reason: "Flight stream changed after peek".to_string(), + } + .fail(), + }, + } + } +} pub(crate) fn decode_flight_data( decoder: &mut FlightDecoder, diff --git a/src/client/src/region.rs b/src/client/src/region.rs index 2bb5f7dd0a6..cd155810ea9 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -47,7 +47,7 @@ use crate::error::{ self, FlightGetSnafu, IllegalDatabaseResponseSnafu, IllegalFlightMessagesSnafu, MissingFieldSnafu, Result, ServerSnafu, }; -use crate::flight::decode_flight_data; +use crate::flight::{FlightMessageKind, FlightMessageReader, decode_flight_data}; use crate::{Client, metrics}; const FLIGHT_DO_GET_TIMEOUT: Duration = Duration::from_secs(10); @@ -203,26 +203,22 @@ impl RegionRequester { async fn recordbatches_from_flight_message_stream( addr: String, - mut flight_message_stream: S, + flight_message_stream: S, ) -> Result where S: Stream> + Send + Unpin + 'static, { - let Some(first_flight_message) = flight_message_stream.next().await else { - return IllegalFlightMessagesSnafu { - reason: "Expect the response not to be empty", - } - .fail() - .map_err(|error| flight_stream_error(&addr, error)); - }; - let FlightMessage::Schema(schema) = - first_flight_message.map_err(|e| flight_stream_error(&addr, e))? + let mut reader = FlightMessageReader::new(addr.clone(), flight_message_stream); + let FlightMessage::Schema(schema) = reader + .read_first() + .await + .map_err(|error| flight_stream_error(reader.remote_addr(), error))? else { return IllegalFlightMessagesSnafu { reason: "Expect schema to be the first flight message", } .fail() - .map_err(|error| flight_stream_error(&addr, error)); + .map_err(|error| flight_stream_error(reader.remote_addr(), error)); }; let metrics = Arc::new(ArcSwapOption::from(None)); @@ -233,31 +229,23 @@ where let schema = Arc::new(datatypes::schema::Schema::try_from(schema).context(error::ConvertSchemaSnafu)?); let schema_cloned = schema.clone(); - let stream_addr = addr.clone(); + let stream_addr = addr; let stream = Box::pin(stream!({ let _span = tracing_context.attach(common_telemetry::tracing::info_span!( "poll_flight_data_stream" )); - let mut buffered_message: Option = None; let mut stream_ended = false; while !stream_ended { - // get the next message from the buffered message or read from the flight message stream - let flight_message_item = if let Some(msg) = buffered_message.take() { - Some(Ok(msg)) - } else { - flight_message_stream.next().await - }; - - let flight_message = match flight_message_item { - Some(Ok(message)) => message, - Some(Err(e)) => { - yield Err(BoxedError::new(flight_stream_error(&stream_addr, e))) + let flight_message = match reader.read_next().await { + Ok(Some(message)) => message, + Ok(None) => break, + Err(error) => { + yield Err(BoxedError::new(flight_stream_error(&stream_addr, error))) .context(ExternalSnafu); break; } - None => break, }; match flight_message { @@ -265,42 +253,54 @@ where let result_to_yield = RecordBatch::from_df_record_batch(schema_cloned.clone(), record_batch); - // get the next message from the stream. normally it should be a metrics message. - if let Some(next_flight_message_result) = flight_message_stream.next().await { - match next_flight_message_result { - Ok(FlightMessage::Metrics(s)) => { - let m = serde_json::from_str(&s).ok().map(Arc::new); - metrics_ref.swap(m); - } - Ok(FlightMessage::RecordBatch(rb)) => { - // for some reason it's not a metrics message, so we need to buffer this record batch - // and yield it in the next iteration. - buffered_message = Some(FlightMessage::RecordBatch(rb)); - } - Ok(_) => { - yield IllegalFlightMessagesSnafu { - reason: "A RecordBatch message can only be succeeded by a Metrics message or another RecordBatch message" - } - .fail() - .map_err(BoxedError::new) - .context(ExternalSnafu); - break; - } - Err(e) => { - yield Err(BoxedError::new(flight_stream_error(&stream_addr, e))) + // Metrics follow a batch so MergeScan can observe them before yielding it. + match reader.peek_next_message_kind().await { + Ok(Some(FlightMessageKind::Metrics)) => { + let metrics_message = match reader.read_next().await { + Ok(Some(FlightMessage::Metrics(metrics))) => metrics, + Ok(Some(_) | None) => { + yield IllegalFlightMessagesSnafu { + reason: "Flight stream changed after peek", + } + .fail() + .map_err(BoxedError::new) .context(ExternalSnafu); - break; + break; + } + Err(error) => { + yield Err(BoxedError::new(flight_stream_error( + &stream_addr, + error, + ))) + .context(ExternalSnafu); + break; + } + }; + let metrics = serde_json::from_str(&metrics_message).ok().map(Arc::new); + metrics_ref.swap(metrics); + } + Ok(Some(FlightMessageKind::RecordBatch)) => {} + Ok(Some(FlightMessageKind::Schema | FlightMessageKind::AffectedRows)) => { + yield IllegalFlightMessagesSnafu { + reason: "A RecordBatch message can only be succeeded by a Metrics message or another RecordBatch message" } + .fail() + .map_err(BoxedError::new) + .context(ExternalSnafu); + break; + } + Ok(None) => stream_ended = true, + Err(error) => { + yield Err(BoxedError::new(flight_stream_error(&stream_addr, error))) + .context(ExternalSnafu); + break; } - } else { - // the stream has ended - stream_ended = true; } yield Ok(result_to_yield); } FlightMessage::Metrics(s) => { - // just a branch in case of some metrics message comes after other things. + // Metrics may arrive before the next RecordBatch. let m = serde_json::from_str(&s).ok().map(Arc::new); metrics_ref.swap(m); continue; @@ -412,6 +412,7 @@ mod test { RemoteDynFilterUnregister, RemoteDynFilterUpdate, region_request, remote_dyn_filter_request, }; use common_recordbatch::adapter::RecordBatchMetrics; + use datatypes::arrow::array::Int32Array; use datatypes::prelude::{ConcreteDataType, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::Int32Vector; @@ -598,6 +599,86 @@ mod test { assert_eq!(metrics.elapsed_compute, 7); } + #[tokio::test] + async fn test_record_batch_stream_updates_following_metrics_before_yielding_batch() { + let schema = test_schema(); + let batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], + ) + .unwrap(); + + let mut recordbatches = recordbatches_from_flight_message_stream( + "test-peer".to_string(), + stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + Ok(FlightMessage::Metrics(test_metrics_json())), + ]), + ) + .await + .unwrap(); + + let batch = recordbatches.next().await.unwrap().unwrap(); + assert_eq!(batch.num_rows(), 1); + + let metrics = recordbatches.metrics().unwrap(); + assert_eq!(metrics.elapsed_compute, 7); + assert!(recordbatches.next().await.is_none()); + } + + #[tokio::test] + async fn test_record_batch_stream_preserves_peeked_record_batch() { + let schema = test_schema(); + let first_batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], + ) + .unwrap(); + let second_batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([2])) as VectorRef], + ) + .unwrap(); + + let mut recordbatches = recordbatches_from_flight_message_stream( + "test-peer".to_string(), + stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch( + first_batch.into_df_record_batch(), + )), + Ok(FlightMessage::RecordBatch( + second_batch.into_df_record_batch(), + )), + ]), + ) + .await + .unwrap(); + + let first_batch = recordbatches.next().await.unwrap().unwrap(); + let second_batch = recordbatches.next().await.unwrap().unwrap(); + assert_eq!( + first_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 1 + ); + assert_eq!( + second_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 2 + ); + assert!(recordbatches.next().await.is_none()); + } + #[tokio::test] async fn test_record_batch_stream_exposes_error_after_pre_batch_metrics() { let schema = test_schema(); diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 0d9398b0889..79d0dd19152 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -67,7 +67,8 @@ use servers::error::{ }; use servers::grpc::FlightCompression; use servers::grpc::flight::{ - FlightCraft, FlightRecordBatchSource, FlightRecordBatchStream, TonicStream, + FlightCraft, FlightRecordBatchSource, FlightRecordBatchStream, FlightRecordBatchStreamInput, + TonicStream, }; use servers::grpc::region_server::RegionServerHandler; use session::context::{ @@ -990,7 +991,11 @@ impl FlightCraft for RegionServer { }; let stream = Box::pin(FlightRecordBatchStream::new( - FlightRecordBatchSource::initializer(initializer), + FlightRecordBatchStreamInput::initializer(async move { + initializer + .await + .map(FlightRecordBatchSource::RecordBatches) + }), tracing_context, self.flight_compression, query_ctx, diff --git a/src/flow/src/batching_mode.rs b/src/flow/src/batching_mode.rs index bb946e75740..3981446721b 100644 --- a/src/flow/src/batching_mode.rs +++ b/src/flow/src/batching_mode.rs @@ -44,6 +44,8 @@ pub struct BatchingModeOptions { /// The gRPC connection timeout #[serde(with = "humantime_serde")] pub grpc_conn_timeout: Duration, + #[serde(with = "humantime_serde")] + pub experimental_flight_do_get_timeout: Duration, /// The gRPC max retry number pub experimental_grpc_max_retries: u32, /// Flow wait for available frontend timeout, @@ -72,6 +74,7 @@ impl Default for BatchingModeOptions { slow_query_threshold: Duration::from_secs(60), experimental_min_refresh_duration: Duration::new(5, 0), grpc_conn_timeout: Duration::from_secs(5), + experimental_flight_do_get_timeout: Duration::from_secs(10), experimental_grpc_max_retries: 3, experimental_frontend_scan_timeout: Duration::from_secs(30), experimental_max_filter_num_per_query: 20, diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index aa94b91b2ae..27db96748a6 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -18,7 +18,6 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex, RwLock, Weak}; use api::v1::greptime_request::Request; -use api::v1::query_request::Query; use api::v1::{CreateTableExpr, QueryRequest}; use client::{Client, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, Database, OutputWithMetrics}; use common_error::ext::BoxedError; @@ -340,53 +339,6 @@ impl FrontendClient { }) } - /// Execute a SQL statement on the frontend. - pub async fn sql(&self, catalog: &str, schema: &str, sql: &str) -> Result { - match self { - FrontendClient::Distributed { .. } => { - let db = self.get_random_active_frontend(catalog, schema).await?; - db.database - .sql(sql) - .await - .map_err(BoxedError::new) - .context(ExternalSnafu) - } - FrontendClient::Standalone { - database_client, .. - } => { - let ctx = QueryContextBuilder::default() - .current_catalog(catalog.to_string()) - .current_schema(schema.to_string()) - .build(); - let ctx = Arc::new(ctx); - { - let database_client = { - database_client - .handler - .lock() - .unwrap() - .as_ref() - .context(UnexpectedSnafu { - reason: "Standalone's frontend instance is not set", - })? - .upgrade() - .context(UnexpectedSnafu { - reason: "Failed to upgrade database client", - })? - }; - let req = Request::Query(QueryRequest { - query: Some(Query::Sql(sql.to_string())), - }); - database_client - .do_query(req, ctx) - .await - .map_err(BoxedError::new) - .context(ExternalSnafu) - } - } - } - } - /// Execute a flow query and return terminal metrics. `snapshot_seqs` are /// optional read upper bounds used only by snapshot-fenced repair chunks. pub(crate) async fn query_with_terminal_metrics( @@ -413,12 +365,12 @@ impl FrontendClient { peer: db.peer.clone(), }); db.database - .query_with_terminal_metrics_and_flow_extensions( - request, - &hints, - extensions, - snapshot_seqs, - ) + .flight_request() + .with_hints(&hints) + .with_flow_extensions(extensions) + .with_snapshot_seqs(snapshot_seqs) + .with_timeout(batch_opts.experimental_flight_do_get_timeout) + .query_with_terminal_metrics(request) .await .map_err(BoxedError::new) .context(ExternalSnafu) @@ -623,6 +575,7 @@ mod tests { use std::task::{Context, Poll}; use std::time::Duration; + use api::v1::query_request::Query; use arrow_flight::flight_service_server::FlightServiceServer; use arrow_flight::{FlightData, Ticket}; use common_query::{Output, OutputData}; diff --git a/src/servers/src/grpc/flight.rs b/src/servers/src/grpc/flight.rs index e6f35e17b29..0b7a9dd5704 100644 --- a/src/servers/src/grpc/flight.rs +++ b/src/servers/src/grpc/flight.rs @@ -30,8 +30,7 @@ use bytes::{self, Bytes}; use common_error::ext::ErrorExt; use common_grpc::flight::do_put::{DoPutMetadata, DoPutResponse}; use common_grpc::flight::{ - FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightEncoder, FlightMessage, - SNAPSHOT_SEQS_METADATA_KEY, + FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightMessage, SNAPSHOT_SEQS_METADATA_KEY, }; use common_memory_manager::MemoryGuard; use common_query::{Output, OutputData}; @@ -46,14 +45,16 @@ use prost::Message; use query::metrics::terminal_recordbatch_metrics_from_plan_if_requested; use query::options::FlowQueryExtensions; use session::context::{Channel, QueryContextRef}; -use snafu::{IntoError, ResultExt, ensure}; +use snafu::{IntoError, OptionExt, ResultExt, ensure}; use table::table_name::TableName; use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; use tonic::{Request, Response, Status, Streaming}; -use crate::error::{InvalidParameterSnafu, Result, ToJsonSnafu}; -pub use crate::grpc::flight::stream::{FlightRecordBatchSource, FlightRecordBatchStream}; +use crate::error::{InvalidParameterSnafu, InvalidQuerySnafu, Result, ToJsonSnafu}; +pub use crate::grpc::flight::stream::{ + FlightRecordBatchSource, FlightRecordBatchStream, FlightRecordBatchStreamInput, +}; use crate::grpc::greptime_handler::{ GreptimeRequestHandler, create_query_context, get_request_type, }; @@ -221,9 +222,12 @@ impl FlightCraft for GreptimeRequestHandler { ); let flight_compression = self.flight_compression; async { - let output = self - .handle_request_with_query_ctx(request, query_ctx.clone()) + let query = request.request.context(InvalidQuerySnafu { + reason: "Expecting non-empty GreptimeRequest.", + })?; + self.authenticate_request_with_query_ctx(request.header.as_ref(), &query_ctx) .await?; + let output = self.handle_request_with_query_ctx(query, query_ctx.clone()); let stream = to_flight_data_stream( output, TracingContext::from_current_span(), @@ -573,32 +577,36 @@ fn extract_json_metadata( Ok(Some(parsed)) } -fn to_flight_data_stream( - output: Output, +fn to_flight_data_stream( + output: F, tracing_context: TracingContext, flight_compression: FlightCompression, query_ctx: QueryContextRef, should_emit_terminal_metrics: bool, -) -> TonicStream { +) -> TonicStream +where + F: std::future::Future> + Send + 'static, +{ + let initializer = async move { + let output = output.await.map_err(Status::from)?; + output_to_flight_record_batch_source(output, should_emit_terminal_metrics) + }; + let stream = FlightRecordBatchStream::new( + FlightRecordBatchStreamInput::initializer(initializer), + tracing_context, + flight_compression, + query_ctx, + ); + Box::pin(stream) as _ +} + +fn output_to_flight_record_batch_source( + output: Output, + should_emit_terminal_metrics: bool, +) -> TonicResult { match output.data { - OutputData::Stream(stream) => { - let stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(stream), - tracing_context, - flight_compression, - query_ctx, - ); - Box::pin(stream) as _ - } - OutputData::RecordBatches(x) => { - let stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(x.as_stream()), - tracing_context, - flight_compression, - query_ctx, - ); - Box::pin(stream) as _ - } + OutputData::Stream(stream) => Ok(FlightRecordBatchSource::RecordBatches(stream)), + OutputData::RecordBatches(x) => Ok(FlightRecordBatchSource::RecordBatches(x.as_stream())), OutputData::AffectedRows(rows) => { let terminal_metrics = match terminal_recordbatch_metrics_from_plan_if_requested( output.meta.plan, @@ -607,20 +615,17 @@ fn to_flight_data_stream( Some(metrics) => match serde_json::to_string(&metrics) { Ok(metrics) => Some(metrics), Err(e) => { - let stream = tokio_stream::once(Err(Status::internal(format!( + return Err(Status::internal(format!( "Failed to serialize terminal metrics: {e}" - )))); - return Box::pin(stream) as _; + ))); } }, None => None, }; - let affected_rows = FlightEncoder::default().encode(FlightMessage::AffectedRows { + Ok(FlightRecordBatchSource::AffectedRows { rows, metrics: terminal_metrics, - }); - let stream = tokio_stream::iter(affected_rows.into_iter().map(Ok)); - Box::pin(stream) as _ + }) } } } diff --git a/src/servers/src/grpc/flight/stream.rs b/src/servers/src/grpc/flight/stream.rs index 319a797d9ca..b740f12d02e 100644 --- a/src/servers/src/grpc/flight/stream.rs +++ b/src/servers/src/grpc/flight/stream.rs @@ -41,27 +41,34 @@ use crate::error; use crate::grpc::FlightCompression; use crate::grpc::flight::TonicResult; -pub type FlightRecordBatchStreamInitializer = Pin< - Box< - dyn Future> - + Send - + 'static, - >, ->; - pub enum FlightRecordBatchSource { RecordBatches(SendableRecordBatchStream), - Initializer(FlightRecordBatchStreamInitializer), + AffectedRows { + rows: usize, + metrics: Option, + }, } -impl FlightRecordBatchSource { - pub fn initializer(initializer: F) -> Self - where - F: Future> - + Send - + 'static, - { - Self::Initializer(Box::pin(initializer)) +/// Determines whether a Flight result is ready now or initialized asynchronously. +pub enum FlightRecordBatchStreamInput>> +{ + Ready(FlightRecordBatchSource), + Initializer(F), +} + +impl FlightRecordBatchStreamInput { + /// Creates an input from a source that is already available. + pub fn ready(source: FlightRecordBatchSource) -> Self { + Self::Ready(source) + } +} + +impl FlightRecordBatchStreamInput { + /// Creates an input that obtains its source asynchronously. + /// + /// Errors from the initializer are returned through the Flight response stream. + pub fn initializer(initializer: F) -> Self { + Self::Initializer(initializer) } } @@ -160,31 +167,32 @@ impl FlightRecordBatchStream { Self::send_metrics(tx, metrics, metrics_str).await } - pub fn new( - source: FlightRecordBatchSource, + pub fn new( + input: FlightRecordBatchStreamInput, tracing_context: TracingContext, compression: FlightCompression, query_ctx: QueryContextRef, - ) -> Self { - let should_send_partial_metrics = query_ctx.explain_verbose(); - let can_send_metrics_before_batch = query_ctx.explain_verbose() - && query_ctx.live_analyze_metrics_enabled() - && query_ctx - .remote_query_id() - .zip(query_ctx.extension(SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY)) - .is_some_and(|(remote_query_id, capability)| capability == remote_query_id); + ) -> Self + where + F: Future> + Send + 'static, + { let (mut tx, rx) = mpsc::channel::>(1); - let source_type = match &source { - FlightRecordBatchSource::RecordBatches(_) => "record_batches", - FlightRecordBatchSource::Initializer(_) => "initializer", + let source_type = match &input { + FlightRecordBatchStreamInput::Ready(FlightRecordBatchSource::RecordBatches(_)) => { + "record_batches" + } + FlightRecordBatchStreamInput::Ready(FlightRecordBatchSource::AffectedRows { + .. + }) => "affected_rows", + FlightRecordBatchStreamInput::Initializer(_) => "initializer", }; let initializer_tracing_context = tracing_context.clone(); let join_handle = common_runtime::spawn_global( async move { - let recordbatches = async move { - match source { - FlightRecordBatchSource::RecordBatches(recordbatches) => Ok(recordbatches), - FlightRecordBatchSource::Initializer(initializer) => initializer.await, + let source = async move { + match input { + FlightRecordBatchStreamInput::Ready(source) => Ok(source), + FlightRecordBatchStreamInput::Initializer(initializer) => initializer.await, } } .trace( @@ -193,8 +201,20 @@ impl FlightRecordBatchStream { ) .await; - match recordbatches { - Ok(recordbatches) => { + match source { + Ok(FlightRecordBatchSource::RecordBatches(recordbatches)) => { + let should_send_partial_metrics = query_ctx.explain_verbose(); + let can_send_metrics_before_batch = + query_ctx.explain_verbose() + && query_ctx.live_analyze_metrics_enabled() + && query_ctx + .remote_query_id() + .zip(query_ctx.extension( + SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY, + )) + .is_some_and(|(remote_query_id, capability)| { + capability == remote_query_id + }); Self::flight_data_stream( recordbatches, tx, @@ -203,6 +223,11 @@ impl FlightRecordBatchStream { ) .await; } + Ok(FlightRecordBatchSource::AffectedRows { rows, metrics }) => { + let _ = tx + .send(Ok(FlightMessage::AffectedRows { rows, metrics })) + .await; + } Err(status) => { let _ = tx.send(Err(status)).await; } @@ -485,7 +510,9 @@ mod test { .unwrap() .as_stream(); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), QueryContext::arc(), @@ -519,10 +546,34 @@ mod test { } } + #[tokio::test] + async fn test_flight_record_batch_stream_encodes_affected_rows() { + let mut stream = FlightRecordBatchStream::new( + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::AffectedRows { + rows: 42, + metrics: Some(r#"{"region_watermarks":[]}"#.to_string()), + }), + TracingContext::default(), + FlightCompression::default(), + QueryContext::arc(), + ); + + let data = stream.next().await.unwrap().unwrap(); + let message = FlightDecoder::default().try_decode(&data).unwrap().unwrap(); + assert!(matches!( + message, + FlightMessage::AffectedRows { + rows: 42, + metrics: Some(_), + } + )); + assert!(stream.next().await.is_none()); + } + #[tokio::test] async fn test_flight_record_batch_stream_forwards_initializer_error() { let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::initializer(async { + FlightRecordBatchStreamInput::initializer(async { Err(tonic::Status::unavailable( "remote read initialization failed", )) @@ -552,9 +603,12 @@ mod test { metrics, }); let query_ctx = query_context_with_live_metrics_and_matching_capability(); - query_ctx.set_explain_verbose(true); + let initializer_query_ctx = query_ctx.clone(); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::initializer(async move { + initializer_query_ctx.set_explain_verbose(true); + Ok(FlightRecordBatchSource::RecordBatches(recordbatches)) + }), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -609,7 +663,9 @@ mod test { let query_ctx = query_context_with_live_metrics_and_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -663,7 +719,9 @@ mod test { let query_ctx = query_context_with_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -706,7 +764,9 @@ mod test { let query_ctx = Arc::new(query_ctx); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -745,7 +805,9 @@ mod test { }); let query_ctx = query_context_with_matching_capability(); let mut stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::RecordBatches(recordbatches), + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), TracingContext::default(), FlightCompression::default(), query_ctx, diff --git a/src/servers/src/grpc/greptime_handler.rs b/src/servers/src/grpc/greptime_handler.rs index ad7779435b8..98786d30adc 100644 --- a/src/servers/src/grpc/greptime_handler.rs +++ b/src/servers/src/grpc/greptime_handler.rs @@ -15,11 +15,13 @@ //! Handler for Greptime Database service. It's implemented by frontend. use std::collections::HashMap; +use std::future::Future; use std::str::FromStr; use std::sync::{Arc, RwLock}; use std::time::Instant; use api::helper::request_type; +use api::v1::greptime_request::Request as QueryRequest; use api::v1::{GreptimeRequest, RequestHeader}; use auth::UserProviderRef; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; @@ -79,30 +81,39 @@ impl GreptimeRequestHandler { ) -> Result { let header = request.header.as_ref(); let query_ctx = create_query_context(Channel::Grpc, header, hints, HashMap::new())?; - self.handle_request_with_query_ctx(request, query_ctx).await - } - - pub(crate) async fn handle_request_with_query_ctx( - &self, - request: GreptimeRequest, - query_ctx: QueryContextRef, - ) -> Result { let query = request.request.context(InvalidQuerySnafu { reason: "Expecting non-empty GreptimeRequest.", })?; + self.authenticate_request_with_query_ctx(request.header.as_ref(), &query_ctx) + .await?; + self.handle_request_with_query_ctx(query, query_ctx).await + } - let header = request.header.as_ref(); - let user_info = context_auth::auth(self.user_provider.clone(), header, &query_ctx).await?; + pub(crate) async fn authenticate_request_with_query_ctx( + &self, + header: Option<&RequestHeader>, + query_ctx: &QueryContextRef, + ) -> Result<()> { + let user_info = context_auth::auth(self.user_provider.clone(), header, query_ctx).await?; query_ctx.set_current_user(user_info); + Ok(()) + } + pub(crate) fn handle_request_with_query_ctx( + &self, + query: QueryRequest, + query_ctx: QueryContextRef, + ) -> impl Future> + Send + 'static { let handler = self.handler.clone(); + let runtime = self.runtime.clone(); let request_type = request_type(&query).to_string(); let db = query_ctx.get_db_string(); let timer = RequestTimer::new(db.clone(), request_type); let tracing_context = TracingContext::from_current_span(); - let result_future = async move { - handler + async move { + let result_future = async move { + handler .do_query(query, query_ctx) .trace(tracing_context.attach(tracing::info_span!( "GreptimeRequestHandler::handle_request_runtime" @@ -118,26 +129,27 @@ impl GreptimeRequestHandler { } e }) - }; + }; - match &self.runtime { - Some(runtime) => { - // Executes requests in another runtime to - // 1. prevent the execution from being cancelled unexpected by Tonic runtime; - // - Refer to our blog for the rational behind it: - // https://www.greptime.com/blogs/2023-01-12-hidden-control-flow.html - // - Obtaining a `JoinHandle` to get the panic message (if there's any). - // From its docs, `JoinHandle` is cancel safe. The task keeps running even it's handle been dropped. - // 2. avoid the handler blocks the gRPC runtime incidentally. - runtime - .spawn(result_future) - .await - .context(JoinTaskSnafu) - .inspect_err(|e| { - timer.record(e.status_code()); - })? + match runtime { + Some(runtime) => { + // Executes requests in another runtime to + // 1. prevent the execution from being cancelled unexpected by Tonic runtime; + // - Refer to our blog for the rational behind it: + // https://www.greptime.com/blogs/2023-01-12-hidden-control-flow.html + // - Obtaining a `JoinHandle` to get the panic message (if there's any). + // From its docs, `JoinHandle` is cancel safe. The task keeps running even it's handle been dropped. + // 2. avoid the handler blocks the gRPC runtime incidentally. + runtime + .spawn(result_future) + .await + .context(JoinTaskSnafu) + .inspect_err(|e| { + timer.record(e.status_code()); + })? + } + None => result_future.await, } - None => result_future.await, } } diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index cf75ae7fff1..1a2d841600f 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -16,40 +16,45 @@ mod test { use std::collections::HashMap; use std::net::SocketAddr; + use std::pin::Pin; use std::sync::Arc; use std::time::{Duration, Instant}; use api::v1::auth_header::AuthScheme; + use api::v1::greptime_request::Request as GreptimeQueryRequest; use api::v1::query_request::Query; use api::v1::{Basic, ColumnDataType, ColumnDef, CreateTableExpr, QueryRequest, SemanticType}; use arrow_flight::flight_service_server::FlightServiceServer; use arrow_flight::{FlightData, FlightDescriptor, Ticket}; use auth::user_provider_from_option; + use client::region::RegionRequester; use client::{Client, Database}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; - use common_grpc::flight::do_put::DoPutMetadata; + use common_grpc::flight::do_put::{DoPutMetadata, DoPutResponse}; use common_grpc::flight::{FlightEncoder, FlightMessage}; - use common_query::OutputData; + use common_query::{Output, OutputData}; use common_recordbatch::adapter::RegionWatermarkEntry; use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; use common_telemetry::tracing_context::TracingContext; use datatypes::prelude::{ConcreteDataType, ScalarVector, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{Int32Vector, StringVector, TimestampMillisecondVector}; - use futures_util::StreamExt; + use futures_util::{Stream, StreamExt}; use hyper_util::rt::TokioIo; use itertools::Itertools; use servers::grpc::builder::GrpcServerBuilder; use servers::grpc::flight::{ FlightCraft, FlightCraftWrapper, FlightRecordBatchSource, FlightRecordBatchStream, - TonicStream, + FlightRecordBatchStreamInput, PutRecordBatchRequestStream, TonicStream, }; use servers::grpc::greptime_handler::GreptimeRequestHandler; use servers::grpc::{FlightCompression, GrpcServerConfig}; + use servers::query_handler::grpc::GrpcQueryHandler; use servers::server::Server; - use tonic::Response; + use session::context::QueryContextRef; use tonic::transport::Server as TonicServer; + use tonic::{Response, Status}; use tower::service_fn; use crate::cluster::GreptimeDbClusterBuilder; @@ -59,6 +64,12 @@ mod test { struct SlowFlightCraft; + struct ErrorFlightCraft; + + struct SlowRemoteQueryHandler { + region_requester: RegionRequester, + } + fn slow_recordbatch_stream() -> SendableRecordBatchStream { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( "value", @@ -83,9 +94,11 @@ mod test { _: tonic::Request, ) -> std::result::Result>, tonic::Status> { let stream = FlightRecordBatchStream::new( - FlightRecordBatchSource::initializer(async { + FlightRecordBatchStreamInput::initializer(async { tokio::time::sleep(Duration::from_secs(2)).await; - Ok(slow_recordbatch_stream()) + Ok(FlightRecordBatchSource::RecordBatches( + slow_recordbatch_stream(), + )) }), TracingContext::default(), FlightCompression::default(), @@ -96,14 +109,72 @@ mod test { } } - #[tokio::test(flavor = "multi_thread")] - async fn test_do_get_timeout_does_not_cancel_slow_flight_stream() { + #[async_trait::async_trait] + impl FlightCraft for ErrorFlightCraft { + async fn do_get( + &self, + _: tonic::Request, + ) -> std::result::Result>, tonic::Status> { + let stream = FlightRecordBatchStream::new( + FlightRecordBatchStreamInput::initializer(async { + Err(Status::internal("deferred initializer detail")) + }), + TracingContext::default(), + FlightCompression::default(), + session::context::QueryContext::arc(), + ); + Ok(Response::new(Box::pin(stream))) + } + } + + #[async_trait::async_trait] + impl GrpcQueryHandler for SlowRemoteQueryHandler { + async fn do_query( + &self, + _: GreptimeQueryRequest, + _: QueryContextRef, + ) -> servers::error::Result { + let stream = self + .region_requester + .do_get_inner(Ticket::default()) + .await + .unwrap(); + Ok(Output::new_with_stream(stream)) + } + + fn handle_put_record_batch_stream( + &self, + _: PutRecordBatchRequestStream, + _: QueryContextRef, + ) -> Pin> + Send>> { + Box::pin(futures::stream::empty()) + } + } + + fn client_for_flight_craft(addr: &'static str, craft: T) -> Client + where + T: FlightCraft, + { + client_for_flight_craft_with_max_encoding(addr, craft, None) + } + + fn client_for_flight_craft_with_max_encoding( + addr: &'static str, + craft: T, + max_encoding_message_size: Option, + ) -> Client + where + T: FlightCraft, + { let (client_io, server_io) = tokio::io::duplex(1024); tokio::spawn(async move { + let flight_service = FlightServiceServer::new(FlightCraftWrapper(craft)); + let flight_service = match max_encoding_message_size { + Some(size) => flight_service.max_encoding_message_size(size), + None => flight_service, + }; TonicServer::builder() - .add_service(FlightServiceServer::new(FlightCraftWrapper( - SlowFlightCraft, - ))) + .add_service(flight_service) .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>( server_io, )])) @@ -115,7 +186,7 @@ mod test { let mut client_io = Some(client_io); channel_manager .reset_with_connector( - "slow-flight", + addr, service_fn(move |_| { let client_io = client_io.take(); @@ -127,7 +198,12 @@ mod test { }), ) .unwrap(); - let client = Client::with_manager_and_urls(channel_manager, ["slow-flight"]); + Client::with_manager_and_urls(channel_manager, [addr]) + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_do_get_timeout_does_not_cancel_slow_flight_stream() { + let client = client_for_flight_craft("slow-flight", SlowFlightCraft); let mut flight_client = client.make_flight_client(false, false).unwrap(); let start = Instant::now(); @@ -143,6 +219,64 @@ mod test { assert!(stream.message().await.unwrap().is_some()); } + + #[tokio::test(flavor = "multi_thread")] + async fn test_deferred_flight_initializer_error_preserves_message() { + let client = client_for_flight_craft("error-flight", ErrorFlightCraft); + let mut flight_client = client.make_flight_client(false, false).unwrap(); + let mut stream = flight_client + .mut_inner() + .do_get(tonic::Request::new(Ticket::default())) + .await + .unwrap() + .into_inner(); + let error = stream.message().await.unwrap_err(); + assert_eq!(tonic::Code::Internal, error.code()); + assert_eq!("deferred initializer detail", error.message()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_deferred_flight_encode_error_preserves_message() { + let client = + client_for_flight_craft_with_max_encoding("limited-flight", SlowFlightCraft, Some(1)); + let mut flight_client = client.make_flight_client(false, false).unwrap(); + let mut stream = flight_client + .mut_inner() + .do_get(tonic::Request::new(Ticket::default())) + .await + .unwrap() + .into_inner(); + let error = stream.message().await.unwrap_err(); + assert_eq!(tonic::Code::OutOfRange, error.code()); + assert!(error.message().contains("message length too large")); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_flight_request_timeout_does_not_cancel_slow_datanode_stream() { + let datanode_client = client_for_flight_craft("slow-datanode", SlowFlightCraft); + let frontend_handler = GreptimeRequestHandler::new( + Arc::new(SlowRemoteQueryHandler { + region_requester: RegionRequester::new(datanode_client, false, false), + }), + None, + None, + FlightCompression::default(), + ); + let frontend_client = client_for_flight_craft("slow-frontend", frontend_handler); + let database = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, frontend_client); + + let start = Instant::now(); + let output = database + .flight_request() + .with_timeout(Duration::from_secs(1)) + .sql("select 1") + .await + .unwrap(); + // `sql()` waits for the first Flight message. Succeeding after this delay + // proves the request timeout applied only to the response header. + assert!(start.elapsed() >= Duration::from_secs(1)); + assert!(matches!(output.data, OutputData::Stream(_))); + } #[tokio::test(flavor = "multi_thread")] async fn test_standalone_flight_do_put() { common_telemetry::init_default_ut_logging(); @@ -446,16 +580,14 @@ mod test { ); let result = client - .query_with_terminal_metrics_and_flow_extensions( - QueryRequest { - query: Some(Query::Sql( - "select ts, a, `B` from foo order by ts".to_string(), - )), - }, - &[], - &[("flow.return_region_seq", "true")], - &HashMap::from([(region_id.as_u64(), stale_snapshot_seq)]), - ) + .flight_request() + .with_flow_extensions(&[("flow.return_region_seq", "true")]) + .with_snapshot_seqs(&HashMap::from([(region_id.as_u64(), stale_snapshot_seq)])) + .query_with_terminal_metrics(QueryRequest { + query: Some(Query::Sql( + "select ts, a, `B` from foo order by ts".to_string(), + )), + }) .await .unwrap(); diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 4a8f96260e8..43ba8701d76 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2252,6 +2252,7 @@ query_timeout = "10m" slow_query_threshold = "1m" experimental_min_refresh_duration = "5s" grpc_conn_timeout = "5s" +experimental_flight_do_get_timeout = "10s" experimental_grpc_max_retries = 3 experimental_frontend_scan_timeout = "30s" experimental_max_filter_num_per_query = 20 From b31f05eb5964c17342e7606d80dd04e6b93a3e03 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Thu, 27 Aug 2026 03:37:59 +0000 Subject: [PATCH 89/92] fix: update tokio-postgres and correct explain/fetch cursor output schema (#8955) * chore(deps): update tokio-postgres * fix: describing fetch cursor and analyze --- Cargo.lock | 180 ++++++++++++++++++++------- Cargo.toml | 2 +- src/common/meta/Cargo.toml | 2 +- src/common/recordbatch/src/cursor.rs | 9 ++ src/query/src/analyze.rs | 21 +++- src/query/src/lib.rs | 2 +- src/servers/Cargo.toml | 4 +- src/servers/src/postgres/handler.rs | 29 +++++ src/session/src/lib.rs | 9 ++ 9 files changed, 202 insertions(+), 56 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ec477e5c9a3..d77ec340a8d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4628,10 +4628,23 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb" dependencies = [ "const-oid 0.9.6", + "der_derive", + "flagset", "pem-rfc7468", "zeroize", ] +[[package]] +name = "der_derive" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8034092389675178f570469e6c3b0465d3d30b4505c294a6550db47f3c17ad18" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "deranged" version = "0.5.8" @@ -5334,6 +5347,12 @@ version = "0.5.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" +[[package]] +name = "flagset" +version = "0.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7ac824320a75a52197e8f2d787f6a38b6718bb6897a35142d749af3c0e8f4fe" + [[package]] name = "flatbuffers" version = "25.2.10" @@ -5973,7 +5992,7 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "wasm-bindgen", ] @@ -7866,13 +7885,14 @@ checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de" [[package]] name = "libredox" -version = "0.1.4" +version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1580801010e535496706ba011c15f8532df6b42297d2e471fec38ceadd8c0638" +checksum = "28d0a00925a9f930d679b6789b721e3a7f9ed110f41b86d2497caa780c3a070a" dependencies = [ "bitflags 2.12.1", "libc", - "redox_syscall 0.5.13", + "plain", + "redox_syscall 0.9.3", ] [[package]] @@ -8543,7 +8563,7 @@ checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" dependencies = [ "libc", "log", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "windows-sys 0.61.2", ] @@ -9358,6 +9378,15 @@ dependencies = [ "objc2-encode", ] +[[package]] +name = "objc2-core-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags 2.12.1", +] + [[package]] name = "objc2-encode" version = "4.1.0" @@ -9374,6 +9403,15 @@ dependencies = [ "objc2", ] +[[package]] +name = "objc2-system-configuration" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7216bd11cbda54ccabcab84d523dc93b858ec75ecfb3a7d89513fa22464da396" +dependencies = [ + "objc2-core-foundation", +] + [[package]] name = "object" version = "0.36.7" @@ -10552,16 +10590,7 @@ dependencies = [ "tokio", "tokio-rustls", "tokio-util", - "x509-certificate 0.25.0", -] - -[[package]] -name = "phf" -version = "0.11.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078" -dependencies = [ - "phf_shared 0.11.3", + "x509-certificate", ] [[package]] @@ -10814,6 +10843,12 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" +[[package]] +name = "plain" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4596b6d070b27117e987119b4dac604f3c58cfb0b191112e24771b2faeac1a6" + [[package]] name = "plist" version = "1.7.2" @@ -10924,16 +10959,16 @@ dependencies = [ [[package]] name = "postgres-types" -version = "0.2.9" +version = "0.2.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "613283563cd90e1dfc3518d548caee47e0e725455ed619881f5cf21f36de4b48" +checksum = "851ca9db4932932d69f3ea811b1abe63087a0f740a47692619dd40d4899b68be" dependencies = [ "array-init", "bytes", "chrono", "fallible-iterator", "postgres-protocol", - "serde", + "serde_core", "serde_json", ] @@ -11299,7 +11334,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4" dependencies = [ "bytes", - "heck 0.4.1", + "heck 0.5.0", "itertools 0.12.1", "log", "multimap", @@ -11319,7 +11354,7 @@ version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "itertools 0.14.0", "log", "multimap", @@ -12087,6 +12122,15 @@ dependencies = [ "bitflags 2.12.1", ] +[[package]] +name = "redox_syscall" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d678d17679829e73d371e96880897e98fee2ded7acc0a50bdf8af2affa4b2fe5" +dependencies = [ + "bitflags 2.12.1", +] + [[package]] name = "redox_users" version = "0.4.6" @@ -13785,7 +13829,7 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1961e2ef424c1424204d3a5d6975f934f56b6d50ff5732382d84ebf460e147f7" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "proc-macro2", "quote", "syn 2.0.117", @@ -14117,7 +14161,7 @@ dependencies = [ "stringprep", "thiserror 2.0.17", "tracing", - "whoami", + "whoami 1.6.0", ] [[package]] @@ -14156,7 +14200,7 @@ dependencies = [ "stringprep", "thiserror 2.0.17", "tracing", - "whoami", + "whoami 1.6.0", ] [[package]] @@ -15175,6 +15219,27 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" +[[package]] +name = "tls_codec" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0de2e01245e2bb89d6f05801c564fa27624dbd7b1846859876c7dad82e90bf6b" +dependencies = [ + "tls_codec_derive", + "zeroize", +] + +[[package]] +name = "tls_codec_derive" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d2e76690929402faae40aebdda620a2c0e25dd6d3b9afe48867dfd95991f4bd" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "tokio" version = "1.52.3" @@ -15241,9 +15306,9 @@ dependencies = [ [[package]] name = "tokio-postgres" -version = "0.7.13" +version = "0.7.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c95d533c83082bb6490e0189acaa0bbeef9084e60471b696ca6988cd0541fb0" +checksum = "a528f7d280f6d5b9cd149635c8705b0dd049754bc67d81d31fa25169a93809d3" dependencies = [ "async-trait", "byteorder", @@ -15254,29 +15319,29 @@ dependencies = [ "log", "parking_lot 0.12.4", "percent-encoding", - "phf 0.11.3", + "phf 0.13.1", "pin-project-lite", "postgres-protocol", "postgres-types", - "rand 0.9.4", - "socket2 0.5.10", + "rand 0.10.1", + "socket2 0.6.4", "tokio", "tokio-util", - "whoami", + "whoami 2.1.3", ] [[package]] name = "tokio-postgres-rustls" -version = "0.12.0" +version = "0.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04fb792ccd6bbcd4bba408eb8a292f70fc4a3589e5d793626f45190e6454b6ab" +checksum = "4c2ad44aa0ae96db89c4742212ed41645b2f597311ff6e1945542a4d9fadc2fb" dependencies = [ - "ring", "rustls", + "sha2 0.11.0", "tokio", "tokio-postgres", "tokio-rustls", - "x509-certificate 0.23.1", + "x509-cert", ] [[package]] @@ -16284,6 +16349,15 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wasi" +version = "0.14.7+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "883478de20367e224c0090af9cf5f9fa85bed63a95c1abf3afc5c083ebc06e8c" +dependencies = [ + "wasip2", +] + [[package]] name = "wasip2" version = "1.0.2+wasi-0.2.9" @@ -16308,6 +16382,15 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8dad83b4f25e74f184f64c43b150b91efe7647395b42289f38e50566d82855b" +[[package]] +name = "wasite" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66fe902b4a6b8028a753d5424909b764ccf79b7a209eac9bf97e59cda9f71a42" +dependencies = [ + "wasi 0.14.7+wasi-0.2.4", +] + [[package]] name = "wasm-bindgen" version = "0.2.122" @@ -16530,7 +16613,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6994d13118ab492c3c80c1f81928718159254c53c472bf9ce36f8dae4add02a7" dependencies = [ "redox_syscall 0.5.13", - "wasite", + "wasite 0.1.0", +] + +[[package]] +name = "whoami" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "626c4bac6755d76ffc12cb01b2eac751db1996b9e0041de9aa02c8c211ddc82c" +dependencies = [ + "libc", + "libredox", + "objc2-system-configuration", + "wasite 1.0.2", "web-sys", ] @@ -17139,22 +17234,15 @@ dependencies = [ ] [[package]] -name = "x509-certificate" -version = "0.23.1" +name = "x509-cert" +version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "66534846dec7a11d7c50a74b7cdb208b9a581cad890b7866430d438455847c85" +checksum = "1301e935010a701ae5f8655edc0ad17c44bad3ac5ce8c39185f75453b720ae94" dependencies = [ - "bcder", - "bytes", - "chrono", + "const-oid 0.9.6", "der", - "hex", - "pem", - "ring", - "signature", "spki", - "thiserror 1.0.69", - "zeroize", + "tls_codec", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 7c3fccfbbf8..fab7b2aa736 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -255,7 +255,7 @@ strum = { version = "0.27", features = ["derive"] } sysinfo = "0.33" tempfile = "3" tokio = { version = "1.47", features = ["full"] } -tokio-postgres = "0.7" +tokio-postgres = "0.7.18" tokio-rustls = { version = "0.26.2", default-features = false } tokio-stream = "0.1" tokio-util = { version = "0.7", features = ["io-util", "compat"] } diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index 8e0cbbdf354..8deb3d00814 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -89,7 +89,7 @@ strum.workspace = true table = { workspace = true, features = ["testing"] } tokio.workspace = true tokio-postgres = { workspace = true, optional = true } -tokio-postgres-rustls = { version = "0.12", optional = true } +tokio-postgres-rustls = { version = "0.14", optional = true } tonic.workspace = true tracing.workspace = true typetag.workspace = true diff --git a/src/common/recordbatch/src/cursor.rs b/src/common/recordbatch/src/cursor.rs index a741953ccc2..60ae9d54fa2 100644 --- a/src/common/recordbatch/src/cursor.rs +++ b/src/common/recordbatch/src/cursor.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use datatypes::schema::SchemaRef; use futures::StreamExt; use tokio::sync::Mutex; @@ -28,12 +29,15 @@ struct Inner { /// A cursor on RecordBatchStream that fetches data batch by batch pub struct RecordBatchStreamCursor { + schema: SchemaRef, inner: Mutex, } impl RecordBatchStreamCursor { pub fn new(stream: SendableRecordBatchStream) -> RecordBatchStreamCursor { + let schema = stream.schema(); Self { + schema, inner: Mutex::new(Inner { stream, current_row_index: 0, @@ -43,6 +47,11 @@ impl RecordBatchStreamCursor { } } + /// Returns the schema of the underlying record batch stream. + pub fn schema(&self) -> SchemaRef { + self.schema.clone() + } + /// Take `size` of row from the `RecordBatchStream` and create a new /// `RecordBatch` for these rows. pub async fn take(&self, size: usize) -> Result { diff --git a/src/query/src/analyze.rs b/src/query/src/analyze.rs index ec0abf25ba6..b1522fc05d5 100644 --- a/src/query/src/analyze.rs +++ b/src/query/src/analyze.rs @@ -46,6 +46,21 @@ const STAGE: &str = "stage"; const NODE: &str = "node"; const PLAN: &str = "plan"; +/// Returns the fixed output schema of [`DistAnalyzeExec`]: +/// (`stage`: UInt32, `node`: UInt32, `plan`: Utf8). +/// +/// Exposed so protocol-level `Describe` handlers can advertise the schema +/// clients will actually receive (the physical plan is rewritten in +/// `optimize_physical_plan` and its schema differs from the DataFusion +/// `Analyze` logical plan schema). +pub fn dist_analyze_output_schema() -> SchemaRef { + SchemaRef::new(Schema::new(vec![ + Field::new(STAGE, DataType::UInt32, true), + Field::new(NODE, DataType::UInt32, true), + Field::new(PLAN, DataType::Utf8, true), + ])) +} + #[derive(Debug)] pub struct DistAnalyzeExec { input: Arc, @@ -58,11 +73,7 @@ pub struct DistAnalyzeExec { impl DistAnalyzeExec { /// Create a new DistAnalyzeExec pub fn new(input: Arc, verbose: bool, format: AnalyzeFormat) -> Self { - let schema = SchemaRef::new(Schema::new(vec![ - Field::new(STAGE, DataType::UInt32, true), - Field::new(NODE, DataType::UInt32, true), - Field::new(PLAN, DataType::Utf8, true), - ])); + let schema = dist_analyze_output_schema(); let properties = Arc::new(Self::compute_properties(&input, schema.clone())); Self { input, diff --git a/src/query/src/lib.rs b/src/query/src/lib.rs index 68d8ff3a9eb..4b2a2d5264d 100644 --- a/src/query/src/lib.rs +++ b/src/query/src/lib.rs @@ -45,7 +45,7 @@ pub(crate) mod test_util; #[cfg(test)] mod tests; -pub use crate::analyze::analyze_plan_metrics_to_json_value; +pub use crate::analyze::{analyze_plan_metrics_to_json_value, dist_analyze_output_schema}; pub use crate::datafusion::DfContextProviderAdapter; pub use crate::query_engine::{ QueryEngine, QueryEngineContext, QueryEngineFactory, QueryEngineRef, diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 7936c96fbda..bb358ea76a6 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -167,8 +167,8 @@ servers = { workspace = true, features = ["testing"] } session = { workspace = true, features = ["testing"] } table.workspace = true tempfile = "3.0.0" -tokio-postgres = "0.7" -tokio-postgres-rustls = "0.12" +tokio-postgres.workspace = true +tokio-postgres-rustls = "0.14" [target.'cfg(unix)'.dev-dependencies] pprof = { version = "0.14", features = ["criterion", "flamegraph"] } diff --git a/src/servers/src/postgres/handler.rs b/src/servers/src/postgres/handler.rs index 484bb6a1f11..433eb11cdad 100644 --- a/src/servers/src/postgres/handler.rs +++ b/src/servers/src/postgres/handler.rs @@ -40,6 +40,7 @@ use pgwire::error::{ErrorInfo, PgWireError, PgWireResult}; use pgwire::messages::PgWireBackendMessage; use pgwire::messages::copy::CopyData; use pgwire::messages::data::DataRow; +use query::dist_analyze_output_schema; use query::planner::DfLogicalPlanner; use query::query_engine::DescribeResult; use session::Session; @@ -554,6 +555,16 @@ fn describe_fields( session: &Arc, ) -> PgWireResult> { match sql_plan { + // EXPLAIN ANALYZE: at execution time the physical plan is replaced with + // GreptimeDB's `DistAnalyzeExec` (see `optimize_physical_plan`), whose + // output schema (stage/node/plan) differs from the DataFusion `Analyze` + // logical plan schema (plan_type/plan). Describe with the schema the + // client will actually receive so the DataRow field count matches. + SqlPlan::Plan(LogicalPlan::Analyze(_), _) => { + let schema: Schema = + Schema::try_from(dist_analyze_output_schema()).map_err(convert_err)?; + schema_to_pg(&schema, format, None).map_err(convert_err) + } // query SqlPlan::Plan(plan, _) if !matches!(plan, LogicalPlan::Dml(_) | LogicalPlan::Ddl(_)) => { let schema: Schema = plan.schema().clone().try_into().map_err(convert_err)?; @@ -679,6 +690,24 @@ fn describe_fields( Ok(vec![]) } } + // FETCH cursor: return the cursor's schema so the RowDescription + // matches the DataRow field count sent during Execute. + SqlPlan::Statement(Statement::FetchCursor(fetch), _) => { + let cursor_name = fetch.cursor_name.to_string(); + match session.get_cursor(&cursor_name) { + Some(cursor) => { + // `cursor.schema()` is the GreptimeDB `SchemaRef` captured + // when the cursor was declared; `schema_to_pg` accepts it + // directly (no DataFusion -> GreptimeDB conversion needed). + schema_to_pg(&cursor.schema(), format, None).map_err(convert_err) + } + None => { + // Cursor not found (e.g. DECLARE hasn't executed yet in + // an extended-protocol batch). Return NoData as a fallback. + Ok(vec![]) + } + } + } _ => { // NoData Ok(vec![]) diff --git a/src/session/src/lib.rs b/src/session/src/lib.rs index 110560b3243..44f3ee66891 100644 --- a/src/session/src/lib.rs +++ b/src/session/src/lib.rs @@ -111,6 +111,15 @@ impl Session { .into() } + /// Returns the cursor with the given name, if it exists. + /// + /// Cursors are stored in the session's mutable inner data and shared across + /// all query contexts created from this session. + pub fn get_cursor(&self, name: &str) -> Option> { + let guard = self.mutable_inner.read().unwrap(); + guard.cursors.get(name).cloned() + } + pub fn conn_info(&self) -> &ConnInfo { &self.conn_info } From a7ba8b01e1c8ced20a19e39ac68a4637dc64e266 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 27 Aug 2026 04:24:07 +0000 Subject: [PATCH 90/92] fix(flow): restore FrontendClient::sql API (#8963) fix(flow): restore batching frontend SQL Signed-off-by: WenyXu --- src/flow/src/batching_mode/frontend_client.rs | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index 27db96748a6..48314dfa21b 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -18,6 +18,7 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex, RwLock, Weak}; use api::v1::greptime_request::Request; +use api::v1::query_request::Query; use api::v1::{CreateTableExpr, QueryRequest}; use client::{Client, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, Database, OutputWithMetrics}; use common_error::ext::BoxedError; @@ -339,6 +340,53 @@ impl FrontendClient { }) } + /// Execute a SQL statement on the frontend. + pub async fn sql(&self, catalog: &str, schema: &str, sql: &str) -> Result { + match self { + FrontendClient::Distributed { .. } => { + let db = self.get_random_active_frontend(catalog, schema).await?; + db.database + .sql(sql) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu) + } + FrontendClient::Standalone { + database_client, .. + } => { + let ctx = QueryContextBuilder::default() + .current_catalog(catalog.to_string()) + .current_schema(schema.to_string()) + .build(); + let ctx = Arc::new(ctx); + { + let database_client = { + database_client + .handler + .lock() + .unwrap() + .as_ref() + .context(UnexpectedSnafu { + reason: "Standalone's frontend instance is not set", + })? + .upgrade() + .context(UnexpectedSnafu { + reason: "Failed to upgrade database client", + })? + }; + let req = Request::Query(QueryRequest { + query: Some(Query::Sql(sql.to_string())), + }); + database_client + .do_query(req, ctx) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu) + } + } + } + } + /// Execute a flow query and return terminal metrics. `snapshot_seqs` are /// optional read upper bounds used only by snapshot-fenced repair chunks. pub(crate) async fn query_with_terminal_metrics( From beded6e232918147b93f91677e5ec8b8768d1be5 Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 27 Aug 2026 04:51:39 +0000 Subject: [PATCH 91/92] fix(flow): avoid insert select HTTP/2 stalls (#8962) * fix(flow): avoid insert select HTTP/2 stalls Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): share record batch forwarding for DML Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/cmd/src/frontend.rs | 2 + src/query/src/datafusion.rs | 126 +++++++++++++++++++++++++++++------- 2 files changed, 103 insertions(+), 25 deletions(-) diff --git a/src/cmd/src/frontend.rs b/src/cmd/src/frontend.rs index e3737088e8d..a303041b5a2 100644 --- a/src/cmd/src/frontend.rs +++ b/src/cmd/src/frontend.rs @@ -434,6 +434,8 @@ impl StartCommand { // Some queries are expected to take long time. let mut channel_config = opts.datanode.client.channel_config(); channel_config.timeout = None; + // Source Flight streams and sink unary responses share pooled connections. + channel_config.http2_adaptive_window = Some(true); if opts.grpc.flight_compression.transport_compression() { channel_config.accept_compression = true; channel_config.send_compression = true; diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index 754199a78fc..f9d87782a03 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -30,7 +30,7 @@ use common_function::function::FunctionContext; use common_function::function_factory::ScalarFunctionFactory; use common_query::{Output, OutputData, OutputMeta}; use common_recordbatch::adapter::{RecordBatchStreamAdapter, RegionQueryStatCounters}; -use common_recordbatch::{EmptyRecordBatchStream, SendableRecordBatchStream}; +use common_recordbatch::{EmptyRecordBatchStream, RecordBatch, SendableRecordBatchStream}; use common_telemetry::tracing; use datafusion::catalog::TableFunction; use datafusion::dataframe::DataFrame; @@ -44,6 +44,7 @@ use datafusion_expr::{ use datatypes::prelude::VectorRef; use datatypes::schema::Schema; use futures_util::StreamExt; +use futures_util::future::try_join; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt, ensure}; use sqlparser::ast::AnalyzeFormat; @@ -80,6 +81,29 @@ pub const QUERY_PARALLELISM_HINT: &str = "query_parallelism"; /// Whether to fallback to the original plan when failed to push down. pub const QUERY_FALLBACK_HINT: &str = "query_fallback"; +// An unbounded queue keeps draining source RPCs while mutation RPCs on a shared +// HTTP/2 connection are pending, trading bounded memory for request liveness. +async fn forward_record_batches( + mut stream: SendableRecordBatchStream, + batch_tx: tokio::sync::mpsc::UnboundedSender>, +) -> Result<()> { + while let Some(batch) = stream.next().await { + match batch.context(CreateRecordBatchSnafu) { + Ok(batch) => { + if batch_tx.send(Ok(batch)).is_err() { + break; + } + tokio::task::yield_now().await; + } + Err(error) => { + let _ = batch_tx.send(Err(error)); + break; + } + } + } + Ok(()) +} + fn query_load_region_id(plan: &Arc) -> Option { let mut region_id = None; let mut stack = vec![plan.clone()]; @@ -205,7 +229,7 @@ impl DatafusionQueryEngine { let Output { data, meta } = self .exec_query_plan((*dml.input).clone(), query_ctx.clone()) .await?; - let mut stream = match data { + let stream = match data { OutputData::RecordBatches(batches) => batches.as_stream(), OutputData::Stream(stream) => stream, _ => unreachable!(), @@ -214,30 +238,35 @@ impl DatafusionQueryEngine { let mut affected_rows = 0; let mut insert_cost = 0; - while let Some(batch) = stream.next().await { - let batch = batch.context(CreateRecordBatchSnafu)?; - let column_vectors = batch - .column_vectors(&table_name.to_string(), table.schema()) - .map_err(BoxedError::new) - .context(QueryExecutionSnafu)?; - - match dml.op { - WriteOp::Insert(_) => { - // We ignore the insert op. - let output = self - .insert(&table_name, column_vectors, query_ctx.clone()) - .await?; - let (rows, cost) = output.extract_rows_and_cost(); - affected_rows += rows; - insert_cost += cost; - } - WriteOp::Delete => { - affected_rows += self - .delete(&table_name, &table, column_vectors, query_ctx.clone()) - .await?; - } - _ => unreachable!("guarded by the 'ensure!' at the beginning"), + match dml.op { + WriteOp::Insert(_) => { + let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); + let producer = forward_record_batches(stream, batch_tx); + let consumer = self.consume_insert_record_batches( + batch_rx, + &table_name, + table.schema(), + query_ctx.clone(), + ); + let ((), (rows, cost)) = try_join(producer, consumer).await?; + affected_rows += rows; + insert_cost += cost; } + WriteOp::Delete => { + // Keep DELETE on the same producer/consumer schedule as INSERT so the source + // stream can continue draining while mutation RPCs are pending. + let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); + let producer = forward_record_batches(stream, batch_tx); + let consumer = self.consume_delete_record_batches( + batch_rx, + &table_name, + &table, + query_ctx.clone(), + ); + let (rows, ()) = try_join(consumer, producer).await?; + affected_rows += rows; + } + _ => unreachable!("guarded by the 'ensure!' at the beginning"), } Ok(Output::new( OutputData::AffectedRows(affected_rows), @@ -245,6 +274,53 @@ impl DatafusionQueryEngine { )) } + async fn consume_insert_record_batches( + &self, + mut batch_rx: tokio::sync::mpsc::UnboundedReceiver>, + table_name: &ResolvedTableReference, + table_schema: Arc, + query_ctx: QueryContextRef, + ) -> Result<(usize, usize)> { + let mut affected_rows = 0; + let mut insert_cost = 0; + while let Some(batch) = batch_rx.recv().await { + let batch = batch?; + let column_vectors = batch + .column_vectors(&table_name.to_string(), table_schema.clone()) + .map_err(BoxedError::new) + .context(QueryExecutionSnafu)?; + // We ignore the insert op. + let output = self + .insert(table_name, column_vectors, query_ctx.clone()) + .await?; + let (rows, cost) = output.extract_rows_and_cost(); + affected_rows += rows; + insert_cost += cost; + } + Ok((affected_rows, insert_cost)) + } + + async fn consume_delete_record_batches( + &self, + mut batch_rx: tokio::sync::mpsc::UnboundedReceiver>, + table_name: &ResolvedTableReference, + table: &TableRef, + query_ctx: QueryContextRef, + ) -> Result { + let mut affected_rows = 0; + while let Some(batch) = batch_rx.recv().await { + let batch = batch?; + let column_vectors = batch + .column_vectors(&table_name.to_string(), table.schema()) + .map_err(BoxedError::new) + .context(QueryExecutionSnafu)?; + affected_rows += self + .delete(table_name, table, column_vectors, query_ctx.clone()) + .await?; + } + Ok(affected_rows) + } + #[tracing::instrument(skip_all)] async fn delete( &self, From 9befb7e64cdccf25e96b36fa5ca8fa98f1901dd7 Mon Sep 17 00:00:00 2001 From: sun Date: Thu, 27 Aug 2026 09:57:18 +0000 Subject: [PATCH 92/92] feat: update dashboard to v0.13.14 (#8968) --- src/servers/dashboard/VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/servers/dashboard/VERSION b/src/servers/dashboard/VERSION index b68b26e6f78..0822cdb1e68 100644 --- a/src/servers/dashboard/VERSION +++ b/src/servers/dashboard/VERSION @@ -1 +1 @@ -v0.13.13 +v0.13.14