rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-11-15 07:17:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore: merge main into frontend heartbeat lifecycle

Browse Source 
Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>
This commit is contained in:
Lei, HUANG
2026-08-28 00:30:18 +08:00
parent e589a77a90 9befb7e64c
commit e685d42d16
498 changed files with 49577 additions and 9867 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.agents/README.md
+10 -3
View File
@@ -11,16 +11,23 @@ Shared agent skills live in `.agents/skills`.
- Add or update shared skills under `.agents/skills/<skill-name>/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
.agents/architecture-invariants.md
+25 -26
View File
@@ -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).
@@ -43,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.
@@ -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 `=<version>` 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
.agents/generated-files.md
+11 -3
View File
@@ -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`.
.agents/skills/greptimedb-release-note/SKILL.md
+4 -3
View File
@@ -167,9 +167,10 @@ The release note is also published as a blog post in **`GreptimeTeam/docs`**.
git -C <docs> fetch origin main
git -C <docs> 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
.cargo/config.toml
+3
View File
@@ -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 --"
.github/CODEOWNERS
+5 -5
View File
@@ -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
.github/actions/build-dev-builder-images/action.yml
+19
View File
@@ -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' }}
.github/runner-scale-sets/query-regression/README.md
+158 -295
View File
@@ -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 <region> --vswitch-id <vsw-...> --security-group-id <sg-...> \
--base-image-id <ubuntu-24.04-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=<app-id> \
--from-literal=github_app_installation_id=<installation-id> \
--from-file=github_app_private_key=<private-key.pem>
```
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.
.github/runner-scale-sets/query-regression/cache-pvc.yaml
-15
View File
@@ -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
.github/runner-scale-sets/query-regression/ecs-image/bootstrap-runner-host.sh
+195
View File
@@ -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=<token> bash bootstrap-runner-host.sh --repo-dir . \
# --register --repo-url https://github.com/<owner>/<repo>
#
# 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 <<EOF
Types: deb
URIs: https://download.docker.com/linux/ubuntu
Suites: ${UBUNTU_CODENAME:-$VERSION_CODENAME}
Components: stable
Architectures: $(dpkg --print-architecture)
Signed-By: /etc/apt/keyrings/docker.asc
EOF
apt-get update
apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin
step "Move docker data-root to ${DATA_ROOT}/docker"
systemctl stop docker containerd
mkdir -p "${DATA_ROOT}/docker" "${DATA_ROOT}/containerd"
if [[ -f /etc/docker/daemon.json ]] && ! grep -q '"data-root"' /etc/docker/daemon.json; then
echo "ERROR: /etc/docker/daemon.json exists without data-root; merge manually." >&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"
.github/runner-scale-sets/query-regression/ecs-image/build-ecs-image.py
+392
View File
@@ -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())
.github/runner-scale-sets/query-regression/ecs-image/ephemeral-github-runner.service
+45
View File
@@ -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
.github/runner-scale-sets/query-regression/ecs-image/start-runner.sh
+37
View File
@@ -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
.github/runner-scale-sets/query-regression/values-8-cores.yaml
-146
View File
@@ -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
.github/runner-scale-sets/query-regression/values-paused.yaml
-2
View File
@@ -1,2 +0,0 @@
minRunners: 0
maxRunners: 0
.github/scripts/aliyun-ecs-runner-provision.py
+485
View File
@@ -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"<details><summary>ECS console output ({instance_id})</summary>\n\n```\n")
summary.write(tail)
summary.write("\n```\n</details>\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())
.github/scripts/aliyun-ecs-runner-teardown.py
+237
View File
@@ -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-<run_id>).
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())
.github/scripts/pr-open-limit.js
+144
View File
@@ -0,0 +1,144 @@
// 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 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";
const MARKER = "<!-- pr-open-limit -->";
// GITHUB_TOKEN can only edit comments authored by the Actions bot itself.
const COMMENT_AUTHOR = "github-actions[bot]";
// 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;
function summary(text) {
if (process.env.GITHUB_STEP_SUMMARY) {
appendFileSync(process.env.GITHUB_STEP_SUMMARY, `${text}\n`);
}
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)
.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");
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 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 (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))) {
summary(`Skipping \`${author}\`: not a team member.`);
return;
}
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 && !pr.draft
);
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);
});
.github/scripts/query-regression-nightly-refs.py
+270
View File
@@ -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())
.github/workflows/checks.yml
+8
View File
@@ -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' }}
.github/workflows/pr-open-limit.yml
+48
View File
@@ -0,0 +1,48 @@
name: PR Open Limit
on:
pull_request_target:
types:
- opened
- reopened
- ready_for_review
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
cancel-in-progress: true
permissions:
contents: read
pull-requests: write
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_IS_DRAFT: ${{ github.event.pull_request.draft }}
MAX_OPEN_PRS: ${{ vars.MAX_OPEN_PRS_PER_AUTHOR || 5 }}
.github/workflows/query-regression-janitor.yml
+36
View File
@@ -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
.github/workflows/query-regression-nightly.yml
+104
View File
@@ -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
.github/workflows/query-regression.yml
+306 -127
View File
@@ -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 '<failed to run: %s>' "$*"
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="<unavailable>"
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:-<missing>}"
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
.github/workflows/release-dev-builder-images.yaml
+49 -1
View File
@@ -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
.github/workflows/release.yml
+52 -2
View File
@@ -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:
@@ -464,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
@@ -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
.github/workflows/rust.yml
+26
View File
@@ -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
AGENTS.md
+56 -12
View File
@@ -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 <package>` (preferred over `cargo test`) |
| Full local Rust suite | `make test` |
| SQL tests | `cargo sqlness bare` (single case: `cargo sqlness bare -t <name>`) |
| 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,24 +88,51 @@ 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 <package>` |
| 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 <case>`; 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`)
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.
Cargo.lock
Generated
+510 -318
View File
File diff suppressed because it is too large Load Diff
Cargo.toml
+28 -19
View File
@@ -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 = "32f467fa2ba2b3588a58381a24af83de09fbb00a" }
hex = "0.4"
http = "1"
humantime = "2.1"
@@ -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",
@@ -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"
@@ -202,6 +205,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"
@@ -226,10 +231,12 @@ 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"] }
serde_with = "3"
serde_yaml_ng = "0.10"
simd-json = "0.15"
similar-asserts = "1.6.0"
smallvec = { version = "1", features = ["serde"] }
@@ -248,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"] }
@@ -260,11 +267,13 @@ 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"
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
@@ -343,21 +352,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 = "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]
Makefile
+27 -1
View File
@@ -8,7 +8,10 @@ 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
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
README.md
+49 -38
View File
@@ -6,10 +6,11 @@
</picture>
</p>
<h2 align="center">One database for metrics, logs, and traces<br/>
replacing Prometheus, Loki, and Elasticsearch</h2>
<h2 align="center">Metrics, logs, and traces.<br/>
One engine, on your infrastructure.</h2>
> 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.
<div align="center">
<h3 align="center">
@@ -28,7 +29,7 @@ replacing Prometheus, Loki, and Elasticsearch</h2>
<img src="https://img.shields.io/github/v/release/GreptimeTeam/greptimedb?include_prereleases&filter=*-nightly-*&label=nightly&color=orange" alt="Nightly"/>
</a>
<sub><b>stable</b> for production &nbsp;·&nbsp; <b>latest</b> includes pre-releases &nbsp;·&nbsp; <b>nightly</b> is a weekly snapshot of <code>main</code></sub>
<sub><b>stable</b> for production &nbsp;·&nbsp; <b>canary</b> includes pre-releases &nbsp;·&nbsp; <b>nightly</b> is a weekly snapshot of <code>main</code></sub>
<a href="https://hub.docker.com/r/greptime/greptimedb/">
<img src="https://img.shields.io/docker/pulls/greptime/greptimedb.svg" alt="Docker Pulls"/>
@@ -57,9 +58,11 @@ replacing Prometheus, Loki, and Elasticsearch</h2>
</div>
- [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</h2>
## 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
</a>
</p>
## 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).
@@ -203,13 +213,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).
@@ -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
config/config.md
+13 -6
View File
@@ -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.<br/>Defaults to max(num_cpus / 2, 1). |
| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.<br/>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.<br/>When Prometheus pending-row batching is enabled, a nonzero timeout less than or equal to the<br/>`prom_store.pending_rows_flush_interval` plus 1 second is adjusted to that value. |
@@ -74,7 +74,9 @@
| `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. |
| `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric<br/>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. |
@@ -185,6 +187,7 @@
| `region_engine.mito.min_compaction_interval` | String | `0m` | Minimum time interval between two compactions.<br/>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.<br/><br/>Setting this to "true" is a necessary but not sufficient condition for scheduling compaction after a region edit.<br/>Other constraints, such as "min_compaction_interval", may still prevent compaction from being scheduled.<br/>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<br/>creating the index and staging files for searching the index, defaults to `{data_home}/index_intermediate`.<br/>The default name for this directory is `index_intermediate` for backward compatibility.<br/><br/>This path contains two subdirectories:<br/>- `__intm`: for storing intermediate files used during creating index.<br/>- `staging`: for storing staging files used during searching index. |
| `region_engine.mito.index.staging_size` | String | `2GB` | The max capacity of the staging directory. |
@@ -234,7 +237,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`,<br/>`alter_database`, `drop_database`, `create_flow`, `drop_flow`,<br/>`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,<br/>`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,<br/>`create_view`, `drop_view`.<br/>When omitted, all current and future event types are recorded.<br/>Set to an empty array to disable event recording. |
| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `create_database`,<br/>`alter_database`, `drop_database`, `create_flow`, `drop_flow`,<br/>`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,<br/>`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,<br/>`create_view`, `drop_view`, `admin_function`.<br/>When omitted, all current and future event types are recorded.<br/>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.<br/>When enabled, heap profiling will be activated if the `MALLOC_CONF` environment variable<br/>is set to "prof:true,prof_active:false". The official image adds this env variable.<br/>Default is true. |
@@ -254,7 +257,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.<br/>Defaults to max(num_cpus / 2, 1). |
| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.<br/>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.<br/>When Prometheus pending-row batching is enabled, a nonzero timeout less than or equal to the<br/>`prom_store.pending_rows_flush_interval` plus 1 second is adjusted to that value. |
@@ -315,7 +318,9 @@
| `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. |
| `otlp.experimental_enable_resource_info` | Bool | `false` | Whether to synthesize the `greptime_otel_resource_info` table from OTLP metric<br/>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. |
@@ -372,6 +377,7 @@
| `memory.enable_heap_profiling` | Bool | `true` | Whether to enable heap profiling activation during startup.<br/>When enabled, heap profiling will be activated if the `MALLOC_CONF` environment variable<br/>is set to "prof:true,prof_active:false". The official image adds this env variable.<br/>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`.<br/>When omitted, all current and future event types are recorded.<br/>Set to an empty array to disable event recording. |
### Metasrv
@@ -396,7 +402,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.<br/>Defaults to max(num_cpus / 2, 1). |
| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.<br/>Defaults to max(num_cpus / 2, 2). |
| `backend_tls` | -- | -- | TLS configuration for kv store backend (applicable for etcd, PostgreSQL, and MySQL backends)<br/>When using etcd, PostgreSQL, or MySQL as metadata store, you can configure TLS here<br/><br/>Note: if TLS is configured in both this section and the `store_addrs` connection string, the<br/>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<br/>- "disable" - No TLS<br/>- "prefer" (default) - Try TLS, fallback to plain<br/>- "require" - Require TLS<br/>- "verify_ca" - Require TLS and verify CA<br/>- "verify_full" - Require TLS and verify hostname |
| `backend_tls.cert_path` | String | `""` | Path to client certificate file (for client authentication)<br/>Like "/path/to/client.crt" |
@@ -507,8 +513,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.<br/>Defaults to max(num_cpus / 2, 1). |
| `runtime.query_rt_size` | Integer | `7` | The number of threads to execute datanode query operations.<br/>Defaults to max(num_cpus - 1, 1). |
| `runtime.compact_rt_max_blocking_threads` | Integer | `4` | The maximum number of blocking threads for compact operations.<br/>Defaults to max(num_cpus / 2, 2). |
| `runtime.query_rt_size` | Integer | `7` | The number of threads to execute datanode query operations.<br/>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. |
@@ -605,6 +611,7 @@
| `region_engine.mito.min_compaction_interval` | String | `0m` | Minimum time interval between two compactions.<br/>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.<br/><br/>Setting this to "true" is a necessary but not sufficient condition for scheduling compaction after a region edit.<br/>Other constraints, such as "min_compaction_interval", may still prevent compaction from being scheduled.<br/>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<br/>creating the index and staging files for searching the index, defaults to `{data_home}/index_intermediate`.<br/>The default name for this directory is `index_intermediate` for backward compatibility.<br/><br/>This path contains two subdirectories:<br/>- `__intm`: for storing intermediate files used during creating index.<br/>- `staging`: for storing staging files used during searching index. |
| `region_engine.mito.index.staging_size` | String | `2GB` | The max capacity of the staging directory. |
config/datanode.example.toml
+5 -2
View File
@@ -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
@@ -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]
config/frontend.example.toml
+10 -1
View File
@@ -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.
@@ -237,8 +237,13 @@ 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
## 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]
@@ -405,3 +410,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"]
config/metasrv.example.toml
+1 -1
View File
@@ -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)
config/standalone.example.toml
+10 -2
View File
@@ -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.
@@ -204,8 +204,13 @@ 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
## 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]
@@ -729,6 +734,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]
@@ -922,7 +930,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"]
docker/dev-builder/riscv64/Dockerfile
+55
View File
@@ -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
docs/parquet-tools.md
+255
View File
@@ -0,0 +1,255 @@
# 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 <COMMAND> --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`. 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
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 <ROWS>` 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
`<table-id>:<region-number>`. `--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 <SVG>` 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.
docs/rfcs/2026-06-25-entity-relationships-and-graph-query.md
+4 -4
View File
@@ -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 `<qualifier>/<value>`, 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.
docs/rfcs/2026-08-04-native-histograms.md
+82 -27
View File
@@ -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.
src/auth/src/error.rs
+8
View File
@@ -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,
}
}
src/auth/src/lib.rs
+1 -1
View File
@@ -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;
src/auth/src/permission.rs
+5
View File
@@ -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)]
src/auth/src/user_info.rs
+47 -48
View File
@@ -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<Self> {
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,30 +219,10 @@ 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);
}
}
#[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::<DefaultUserInfo>()
.unwrap();
assert_eq!(default_user.permission_mode, PermissionMode::ReadOnly);
}
#[test]
fn test_user_info_as_any() {
let user_info = DefaultUserInfo::with_name("test_user");
src/auth/src/user_provider.rs
+36 -3
View File
@@ -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<UserInfoRef> {
UnsupportedAuthMethodSnafu {
method: "bearer token",
}
.fail()
}
async fn postgres_auth_info(&self, _id: Identity<'_>) -> Result<PgAuthInfo> {
Ok(PgAuthInfo::Cleartext)
}
@@ -438,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())
};
@@ -741,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!(
src/auth/src/user_provider/static_user_provider.rs
+4
View File
@@ -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;
}
}
src/catalog/src/kvbackend/manager.rs
+6 -2
View File
@@ -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
}
src/catalog/src/system_schema/semantic_graph.rs
+88 -46
View File
@@ -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,
@@ -36,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;
@@ -49,6 +53,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 +64,18 @@ use crate::system_schema::{SystemSchemaProviderInner, SystemTable, SystemTableRe
pub type EntityGraphProviderRef = Arc<dyn EntityGraphProvider>;
/// 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 +84,7 @@ pub trait EntityGraphProvider: Send + Sync {
&self,
catalog: &str,
request: ScanRequest,
query_ctx: Option<QueryContextRef>,
) -> std::result::Result<Option<SendableRecordBatchStream>, BoxedError>;
/// Produces the relationship set (`semantic_relationships`) rows for `catalog`.
@@ -85,6 +92,7 @@ pub trait EntityGraphProvider: Send + Sync {
&self,
catalog: &str,
request: ScanRequest,
query_ctx: Option<QueryContextRef>,
) -> std::result::Result<Option<SendableRecordBatchStream>, BoxedError>;
}
@@ -94,13 +102,21 @@ pub trait EntityGraphProvider: Send + Sync {
pub(crate) struct SemanticGraphTableProvider {
catalog_name: String,
catalog_manager: Weak<dyn CatalogManager>,
/// 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<QueryContextRef>,
}
impl SemanticGraphTableProvider {
pub(crate) fn new(catalog_name: String, catalog_manager: Weak<dyn CatalogManager>) -> Self {
pub(crate) fn new(
catalog_name: String,
catalog_manager: Weak<dyn CatalogManager>,
query_ctx: Option<QueryContextRef>,
) -> Self {
Self {
catalog_name,
catalog_manager,
query_ctx,
}
}
@@ -139,6 +155,7 @@ impl SystemSchemaProviderInner for SemanticGraphTableProvider {
kind,
self.catalog_name.clone(),
self.catalog_manager.clone(),
self.query_ctx.clone(),
)) as _)
}
}
@@ -178,16 +195,16 @@ fn json() -> ConcreteDataType {
/// - `source_tables` — JSON array of the telemetry tables that contributed this entity.
static ENTITIES_SCHEMA: LazyLock<SchemaRef> = 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),
]))
});
@@ -220,22 +237,38 @@ static ENTITIES_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
/// `db.system`, `peer.service`.
static RELATIONSHIPS_SCHEMA: LazyLock<SchemaRef> = 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),
]))
});
@@ -252,6 +285,7 @@ struct SemanticGraphTable {
schema: SchemaRef,
catalog_name: String,
catalog_manager: Weak<dyn CatalogManager>,
query_ctx: Option<QueryContextRef>,
}
impl SemanticGraphTable {
@@ -259,6 +293,7 @@ impl SemanticGraphTable {
kind: GraphTableKind,
catalog_name: String,
catalog_manager: Weak<dyn CatalogManager>,
query_ctx: Option<QueryContextRef>,
) -> Self {
let schema = match kind {
GraphTableKind::Entities => ENTITIES_SCHEMA.clone(),
@@ -269,6 +304,7 @@ impl SemanticGraphTable {
schema,
catalog_name,
catalog_manager,
query_ctx,
}
}
@@ -277,6 +313,7 @@ impl SemanticGraphTable {
catalog: String,
catalog_manager: Weak<dyn CatalogManager>,
request: ScanRequest,
query_ctx: Option<QueryContextRef>,
) -> Result<Option<SendableRecordBatchStream>> {
let provider = utils::entity_graph_provider(&catalog_manager)?;
// No provider (engine not up / non-frontend node): stream empty.
@@ -284,8 +321,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 +379,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)?;
src/client/src/database.rs
+336 -154
View File
@@ -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<Box<dyn Stream<Item = FlightData> + Send>>;
type DoPutResponseStream = Pin<Box<dyn Stream<Item = Result<DoPutResponse>>>>;
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<S>(
mut flight_message_stream: S,
remote_addr: String,
flight_message_stream: S,
) -> Result<OutputWithMetrics>
where
S: Stream<Item = Result<FlightMessage>> + 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<String>,
flow_extensions: Option<String>,
snapshot_seqs: Option<String>,
timeout: Option<Duration>,
}
impl FlightRequestOptions {
fn apply_to<T>(self, request: &mut tonic::Request<T>) -> 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<Channel>,
@@ -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<u32> {
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<String> {
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<u64, u64>,
) -> Result<()> {
if snapshot_seqs.is_empty() {
return Ok(());
}
fn encode_flow_extensions(flow_extensions: &[(&str, &str)]) -> Option<String> {
(!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<u64, u64>) -> Option<String> {
(!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<str>,
{
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<str>,
{
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<str>,
{
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<u8>) -> Result<Output> {
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<u64, u64>,
) -> Result<OutputWithMetrics> {
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<Output> {
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<Output> {
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<u64, u64>,
options: FlightRequestOptions,
) -> Result<OutputWithMetrics> {
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<u64, u64>) -> 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<S>(self, sql: S) -> Result<Output>
where
S: AsRef<str>,
{
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<S>(self, sql: S) -> Result<OutputWithMetrics>
where
S: AsRef<str>,
{
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<u8>) -> Result<Output> {
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<OutputWithMetrics> {
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<Output> {
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<Output> {
self.do_get(Request::Ddl(DdlRequest {
expr: Some(DdlExpr::AlterTable(expr)),
}))
.await
.map(OutputWithMetrics::into_output)
}
async fn do_get(self, request: Request) -> Result<OutputWithMetrics> {
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::<Vec<(String, String)>>(
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::<HashMap<u64, u64>>(
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<Result<FlightMessage>>))
})] as Vec<Result<FlightMessage>>),
)
.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<Result<FlightMessage>>))
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<Result<FlightMessage>>),
)
.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<Result<FlightMessage>>))
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<Result<FlightMessage>>),
)
.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<Result<FlightMessage>>))
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<Result<FlightMessage>>),
)
.await
.unwrap();
let terminal_metrics = output.metrics.clone();
src/client/src/error.rs
+18
View File
@@ -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<Code> {
match self {
Self::FlightGet { tonic_code, .. }
| Self::FlightStream { tonic_code, .. }
| Self::RegionServer {
code: tonic_code, ..
}
src/client/src/flight.rs
+72 -2
View File
@@ -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<S: Stream + Unpin> {
/// Remote Flight peer associated with this response stream.
remote_addr: String,
messages: Peekable<S>,
}
impl<S> FlightMessageReader<S>
where
S: Stream<Item = Result<FlightMessage>> + Unpin,
{
pub(crate) fn new(remote_addr: impl Into<String>, 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<FlightMessage> {
self.read_next().await?.context(IllegalFlightMessagesSnafu {
reason: "Expect the response not to be empty",
})
}
pub(crate) async fn read_next(&mut self) -> Result<Option<FlightMessage>> {
self.messages.next().await.transpose()
}
pub(crate) async fn peek_next_message_kind(&mut self) -> Result<Option<FlightMessageKind>> {
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,
src/client/src/region.rs
+222 -65
View File
@@ -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};
@@ -46,9 +47,11 @@ 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);
#[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<RegionResponse> {
@@ -195,22 +202,23 @@ impl RegionRequester {
}
async fn recordbatches_from_flight_message_stream<S>(
mut flight_message_stream: S,
addr: String,
flight_message_stream: S,
) -> Result<SendableRecordBatchStream>
where
S: Stream<Item = Result<FlightMessage>> + 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 FlightMessage::Schema(schema) = first_flight_message? else {
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();
.fail()
.map_err(|error| flight_stream_error(reader.remote_addr(), error));
};
let metrics = Arc::new(ArcSwapOption::from(None));
@@ -221,29 +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;
let stream = Box::pin(stream!({
let _span = tracing_context.attach(common_telemetry::tracing::info_span!(
"poll_flight_data_stream"
));
let mut buffered_message: Option<FlightMessage> = 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(e)).context(ExternalSnafu);
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 {
@@ -251,41 +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"
// 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;
}
.fail()
.map_err(BoxedError::new)
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;
}
Err(e) => {
yield Err(BoxedError::new(e)).context(ExternalSnafu);
break;
}
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;
@@ -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<String>,
update: RemoteDynFilterUpdate,
@@ -380,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;
@@ -389,6 +422,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::<Result<FlightMessage>>(),
)
.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<Schema> {
Arc::new(Schema::new(vec![ColumnSchema::new(
"v",
@@ -509,11 +580,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();
@@ -525,14 +599,97 @@ 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::<Int32Array>()
.unwrap()
.value(0),
1
);
assert_eq!(
second_batch
.column(0)
.as_any()
.downcast_ref::<Int32Array>()
.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();
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();
src/cmd/Cargo.toml
+1 -1
View File
@@ -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"]
src/cmd/src/bin/greptime.rs
+6
View File
@@ -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)?)
src/cmd/src/bin/query_regression_runner/direct.rs
+104 -2
View File
@@ -166,7 +166,28 @@ fn create_table_sql(table: &Table) -> Result<String> {
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::<Vec<_>>()
.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;"
);
}
src/cmd/src/bin/query_regression_runner/measure.rs
+59 -5
View File
@@ -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([
src/cmd/src/bin/query_regression_runner/model.rs
+13
View File
@@ -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<String>,
}
const fn default_show_create_engine() -> bool {
@@ -203,6 +209,13 @@ pub(super) struct Query {
#[serde(default)]
pub(super) kind: Option<String>,
pub(super) query: String,
/// Prometheus HTTP range query parameters (`kind = "prom_http"` only).
#[serde(default)]
pub(super) start: Option<String>,
#[serde(default)]
pub(super) end: Option<String>,
#[serde(default)]
pub(super) step: Option<String>,
#[serde(default)]
pub(super) warmup: usize,
#[serde(default = "one")]
src/cmd/src/bin/query_regression_runner/sql.rs
+66 -6
View File
@@ -117,10 +117,62 @@ pub(super) fn value_f64(value: Option<&Value>) -> Option<f64> {
}
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<_>>(),
vec![("db".into(), "catalog-schema name".into())]
);
}
#[test]
fn top_level_errors_do_not_inspect_rows() {
assert!(response_has_error(&json!({"error_code": 7})));
src/cmd/src/datanode.rs
+46
View File
@@ -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);
}
}
}
}
src/cmd/src/datanode/builder.rs
+3 -3
View File
@@ -61,9 +61,6 @@ impl InstanceBuilder {
}
async fn init(opts: &mut DatanodeOptions, plugins: &mut Plugins) -> Result<Vec<WorkerGuard>> {
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);
src/cmd/src/datanode/objbench.rs
+8 -99
View File
@@ -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<PathBuf>,
}
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::<u64>()
})
.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<FileDirComponents> {
})
}
pub(crate) fn extract_region_metadata(
file_path: &str,
meta: &parquet::file::metadata::ParquetMetaData,
) -> error::Result<RegionMetadataRef> {
let kvs: Option<&Vec<KeyValue>> = 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<ObjectStore> {
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() {
src/cmd/src/datanode/parquet_meta.rs
+394
View File
@@ -0,0 +1,394 @@
// 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 parquet::file::page_index::offset_index::PageLocation;
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<String, Option<String>>,
row_groups: Vec<RowGroupMetaView>,
}
#[derive(Debug, Serialize)]
struct RowGroupMetaView {
index: usize,
num_rows: i64,
uncompressed_size: i64,
compressed_size: i64,
compression_ratio: Option<f64>,
data_pages: Option<usize>,
dictionary_page_bytes: i64,
columns: Vec<ColumnChunkMetaView>,
}
#[derive(Debug, Serialize)]
struct ColumnChunkMetaView {
index: usize,
path: String,
physical_type: String,
encodings: Vec<String>,
compression: String,
num_values: i64,
uncompressed_size: i64,
compressed_size: i64,
compression_ratio: Option<f64>,
data_page_offset: i64,
dictionary_page_offset: Option<i64>,
dictionary_page_bytes: Option<i64>,
data_pages: Option<usize>,
max_data_page_row_count: Option<i64>,
has_statistics: bool,
column_index_offset: Option<i64>,
column_index_length: Option<i32>,
offset_index_offset: Option<i64>,
offset_index_length: Option<i32>,
bloom_filter_offset: Option<i64>,
bloom_filter_length: Option<i32>,
}
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 page_locations = offset_index
.and_then(|index| index.get(row_group_idx))
.and_then(|columns| columns.get(column_idx))
.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(),
);
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,
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(),
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={}, 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,
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.max_data_page_row_count),
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 max_data_page_row_count(
page_locations: &[PageLocation],
row_group_row_count: i64,
) -> Option<i64> {
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<i64>,
data_page_offset: i64,
) -> Option<i64> {
dictionary_page_offset.map(|offset| data_page_offset.saturating_sub(offset))
}
fn compression_ratio(uncompressed: i64, compressed: i64) -> Option<f64> {
(uncompressed > 0).then(|| compressed as f64 / uncompressed as f64)
}
fn format_ratio(ratio: Option<f64>) -> String {
ratio
.map(|value| format!("{value:.4}"))
.unwrap_or_else(|| "unknown".to_string())
}
fn format_optional_usize(value: Option<usize>) -> String {
value
.map(|value| value.to_string())
.unwrap_or_else(|| "unknown".to_string())
}
fn format_optional_i64(value: Option<i64>) -> String {
value
.map(|value| value.to_string())
.unwrap_or_else(|| "none".to_string())
}
fn format_optional_i32(value: Option<i32>) -> 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, None);
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);
}
#[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<usize>) {
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 mut props =
WriterProperties::builder().set_key_value_metadata(Some(vec![KeyValue::new(
"greptime:test".to_string(),
"value".to_string(),
)]));
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();
writer.close().unwrap();
}
}
src/cmd/src/datanode/parquet_rewrite.rs
+632
View File
@@ -0,0 +1,632 @@
// 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<PathBuf>,
/// Path to writer properties TOML in rewrite mode.
#[clap(long, value_name = "FILE")]
properties: Option<PathBuf>,
/// Dump writer properties TOML inferred from the input parquet file.
#[clap(long, value_name = "FILE")]
dump_properties: Option<PathBuf>,
/// Number of rows per record batch.
#[clap(long)]
batch_size: Option<usize>,
/// 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<ColumnConfig>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
#[serde(default, deny_unknown_fields)]
struct WriterConfig {
dictionary_enabled: Option<bool>,
compression: Option<CompressionConfig>,
compression_level: Option<u32>,
encoding: Option<EncodingConfig>,
max_row_group_row_count: Option<usize>,
data_page_size_limit: Option<usize>,
data_page_row_count_limit: Option<usize>,
dictionary_page_size_limit: Option<usize>,
column_index_truncate_length: Option<usize>,
statistics_truncate_length: Option<usize>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
#[serde(default, deny_unknown_fields)]
struct ColumnConfig {
path: Vec<String>,
dictionary_enabled: Option<bool>,
compression: Option<CompressionConfig>,
compression_level: Option<u32>,
encoding: Option<EncodingConfig>,
}
#[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 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()
.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,
column_index_truncate_length: truncate_length,
statistics_truncate_length: truncate_length,
},
columns,
}
}
fn load_rewrite_properties(path: &Path) -> error::Result<RewriteProperties> {
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<Vec<parquet::file::metadata::KeyValue>>,
) -> error::Result<WriterProperties> {
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);
}
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() {
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> {
File::open(path).context(error::FileIoSnafu)
}
fn create_file(path: &Path) -> error::Result<File> {
File::create(path).context(error::FileIoSnafu)
}
fn illegal_config<T>(msg: String) -> error::Result<T> {
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<u32>,
) -> error::Result<Compression> {
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<EncodingConfig> {
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, Some("greptime:test"));
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_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()]);
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, metadata_key: Option<&str>) {
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 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(key_value_metadata)
.build();
let mut writer =
ArrowWriter::try_new(File::create(path).unwrap(), schema, Some(props)).unwrap();
writer.write(&batch).unwrap();
writer.close().unwrap();
}
}
src/cmd/src/datanode/parquetbench.rs
+293 -108
View File
@@ -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<PathBuf>,
/// 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<String>,
/// Table directory relative to data home (e.g. "data/greptime/public/1024/")
#[clap(long)]
table_dir: String,
table_dir: Option<String>,
/// SST file id to benchmark.
#[clap(long)]
file_id: String,
file_id: Option<String>,
/// Local parquet SST file to benchmark with the direct reader.
#[clap(long, value_name = "FILE")]
file_path: Option<PathBuf>,
/// 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<PathType>,
table_dir: Option<String>,
}
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, &region_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<ParquetbenchInput> {
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<ParquetScanConfig> {
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<ParquetbenchSource> {
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, &region_id, table_dir, &file_id, path_type).await,
}
}
fn build_local_file_source(file_path: &Path) -> error::Result<ParquetbenchSource> {
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<ParquetbenchSource> {
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,
@@ -554,7 +766,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();
@@ -773,60 +985,6 @@ fn parse_batch_size(s: &str) -> Result<usize, String> {
Ok(batch_size)
}
fn parse_region_id(s: &str) -> error::Result<RegionId> {
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<PathType> {
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]
src/cmd/src/datanode/scanbench.rs
+23 -56
View File
@@ -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;
@@ -52,13 +54,24 @@ 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;
/// Displays a scanner using DataFusion's verbose explain format.
struct VerboseScannerDisplay<'a, T: ?Sized>(&'a T);
impl<T: DisplayAs + ?Sized> 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 {
@@ -164,60 +177,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<RegionId> {
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<PathType> {
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,
@@ -740,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
src/cmd/src/datanode/sst_replace.rs
+361
View File
@@ -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<PathBuf>,
/// Object-store parquet path used as the replacement.
#[clap(long, value_name = "PATH", conflicts_with = "replacement_file")]
replacement_object: Option<String>,
/// 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, &region_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<Replacement<'_>> {
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<Vec<u8>> {
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<parquet::file::metadata::ParquetMetaData, Box<dyn std::error::Error + Send + Sync>> {
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<Option<PathType>> {
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());
}
}
src/cmd/src/datanode/tool_util.rs
+236
View File
@@ -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<ObjectStore> {
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<RegionMetadataRef> {
let key_values: Option<&Vec<KeyValue>> = 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<RegionId> {
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> {
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<PathType> {
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::<u64>()
})
.max()
.unwrap_or(0)
}
#[cfg(feature = "dev-tools")]
pub(crate) fn load_local_parquet_metadata(path: &Path) -> error::Result<ParquetMetaData> {
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);
}
}
src/cmd/src/flownode.rs
+2 -2
View File
@@ -248,8 +248,6 @@ impl StartCommand {
}
async fn build(&self, opts: FlownodeOptions) -> Result<Instance> {
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);
src/cmd/src/frontend.rs
+4 -2
View File
@@ -370,8 +370,6 @@ impl StartCommand {
opts: FrontendOptions,
plugins: Plugins,
) -> Result<Instance> {
common_runtime::init_global_runtimes(&opts.runtime);
let guard = common_telemetry::init_global_logging(
APP_NAME,
&opts.component.logging,
@@ -380,6 +378,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);
@@ -467,6 +467,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;
src/cmd/src/metasrv.rs
+2 -2
View File
@@ -318,8 +318,6 @@ impl StartCommand {
}
pub async fn build(&self, opts: MetasrvOptions) -> Result<Instance> {
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);
src/cmd/src/standalone.rs
+9 -5
View File
@@ -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};
@@ -393,8 +395,6 @@ impl StartCommand {
#[allow(clippy::diverging_sub_expression)]
/// Build GreptimeDB instance with the loaded options.
pub async fn build(&self, opts: GreptimeOptions<StandaloneOptions>) -> Result<Instance> {
common_runtime::init_global_runtimes(&opts.runtime);
let guard = common_telemetry::init_global_logging(
APP_NAME,
&opts.component.logging,
@@ -403,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);
@@ -1127,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]
src/cmd/tests/load_config_test.rs
+54
View File
@@ -153,6 +153,12 @@ fn test_load_frontend_example_config() {
let options =
GreptimeOptions::<FrontendOptions>::load_layered_options(example_config.to_str(), "")
.unwrap();
assert!(
!options
.component
.otlp
.experimental_enable_exponential_histogram
);
let expected = GreptimeOptions::<FrontendOptions> {
component: FrontendOptions {
default_timezone: Some("UTC".to_string()),
@@ -329,6 +335,12 @@ fn test_load_standalone_example_config() {
let options =
GreptimeOptions::<StandaloneOptions>::load_layered_options(example_config.to_str(), "")
.unwrap();
assert!(
!options
.component
.otlp
.experimental_enable_exponential_histogram
);
let expected = GreptimeOptions::<StandaloneOptions> {
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::<FrontendOptions>::load_layered_options(config.path().to_str(), "")
.unwrap();
assert!(
frontend
.component
.otlp
.experimental_enable_exponential_histogram
);
let standalone =
GreptimeOptions::<StandaloneOptions>::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();
@@ -445,6 +494,11 @@ fn test_load_event_types_from_env() {
.component
.event_recorder
.event_types,
GreptimeOptions::<FrontendOptions>::load_layered_options(None, env_prefix)
.unwrap()
.component
.event_recorder
.event_types,
] {
assert!(event_types.allows("region_migration"));
assert!(!event_types.allows("other_event"));
src/common/base/Cargo.toml
+1
View File
@@ -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"] }
src/common/base/src/lib.rs
+1
View File
@@ -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;
src/common/base/src/protocol.rs
+106
View File
@@ -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<u32> 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<str> 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());
}
}
src/common/catalog/src/consts.rs
+45
View File
@@ -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,41 @@ 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";
/// 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";
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 ----
src/common/event-recorder/Cargo.toml
+1
View File
@@ -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
src/common/event-recorder/src/context.rs
+223
View File
@@ -0,0 +1,223 @@
// 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::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<String, String>,
}
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<PbProcedureEventContext> 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 {
pub reason: TriggerReason,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub protocol: Option<String>,
#[serde(default, skip_serializing_if = "serde_json::Map::is_empty")]
pub extensions: serde_json::Map<String, serde_json::Value>,
}
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<String>) -> Self {
self.protocol = Some(protocol.into());
self
}
}
impl From<(ProcedureEventInput, Option<String>)> for PersistentEventContext {
fn from((input, protocol): (ProcedureEventInput, Option<String>)) -> 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<PbProcedureEventContext> 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,
)]
#[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_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();
assert_eq!(TriggerReason::Unknown, reason);
}
}
src/common/event-recorder/src/event_table.rs
+22 -242
View File
@@ -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);
@@ -288,245 +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 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(
[
&REGION_ID_COLUMN,
&REGION_NUMBER_COLUMN,
&REGION_MIGRATION_TRIGGER_REASON_COLUMN,
&REGION_MIGRATION_SRC_NODE_ID_COLUMN,
&REGION_MIGRATION_SRC_PEER_ADDR_COLUMN,
&REGION_MIGRATION_DST_NODE_ID_COLUMN,
&REGION_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<const N: usize>(
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()
})
);
}
}
src/common/event-recorder/src/lib.rs
+2
View File
@@ -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, ProcedureEventInput, TriggerReason};
pub use recorder::*;
src/common/function/Cargo.toml
+7 -1
View File
@@ -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
src/common/function/benches/uddsketch.rs
+201
View File
@@ -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<Arc<dyn PhysicalExpr>>,
expr_fields: Vec<Arc<Field>>,
return_field: Arc<Field>,
}
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<Arc<dyn PhysicalExpr>> = 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::<Vec<_>>();
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<dyn Accumulator> {
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<ArrayRef> {
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::<Vec<_>>();
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::<Vec<_>>();
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);
src/common/function/src/admin.rs
+3 -1
View File
@@ -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());
}
src/common/function/src/admin/discard_unflushed_data.rs
+232
View File
@@ -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<Value> {
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(&params[0])? else {
return UnsupportedInputDataTypeSnafu {
function: "discard_unflushed",
datatypes: params
.iter()
.map(|value| value.data_type())
.collect::<Vec<_>>(),
}
.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::<Vec<_>>(),
),
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::<UInt64Array>().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::<UInt64Array>().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"));
}
}
src/common/function/src/admin/gc.rs
+35 -17
View File
@@ -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<Value> {
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<Option<String>> {
async fn migrate_region(
&self,
_query_ctx: QueryContextRef,
_request: MigrateRegionRequest,
) -> Result<Option<String>> {
unreachable!()
}
@@ -297,12 +307,20 @@ mod tests {
unreachable!()
}
async fn gc_regions(&self, request: GcRegionsRequest) -> Result<GcResponse> {
async fn gc_regions(
&self,
_query_ctx: QueryContextRef,
request: GcRegionsRequest,
) -> Result<GcResponse> {
*self.gc_regions_request.lock().unwrap() = Some(request);
Ok(GcResponse::default())
}
async fn gc_table(&self, request: GcTableRequest) -> Result<GcResponse> {
async fn gc_table(
&self,
_query_ctx: QueryContextRef,
request: GcTableRequest,
) -> Result<GcResponse> {
*self.gc_table_request.lock().unwrap() = Some(request);
Ok(GcResponse::default())
}
src/common/function/src/admin/migrate_region.rs
+10 -7
View File
@@ -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<Value> {
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 {
src/common/function/src/admin/purge_table.rs
+9 -5
View File
@@ -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<Option<String>> {
async fn migrate_region(
&self,
_: QueryContextRef,
_: MigrateRegionRequest,
) -> Result<Option<String>> {
unreachable!()
}
async fn reconcile(&self, _: ReconcileRequest) -> Result<Option<String>> {
@@ -159,10 +163,10 @@ mod tests {
fn catalog_manager(&self) -> &CatalogManagerRef {
unreachable!()
}
async fn gc_regions(&self, _: GcRegionsRequest) -> Result<GcResponse> {
async fn gc_regions(&self, _: QueryContextRef, _: GcRegionsRequest) -> Result<GcResponse> {
unreachable!()
}
async fn gc_table(&self, _: GcTableRequest) -> Result<GcResponse> {
async fn gc_table(&self, _: QueryContextRef, _: GcTableRequest) -> Result<GcResponse> {
unreachable!()
}
}
src/common/function/src/aggrs/approximate/uddsketch.rs
+185 -114
View File
@@ -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<f64>,
}
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<Self> {
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::<Self>(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::<Literal>()
.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::<Float64Type>(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<ScalarValue> {
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::<f64>() * 3 + // alpha, gamma, values_sum
std::mem::size_of::<u32>() + // compactions
std::mem::size_of::<u64>() * 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::<SketchHashKey>() + // key
std::mem::size_of::<u64>() + // count
std::mem::size_of::<SketchHashKey>(); // next
total_size += self.uddsketch.current_buckets_count() * bucket_entry_size;
total_size
std::mem::size_of::<Self>() - std::mem::size_of::<UddSketch>()
+ self.uddsketch.allocated_size()
+ self.workspace.allocated_size()
+ self.values.capacity() * std::mem::size_of::<f64>()
}
fn state(&mut self) -> DfResult<Vec<ScalarValue>> {
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"));
}
}
src/common/function/src/function_registry.rs
+7
View File
@@ -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<Arc<FunctionRegistry>> = 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();
src/common/function/src/handlers.rs
+30 -4
View File
@@ -69,16 +69,34 @@ pub trait TableMutationHandler: Send + Sync {
region_id: RegionId,
ctx: QueryContextRef,
) -> Result<AffectedRows>;
/// Discard all unflushed data from a table region.
async fn discard_unflushed_data(
&self,
region_id: RegionId,
ctx: QueryContextRef,
) -> Result<AffectedRows>;
/// Discard all unflushed data from all regions of a table.
async fn discard_unflushed_data_by_table(
&self,
table_name: TableName,
ctx: QueryContextRef,
) -> Result<AffectedRows>;
}
/// A trait for handling procedure service requests in `QueryEngine`.
#[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<Option<String>>;
async fn migrate_region(
&self,
query_ctx: QueryContextRef,
request: MigrateRegionRequest,
) -> Result<Option<String>>;
/// Reconcile a table, database or catalog, returns the procedure id if success.
async fn reconcile(&self, request: ReconcileRequest) -> Result<Option<String>>;
@@ -93,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<MetaGcResponse>;
async fn gc_regions(
&self,
query_ctx: QueryContextRef,
request: MetaGcRegionsRequest,
) -> Result<MetaGcResponse>;
/// Manually trigger GC for a table.
async fn gc_table(&self, request: MetaGcTableRequest) -> Result<MetaGcResponse>;
async fn gc_table(
&self,
query_ctx: QueryContextRef,
request: MetaGcTableRequest,
) -> Result<MetaGcResponse>;
}
/// This flow service handler is only use for flush flow for now.
src/common/function/src/lib.rs
+1
View File
@@ -18,6 +18,7 @@ mod admin;
mod flush_flow;
mod macros;
mod system;
mod uddsketch_compat;
pub mod aggrs;
pub mod function;
src/common/function/src/scalars.rs
+1
View File
@@ -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;
src/common/function/src/scalars/json.rs
+2
View File
@@ -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());
src/common/function/src/scalars/json/json_get.rs
+53 -17
View File
@@ -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<Box<dyn JsonGetResultBuilder>> {
fn result_builder(
len: usize,
with_type: &DataType,
is_json2: bool,
) -> Result<Box<dyn JsonGetResultBuilder>> {
let builder = match with_type {
DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => {
Box::new(StringResultBuilder(StringViewBuilder::with_capacity(len)))
as Box<dyn JsonGetResultBuilder>
}
DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Box::new(StringResultBuilder {
inner: StringViewBuilder::with_capacity(len),
is_json2,
})
as Box<dyn JsonGetResultBuilder>,
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<Box<dyn JsonGetRes
}
// TODO: refactor this to StringLikeArrayBuilder from Arrow 57
struct StringResultBuilder(StringViewBuilder);
struct StringResultBuilder {
inner: StringViewBuilder,
is_json2: bool,
}
impl JsonGetResultBuilder for StringResultBuilder {
fn append_value(&mut self, value: &[u8]) -> 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<Field> {
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()),
src/common/function/src/scalars/json/json_object.rs
+281
View File
@@ -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<Self> {
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::<Int64Type>().value(row).into(),
ValueColumn::UInt(array) => array.as_primitive::<UInt64Type>().value(row).into(),
ValueColumn::Float(array) => array.as_primitive::<Float64Type>().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<DataType> {
Ok(DataType::BinaryView)
}
fn signature(&self) -> &Signature {
&self.signature
}
fn invoke_with_args(
&self,
args: ScalarFunctionArgs,
) -> datafusion_common::Result<ColumnarValue> {
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::<datafusion_common::Result<Vec<_>>>()?;
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<ColumnarValue>, rows: usize) -> datafusion_common::Result<Vec<String>> {
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"));
}
}
src/common/function/src/scalars/uddsketch_calc.rs
+70 -29
View File
@@ -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::<Float64Type>();
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::<Float64Type>();
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::<Float64Type>().null_count(), 3);
}
}
src/common/function/src/scalars/uddsketch_rank.rs
+224
View File
@@ -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<DataType> {
Ok(DataType::Float64)
}
fn signature(&self) -> &Signature {
&self.signature
}
fn invoke_with_args(
&self,
args: ScalarFunctionArgs,
) -> datafusion_common::Result<ColumnarValue> {
let [arg0, arg1] = extract_args(self.name(), &args)?;
let Some(values) = arg0.as_primitive_opt::<Float64Type>() 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::<i32>() 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::<Float64Type>().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);
}
}
src/common/function/src/state.rs
+29 -3
View File
@@ -48,6 +48,7 @@ impl FunctionState {
BuildIndexTableRequest, CompactTableRequest, DeleteRequest, FlushTableRequest,
InsertRequest,
};
use table::table_name::TableName;
use crate::handlers::{FlowServiceHandler, ProcedureServiceHandler, TableMutationHandler};
struct MockProcedureServiceHandler;
@@ -59,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<Option<String>> {
Ok(Some("test_pid".to_string()))
@@ -91,7 +93,11 @@ impl FunctionState {
Ok(())
}
async fn gc_regions(&self, _request: GcRegionsRequest) -> Result<GcResponse> {
async fn gc_regions(
&self,
_context: QueryContextRef,
_request: GcRegionsRequest,
) -> Result<GcResponse> {
Ok(GcResponse {
processed_regions: 1,
need_retry_regions: vec![],
@@ -100,7 +106,11 @@ impl FunctionState {
})
}
async fn gc_table(&self, _request: GcTableRequest) -> Result<GcResponse> {
async fn gc_table(
&self,
_context: QueryContextRef,
_request: GcTableRequest,
) -> Result<GcResponse> {
Ok(GcResponse {
processed_regions: 1,
need_retry_regions: vec![],
@@ -171,6 +181,22 @@ impl FunctionState {
) -> Result<AffectedRows> {
Ok(ROWS)
}
async fn discard_unflushed_data(
&self,
_region_id: RegionId,
_ctx: QueryContextRef,
) -> Result<AffectedRows> {
Ok(ROWS)
}
async fn discard_unflushed_data_by_table(
&self,
_table_name: TableName,
_ctx: QueryContextRef,
) -> Result<AffectedRows> {
Ok(ROWS)
}
}
#[async_trait]
src/common/function/src/uddsketch_compat.rs
+913
View File
@@ -0,0 +1,913 @@
// 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<LegacyBucketKey, LegacyBucket>,
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,
}
struct ValidatedLegacyUddSketch {
buckets: Vec<(LegacyBucketKey, u64)>,
alpha: f64,
gamma: f64,
count: u64,
}
pub(crate) fn decode(raw: &[u8]) -> Result<UddSketch, String> {
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<Option<f64>, 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}"
)
}),
}
}
pub(crate) fn rank(raw: &[u8], value: f64) -> Result<Option<f64>, 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());
}
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<UddSketch, String> {
validate_legacy_input(raw)?;
let state = legacy_options()
.deserialize::<LegacyUddSketchState>(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<LegacyUddSketch, String> {
validate_legacy_input(raw)?;
legacy_options()
.deserialize::<LegacyUddSketchState>(raw)
.map(|state| state.uddsketch)
.or_else(|state_error| {
legacy_options()
.deserialize::<LegacyUddSketch>(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<Vec<u8>, 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<Option<f64>, String> {
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<Option<f64>, 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<ValidatedLegacyUddSketch, 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 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 != self.count {
return Err("legacy bucket counts do not match the value count".to_string());
}
Ok(ValidatedLegacyUddSketch {
buckets,
alpha: self.alpha,
gamma: self.gamma,
count: self.count,
})
}
}
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<f64, String> {
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<Vec<(LegacyBucketKey, u64)>, 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<Vec<u8>, String> {
let negative = buckets
.iter()
.filter_map(|(key, count)| match key {
LegacyBucketKey::Negative(index) => Some((*index, *count)),
_ => None,
})
.collect::<Vec<_>>();
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::<Vec<_>>();
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<u8>,
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<u8>, 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::*;
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::<f64>()];
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::<f64>();
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();
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::<f64>()];
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::<f64>();
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));
}
}

Some files were not shown because too many files have changed in this diff Show More