diff --git a/.github/runner-scale-sets/query-regression/README.md b/.github/runner-scale-sets/query-regression/README.md index 3e0159af60b..fa8d436d2c2 100644 --- a/.github/runner-scale-sets/query-regression/README.md +++ b/.github/runner-scale-sets/query-regression/README.md @@ -51,10 +51,10 @@ Configuration lives in repository variables/secrets: | vars | `QUERY_REGRESSION_ECS_IMAGE_ID` | Custom image built by `ecs-image/build-ecs-image.py`. | | vars | `QUERY_REGRESSION_COMMENT_ALLOWLIST` | Comma/whitespace-separated GitHub logins allowed to comment `/query-regression` on a PR. Each login must also have repository `admin` permission. Empty denies all comment commands. | -The system disk is 50 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 +The system disk is 80 GiB, providing capacity headroom for 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` @@ -108,7 +108,7 @@ agent, which Aliyun public Ubuntu images include. An allowlisted repository admin commenting `/query-regression` on the PR is **trust admission for that exact PR revision**. `/query-regression` runs the -six routine default cases; `/query-regression heavy` runs only the +nine routine default cases; `/query-regression heavy` runs only the high-cardinality `prom_remote_write_7913` remote-write case. `slash-command-dispatch.yml` (`issue_comment` on the default branch) parses the command and `repository_dispatch`es; `query-regression-slash.yml` admits @@ -180,7 +180,7 @@ reintroduce those actions unless the image contract changes. 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 50 GiB system disk and is discarded with the VM. +metadata) lives on the 80 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 @@ -189,7 +189,7 @@ at 10G. Base and candidate builds share the target on that disk; Cargo fingerprints invalidate source and dependency changes. The workflow reports `du` / `df` in telemetry. It does not try to reclaim -space across runs: a cold 50 GiB disk that fills up fails the build. +space across runs: a cold 80 GiB disk that fills up fails the build. ## Future optional phases diff --git a/.github/scripts/aliyun-ecs-runner-provision.py b/.github/scripts/aliyun-ecs-runner-provision.py index 307ee336ebf..b5dee7f51fe 100644 --- a/.github/scripts/aliyun-ecs-runner-provision.py +++ b/.github/scripts/aliyun-ecs-runner-provision.py @@ -81,8 +81,8 @@ POLL_INTERVAL_SECONDS = 5 SWAP_FILE = "/swapfile" SWAP_SIZE_GIB = 16 # Image (~12G), 16G swap, checkout, base+candidate data homes, and cold -# build caches. Deleted with the instance. -SYSTEM_DISK_GIB = 50 +# build caches. 80 GiB provides capacity headroom. Deleted with the instance. +SYSTEM_DISK_GIB = 80 @dataclass(frozen=True) diff --git a/.github/scripts/query-regression-run.py b/.github/scripts/query-regression-run.py index 268cb0bae6c..6da7c896d67 100644 --- a/.github/scripts/query-regression-run.py +++ b/.github/scripts/query-regression-run.py @@ -33,12 +33,14 @@ from typing import Any DEFAULT_CASES = [ "tests/perf/query_cases/smoke_direct_sst/case.toml", + "tests/perf/query_cases/sst_float_bss/case.toml", "tests/perf/query_cases/prom_remote_write_seeded_random/case.toml", "tests/perf/query_cases/prom_remote_write_run_heavy/case.toml", "tests/perf/query_cases/prom_remote_write_mixed_every/case.toml", "tests/perf/query_cases/prom_remote_write_integer_counter/case.toml", "tests/perf/query_cases/promql_range_boundary/case.toml", "tests/perf/query_cases/promql_instant_last_row_9034/case.toml", + "tests/perf/query_cases/mito_prefilter_all_match/case.toml", ] HEAVY_CASES = [ diff --git a/.github/workflows/backport.yml b/.github/workflows/backport.yml index 8cf2da32fcb..013aa320b20 100644 --- a/.github/workflows/backport.yml +++ b/.github/workflows/backport.yml @@ -45,6 +45,7 @@ jobs: PR_TITLE: ${{ github.event.pull_request.title }} PR_BODY: ${{ github.event.pull_request.body }} AUTHOR: ${{ github.event.pull_request.merged_by.login }} + PR_AUTHOR: ${{ github.event.pull_request.user.login }} BASE_BRANCH: ${{ github.event.pull_request.base.ref }} run: | git config user.name "github-actions[bot]" @@ -76,6 +77,7 @@ jobs: git branch -D "${head_branch}" gh issue create \ --title "Backport failure: PR #${PR_NUMBER} to ${release_branch}" \ + --assignee "${PR_AUTHOR}" \ --body "Cherry-pick of ${MERGE_COMMIT} (PR #${PR_NUMBER}) failed on ${release_branch}. Manual backport required." \ || true continue diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 6d12566b684..c9bba9c7726 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -122,6 +122,35 @@ jobs: - name: Run cargo check for riscv64 run: cargo check --locked --workspace --all-targets --features servers/dashboard --target riscv64gc-unknown-linux-gnu + check-windows: + if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} + name: Check Windows build + runs-on: windows-2022 + timeout-minutes: 60 + steps: + - run: git config --global core.autocrlf false + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: ./.github/actions/setup-cyborg + - uses: arduino/setup-protoc@v3 + with: + repo-token: ${{ secrets.GITHUB_TOKEN }} + - uses: KyleMayes/install-llvm-action@v1 + with: + version: "14.0" + - uses: actions-rust-lang/setup-rust-toolchain@v1 + - name: Rust Cache + uses: Swatinem/rust-cache@v2 + with: + shared-key: "check-windows" + cache-all-crates: "true" + save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Run cargo check for Windows + run: cargo check --locked --workspace --all-targets + env: + RUSTFLAGS: "-D warnings" + conflict-check: if: ${{ github.repository == 'GreptimeTeam/greptimedb' }} name: Check for conflict diff --git a/AGENTS.md b/AGENTS.md index 819679f3048..e8feb4abe4b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -69,6 +69,11 @@ a module map, read/write paths, change-coupling points, and gotchas: - [`tests/compatibility/AGENTS.md`](tests/compatibility/AGENTS.md) - [`tests/perf/AGENTS.md`](tests/perf/AGENTS.md) +## Rust imports + +Prefer crate-rooted imports (`use crate::...`) over `use super::...` in +production code. Relative imports using `super` are allowed in tests. + ## Read before changing code - [`.agents/architecture-invariants.md`](.agents/architecture-invariants.md) — diff --git a/Cargo.lock b/Cargo.lock index 00e3d854421..74617c08cbb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1172,7 +1172,7 @@ dependencies = [ "cfg-if", "libc", "miniz_oxide", - "object", + "object 0.36.7", "rustc-demangle", "windows-targets 0.52.6", ] @@ -1623,7 +1623,7 @@ dependencies = [ "maybe-owned", "rustix 1.0.7", "rustix-linux-procfs", - "windows-sys 0.60.2", + "windows-sys 0.61.2", "winx", ] @@ -2271,7 +2271,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" dependencies = [ "lazy_static", - "windows-sys 0.59.0", + "windows-sys 0.48.0", ] [[package]] @@ -3768,7 +3768,7 @@ checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" [[package]] name = "datafusion" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "arrow-schema 58.3.0", @@ -3822,7 +3822,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3846,7 +3846,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3868,7 +3868,7 @@ dependencies = [ [[package]] name = "datafusion-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -3892,7 +3892,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "futures", "log", @@ -3902,7 +3902,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "async-compression", @@ -3936,7 +3936,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "arrow-ipc 58.3.0", @@ -3959,7 +3959,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "async-trait", @@ -3981,7 +3981,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "async-trait", @@ -4004,7 +4004,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "async-trait", @@ -4033,12 +4033,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" [[package]] name = "datafusion-execution" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "arrow-buffer 58.3.0", @@ -4060,7 +4060,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "async-trait", @@ -4082,7 +4082,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4094,7 +4094,7 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "arrow-buffer 58.3.0", @@ -4125,7 +4125,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4146,7 +4146,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4158,7 +4158,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "arrow-ord 58.3.0", @@ -4182,7 +4182,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "async-trait", @@ -4197,7 +4197,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4214,7 +4214,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -4223,7 +4223,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "datafusion-doc", "quote", @@ -4233,7 +4233,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "chrono", @@ -4282,7 +4282,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4305,7 +4305,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4319,7 +4319,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4335,7 +4335,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4353,7 +4353,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "ahash 0.8.12", "arrow 58.3.0", @@ -4384,7 +4384,7 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "chrono", @@ -4411,7 +4411,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4421,7 +4421,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "datafusion-common", @@ -4437,7 +4437,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "async-trait", "datafusion-common", @@ -4450,7 +4450,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "arrow 58.3.0", "bigdecimal 0.4.8", @@ -4468,7 +4468,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" version = "53.1.0" -source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a#a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" +source = "git+https://github.com/GreptimeTeam/datafusion.git?rev=bb531e754b3b36d7a4d3222f35ebc88204dafd2b#bb531e754b3b36d7a4d3222f35ebc88204dafd2b" dependencies = [ "async-recursion", "async-trait", @@ -5177,7 +5177,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.52.0", ] [[package]] @@ -5265,6 +5265,12 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4443176a9f2c162692bd3d352d745ef9413eec5782a80d8fd6f8a1ac692a07f7" +[[package]] +name = "fallible-iterator" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" + [[package]] name = "fallible-streaming-iterator" version = "0.1.9" @@ -5601,6 +5607,20 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "28dd6caf6059519a65843af8fe2a3ae298b14b80179855aeb4adc2c1934ee619" +[[package]] +name = "framehop" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09a5a3f0acb82df800ca3aa50c0d60d286c5d13d4cfc3114b3a9663f13b032fe" +dependencies = [ + "arrayvec", + "cfg-if", + "fallible-iterator 0.3.0", + "gimli", + "macho-unwind-info", + "pe-unwind-info", +] + [[package]] name = "frontend" version = "1.3.0-alpha.1" @@ -5746,7 +5766,7 @@ checksum = "94e7099f6313ecacbe1256e8ff9d617b75d1bcb16a6fddef94866d225a01a14a" dependencies = [ "io-lifetimes 2.0.4", "rustix 1.0.7", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -6062,6 +6082,10 @@ name = "gimli" version = "0.31.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f" +dependencies = [ + "fallible-iterator 0.3.0", + "stable_deref_trait", +] [[package]] name = "git2" @@ -6097,7 +6121,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=32f467fa2ba2b3588a58381a24af83de09fbb00a#32f467fa2ba2b3588a58381a24af83de09fbb00a" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=0c4beffaa2b3968034915867593a717c6ab0c6b4#0c4beffaa2b3968034915867593a717c6ab0c6b4" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", @@ -7143,7 +7167,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20fd6de4ccfcc187e38bc21cfa543cb5a302cb86a8b114eb7f0bf0dc9f8ac00f" dependencies = [ "io-lifetimes 3.0.1", - "windows-sys 0.60.2", + "windows-sys 0.52.0", ] [[package]] @@ -7181,7 +7205,7 @@ checksum = "e04d7f318608d35d4b61ddd75cbdaee86b023ebe2bd5a66ee0915f0bf93095a9" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -7884,7 +7908,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667" dependencies = [ "cfg-if", - "windows-targets 0.53.5", + "windows-targets 0.48.5", ] [[package]] @@ -8230,6 +8254,17 @@ dependencies = [ "thiserror 1.0.69", ] +[[package]] +name = "macho-unwind-info" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb4bdc8b0ce69932332cf76d24af69c3a155242af95c226b2ab6c2e371ed1149" +dependencies = [ + "thiserror 2.0.17", + "zerocopy", + "zerocopy-derive", +] + [[package]] name = "mappings" version = "0.7.1" @@ -8346,6 +8381,15 @@ dependencies = [ "thiserror 1.0.69", ] +[[package]] +name = "memmap2" +version = "0.5.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83faa42c0a078c393f6b29d5db232d8be22776a891f8f56e5284faee4a20b327" +dependencies = [ + "libc", +] + [[package]] name = "memmap2" version = "0.9.5" @@ -9175,7 +9219,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -9443,6 +9487,16 @@ dependencies = [ "objc2-core-foundation", ] +[[package]] +name = "object" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "21158b2c33aa6d4561f1c0a6ea283ca92bc54802a93b263e910746d679a7eb53" +dependencies = [ + "flate2", + "memchr", +] + [[package]] name = "object" version = "0.36.7" @@ -10003,6 +10057,7 @@ dependencies = [ "arrow-ipc 58.3.0", "async-stream", "async-trait", + "axum 0.8.4", "bytes", "catalog", "chrono", @@ -10063,6 +10118,7 @@ dependencies = [ "tokio-util", "tonic 0.14.2", "tracing", + "url", ] [[package]] @@ -10484,6 +10540,19 @@ dependencies = [ "hmac 0.12.1", ] +[[package]] +name = "pe-unwind-info" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "500fa4cdeacd98997c5865e3d0d1cb8fe7e9d7d75ecc775e07989a433a9a9a59" +dependencies = [ + "arrayvec", + "bitflags 2.12.1", + "thiserror 2.0.17", + "zerocopy", + "zerocopy-derive", +] + [[package]] name = "peeking_take_while" version = "1.0.0" @@ -10997,7 +11066,7 @@ dependencies = [ "base64 0.22.1", "byteorder", "bytes", - "fallible-iterator", + "fallible-iterator 0.2.0", "hmac 0.13.0", "md-5 0.11.0", "memchr", @@ -11015,7 +11084,7 @@ dependencies = [ "array-init", "bytes", "chrono", - "fallible-iterator", + "fallible-iterator 0.2.0", "postgres-protocol", "serde_core", "serde_json", @@ -11038,30 +11107,33 @@ checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" [[package]] name = "pprof" -version = "0.14.0" +version = "0.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ebbe2f8898beba44815fdc9e5a4ae9c929e21c5dc29b0c774a15555f7f58d6d0" +checksum = "38a01da47675efa7673b032bf8efd8214f1917d89685e07e395ab125ea42b187" dependencies = [ "aligned-vec", "backtrace", "cfg-if", "criterion 0.5.1", "findshlibs", + "framehop", "inferno 0.11.21", "libc", "log", + "memmap2 0.5.10", "nix 0.26.4", + "object 0.29.0", "once_cell", - "parking_lot 0.12.4", "prost 0.12.6", "prost-build 0.12.6", "prost-derive 0.12.6", - "protobuf 2.28.0", + "protobuf 3.7.2", "sha2 0.10.9", "smallvec", + "spin 0.10.1", "symbolic-demangle", "tempfile", - "thiserror 1.0.69", + "thiserror 2.0.17", ] [[package]] @@ -11383,7 +11455,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4" dependencies = [ "bytes", - "heck 0.5.0", + "heck 0.4.1", "itertools 0.12.1", "log", "multimap", @@ -11403,7 +11475,7 @@ version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ - "heck 0.5.0", + "heck 0.4.1", "itertools 0.14.0", "log", "multimap", @@ -11922,7 +11994,7 @@ dependencies = [ "once_cell", "socket2 0.5.10", "tracing", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -12893,7 +12965,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.4.15", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -12906,7 +12978,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.9.4", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -12997,7 +13069,7 @@ dependencies = [ "security-framework 3.7.0", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.60.2", + "windows-sys 0.52.0", ] [[package]] @@ -13584,6 +13656,7 @@ dependencies = [ "table", "tempfile", "tikv-jemalloc-ctl", + "tikv-jemallocator", "tokio", "tokio-postgres", "tokio-postgres-rustls", @@ -13882,7 +13955,7 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1961e2ef424c1424204d3a5d6975f934f56b6d50ff5732382d84ebf460e147f7" dependencies = [ - "heck 0.5.0", + "heck 0.4.1", "proc-macro2", "quote", "syn 2.0.117", @@ -13911,7 +13984,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -13952,6 +14025,9 @@ name = "spin" version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "023a211cb3138dbc438680b32560ad89f699977624c9f8dbb95a47d5b4c07dd3" +dependencies = [ + "lock_api", +] [[package]] name = "spki" @@ -14303,7 +14379,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -14582,7 +14658,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6a1150bdda9314f6cfeeea801c23f5593c6e6a6c72e64f67e48d723a12b8efdb" dependencies = [ "debugid", - "memmap2", + "memmap2 0.9.5", "stable_deref_trait", "uuid", ] @@ -14751,7 +14827,7 @@ dependencies = [ "lru 0.16.4", "lz4_flex 0.13.1", "measure_time", - "memmap2", + "memmap2 0.9.5", "once_cell", "oneshot", "rayon", @@ -14920,7 +14996,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix 1.0.7", - "windows-sys 0.60.2", + "windows-sys 0.52.0", ] [[package]] @@ -15179,9 +15255,9 @@ dependencies = [ [[package]] name = "tikv-jemalloc-ctl" -version = "0.6.0" +version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f21f216790c8df74ce3ab25b534e0718da5a1916719771d3fec23315c99e468b" +checksum = "3a184c43b8ab2f41df2733b55556e3f5f632f4aeaa205b1bb018f574b7f5f142" dependencies = [ "libc", "paste", @@ -15190,9 +15266,8 @@ dependencies = [ [[package]] name = "tikv-jemalloc-sys" -version = "0.6.0+5.3.0-1-ge13ca993e8ccb9ba9847cc330696e02839f328f7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd3c60906412afa9c2b5b5a48ca6a5abe5736aec9eb48ad05037a677e52e4e2d" +version = "0.7.0+5.3.1-161-gff80bf2d0c8f162ea445533354cad8b7e5830f5c" +source = "git+https://github.com/GreptimeTeam/jemallocator?rev=e254a7eaf13a9ede04953ddefd173b5303018530#e254a7eaf13a9ede04953ddefd173b5303018530" dependencies = [ "cc", "libc", @@ -15200,9 +15275,9 @@ dependencies = [ [[package]] name = "tikv-jemallocator" -version = "0.6.0" +version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4cec5ff18518d81584f477e9bfdf957f5bb0979b0bac3af4ca30b5b3ae2d2865" +checksum = "249f09e49ab1609436f34c776e84231bead18d6a955f119f939bdc1d847561bd" dependencies = [ "libc", "tikv-jemalloc-sys", @@ -15379,7 +15454,7 @@ dependencies = [ "async-trait", "byteorder", "bytes", - "fallible-iterator", + "fallible-iterator 0.2.0", "futures-channel", "futures-util", "log", @@ -16728,7 +16803,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.48.0", ] [[package]] @@ -17178,7 +17253,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f3fd376f71958b862e7afb20cfe5a22830e1963462f3a17f49d82a6c1d1f42d" dependencies = [ "bitflags 2.12.1", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 55181077f8d..b39064706c4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -159,7 +159,7 @@ fs2 = "0.4" fst = "0.4.7" futures = "0.3" futures-util = "0.3" -greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "32f467fa2ba2b3588a58381a24af83de09fbb00a" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "0c4beffaa2b3968034915867593a717c6ab0c6b4" } hex = "0.4" hostname = "0.4.0" http = "1" @@ -353,23 +353,30 @@ git = "https://github.com/GreptimeTeam/greptime-meter.git" rev = "5618e779cf2bb4755b499c630fba4c35e91898cb" [patch.crates-io] -datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" } -datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" } -datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" } -datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" } -datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" } -datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" } -datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" } -datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" } -datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" } -datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" } -datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" } -datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" } -datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" } -datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" } -datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "a281d9f0ea4b3fb2ec88bc3553d8dcef307e072a" } +datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" } +datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" } +datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" } +datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" } +datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" } +datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" } +datafusion-functions-aggregate-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" } +datafusion-functions-window-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" } +datafusion-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" } +datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" } +datafusion-physical-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" } +datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" } +datafusion-proto = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" } +datafusion-sql = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" } +datafusion-substrait = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "bb531e754b3b36d7a4d3222f35ebc88204dafd2b" } sqlparser = { git = "https://github.com/GreptimeTeam/sqlparser-rs.git", rev = "2aefa08a8d69c96eec2d6d6703598a009bba6e4c" } # on branch v0.61.x +# Temporary: use the GreptimeTeam fork of tikv-jemalloc-sys embedding +# jemalloc 5.3.1 + backport of 54f22c83 ("Initialize TSD tcache before +# enabling it"), fixing the TSD/tcache enable-before-init reentrancy +# window. Remove once tikv/jemallocator ships a jemalloc snapshot that +# includes the fix. See GreptimeTeam/jemalloc#1. +tikv-jemalloc-sys = { git = "https://github.com/GreptimeTeam/jemallocator", rev = "e254a7eaf13a9ede04953ddefd173b5303018530" } + [profile.release] debug = 1 diff --git a/Makefile b/Makefile index 0f280b7b9be..9bed2ddc9b8 100644 --- a/Makefile +++ b/Makefile @@ -4,6 +4,8 @@ FEATURES ?= TARGET_DIR ?= TARGET ?= BUILD_BIN ?= greptime +BUILD_PACKAGE ?= $(if $(filter greptime,$(strip $(BUILD_BIN))),cmd) +BUILD_PACKAGE_OPT = $(if $(strip $(BUILD_PACKAGE)),-p ${BUILD_PACKAGE}) CARGO_BUILD_OPTS := --locked IMAGE_REGISTRY ?= docker.io IMAGE_NAMESPACE ?= greptime @@ -84,7 +86,7 @@ endif .PHONY: build build: ## Build debug version greptime. - cargo ${CARGO_EXTENSION} build ${CARGO_BUILD_OPTS} + cargo ${CARGO_EXTENSION} build ${BUILD_PACKAGE_OPT} ${CARGO_BUILD_OPTS} .PHONY: build-by-dev-builder build-by-dev-builder: ## Build greptime by dev-builder. @@ -93,6 +95,8 @@ build-by-dev-builder: ## Build greptime by dev-builder. -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-${BASE_IMAGE}:${DEV_BUILDER_IMAGE_TAG} \ make build \ + BUILD_PACKAGE="${BUILD_PACKAGE}" \ + BUILD_BIN="${BUILD_BIN}" \ CARGO_EXTENSION="${CARGO_EXTENSION}" \ CARGO_PROFILE=${CARGO_PROFILE} \ FEATURES=${FEATURES} \ @@ -108,6 +112,8 @@ build-riscv64-bin: ## Build greptime binary for riscv64 (linux-gnu) by the riscv -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 \ + BUILD_PACKAGE="${BUILD_PACKAGE}" \ + BUILD_BIN="${BUILD_BIN}" \ CARGO_EXTENSION="${CARGO_EXTENSION}" \ CARGO_PROFILE=${CARGO_PROFILE} \ FEATURES=${FEATURES} \ @@ -122,6 +128,8 @@ build-android-bin: ## Build greptime binary for android. -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-android:${DEV_BUILDER_IMAGE_TAG} \ make build \ + BUILD_PACKAGE="${BUILD_PACKAGE}" \ + BUILD_BIN="${BUILD_BIN}" \ CARGO_EXTENSION="ndk --platform 23 -t aarch64-linux-android" \ CARGO_PROFILE=release \ FEATURES="${FEATURES}" \ @@ -278,6 +286,7 @@ dashboards: ## Generate the Grafana dashboards for standalone mode and intermedi @./grafana/scripts/gen-dashboards.sh ##@ Docs +.PHONY: config-docs config-docs: ## Generate configuration documentation from toml files. docker run --rm \ -v ${PWD}:/greptimedb \ diff --git a/README.md b/README.md index e95f7849e4b..1c8ed985e1d 100644 --- a/README.md +++ b/README.md @@ -13,95 +13,65 @@ One engine, on your infrastructure. > Apache-2.0 licensed core.
-

- User Guide | - API Docs | - Roadmap 2026 -

- -Stable - - -Canary - - -Nightly - +User Guide  ·  +API Docs  ·  +Roadmap 2026  ·  +Slack + +Stable +Canary +Nightly +Docker Pulls +License stable for production  ·  canary includes pre-releases  ·  nightly is a weekly snapshot of main - -Docker Pulls - - -GitHub Actions - - -Codecov - - -License - - -
- - -Slack - - -Twitter - - -LinkedIn -
-- [Introduction](#introduction) -- [Why You Might Use It](#why-you-might-use-it) -- [Overview](#overview) -- [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) -- [Build From Source](#build-from-source) -- [Tools & Extensions](#tools--extensions) -- [Project Status](#project-status) -- [Community](#community) -- [License](#license) -- [Commercial Support](#commercial-support) -- [Contributing](#contributing) -- [Acknowledgement](#acknowledgement) - ## Introduction **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. 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. +## One Query Across Signals + +OpenTelemetry ingestion writes spans to `opentelemetry_traces` and log records to +`opentelemetry_logs`. Both tables carry `trace_id`, so correlating them is a join: + +```sql +-- The slowest failed spans in the last hour, +-- with the log lines emitted inside those same traces. +SELECT + t.service_name, + t.span_name, + t.duration_nano / 1000000 AS duration_ms, + l.timestamp AS log_time, + l.severity_text, + l.body +FROM opentelemetry_traces t +JOIN opentelemetry_logs l ON l.trace_id = t.trace_id +WHERE t.timestamp > now() - INTERVAL '1' HOUR + AND t.span_status_code = 'STATUS_CODE_ERROR' +ORDER BY t.duration_nano DESC +LIMIT 20; +``` + +Metrics join the same way, on any tag the tables share, such as `service`, +`host`, or `pod`. + ## 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 are hitting Loki's query performance limits as log volume grows - 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 - -A quick overview of what GreptimeDB ingests, how it connects to other systems, and what its distributed engine lets you do. - -

- - GreptimeDB Overview - -

- ## What's Supported | | | @@ -113,6 +83,22 @@ A quick overview of what GreptimeDB ingests, how it connects to other systems, a 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. +

+ + GreptimeDB Overview + +

+ +## Benchmarks + +* [Agent RCA Bench](https://rca-bench.greptime.com/): LLM agents doing root cause analysis + over GreptimeDB versus Prometheus + Loki + Tempo. 40% fewer wrong diagnoses, 48% fewer + input tokens, 45% lower cost + ([write-up](https://greptime.com/blogs/2026-09-08-agent-rca-bench-interface-semantic-layer)) +* [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) + ## Compatibility and Migration Compatibility is per protocol, and query-side coverage is narrower than ingestion. @@ -123,11 +109,6 @@ Compatibility is per protocol, and query-side coverage is narrower than ingestio | **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. @@ -221,11 +202,7 @@ See the [Contribution Guidelines](CONTRIBUTING.md) for the full developer workfl 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). - -If GreptimeDB is useful to you, please star the repo. - -Known Users +Release lines and support windows are in the [version reference](https://docs.greptime.com/nightly/reference/about-greptimedb-version). For where the project is going, read the [v1.0 highlights](https://greptime.com/blogs/2025-11-05-greptimedb-v1-highlights) and the [2026 roadmap](https://greptime.com/blogs/2026-02-11-greptimedb-roadmap-2026). ## Community @@ -239,6 +216,10 @@ We invite you to engage and contribute! - [X (Twitter)](https://X.com/greptime) - [YouTube](https://www.youtube.com/@greptime) +If GreptimeDB is useful to you, please star the repo. + +Known Users + ## License GreptimeDB is an open-core project. Its core is licensed under the @@ -262,6 +243,9 @@ security, and support layer for production deployments. - Explore [Internal Concepts](https://docs.greptime.com/contributor-guide/overview.html) and [DeepWiki](https://deepwiki.com/GreptimeTeam/greptimedb). - Pick up a [good first issue](https://github.com/GreptimeTeam/greptimedb/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22) and join the #contributors [Slack](https://greptime.com/slack) channel. +[![Integration CI](https://github.com/GreptimeTeam/greptimedb/actions/workflows/integration.yml/badge.svg)](https://github.com/GreptimeTeam/greptimedb/actions/workflows/integration.yml) +[![Codecov](https://codecov.io/gh/GreptimeTeam/greptimedb/branch/main/graph/badge.svg?token=FITFDI3J3C)](https://codecov.io/gh/GreptimeTeam/greptimedb) + ## Acknowledgement Special thanks to all contributors! See [AUTHOR.md](AUTHOR.md). diff --git a/config/config.md b/config/config.md index a8d12c3557e..2d85ddd7176 100644 --- a/config/config.md +++ b/config/config.md @@ -249,7 +249,7 @@ | `pipeline.cache_ttl` | String | `10s` | Time to live of the local pipeline cache. Default is `10s`. | | `event_recorder` | -- | -- | Configuration options for the event recorder. | | `event_recorder.ttl` | String | `90d` | TTL for the events table that will be used to store the events. Default is `90d`. | -| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `create_database`,
`alter_database`, `drop_database`, `create_flow`, `drop_flow`,
`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,
`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,
`create_view`, `drop_view`, `admin_function`, `reconcile_table`,
`reconcile_logical_tables`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | +| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `create_database`,
`alter_database`, `drop_database`, `create_flow`, `drop_flow`,
`create_table`, `create_logical_tables`, `alter_table`, `alter_logical_tables`,
`drop_table`, `undrop_table`, `purge_dropped_table`, `truncate_table`,
`create_view`, `drop_view`, `admin_function`, `reconcile_catalog`,
`reconcile_database`, `reconcile_table`, `reconcile_logical_tables`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | | `memory` | -- | -- | The memory options. | | `memory.enable_heap_profiling` | Bool | `true` | Whether to enable heap profiling activation during startup.
When enabled, heap profiling will be activated if the `MALLOC_CONF` environment variable
is set to "prof:true,prof_active:false". The official image adds this env variable.
Default is true. | @@ -477,7 +477,7 @@ | `wal.create_topic_timeout` | String | `30s` | The timeout for creating a Kafka topic.
**It's only used when the provider is `kafka`**. | | `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: `region_migration`,
`create_database`, `alter_database`, `drop_database`, `create_flow`,
`drop_flow`, `create_table`, `create_logical_tables`, `alter_table`,
`alter_logical_tables`, `drop_table`, `undrop_table`, `purge_dropped_table`,
`truncate_table`, `create_view`, `drop_view`, `repartition`,
`repartition_group`, `wal_prune`, `batch_gc`, `reconcile_table`,
`reconcile_logical_tables`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | +| `event_recorder.event_types` | Array | -- | Event types to record. Current available event types: `region_migration`,
`create_database`, `alter_database`, `drop_database`, `create_flow`,
`drop_flow`, `create_table`, `create_logical_tables`, `alter_table`,
`alter_logical_tables`, `drop_table`, `undrop_table`, `purge_dropped_table`,
`truncate_table`, `create_view`, `drop_view`, `repartition`,
`repartition_group`, `wal_prune`, `batch_gc`, `reconcile_catalog`,
`reconcile_database`, `reconcile_table`, `reconcile_logical_tables`.
When omitted, all current and future event types are recorded.
Set to an empty array to disable event recording. | | `stats_persistence` | -- | -- | Configuration options for the stats persistence. | | `stats_persistence.ttl` | String | `0s` | TTL for the stats table that will be used to store the stats.
Set to `0s` to disable stats persistence.
Default is `0s`.
If you want to enable stats persistence, set the TTL to a value greater than 0.
It is recommended to set a small value, e.g., `3h`. | | `stats_persistence.interval` | String | `10m` | The interval to persist the stats. Default is `10m`.
The minimum value is `10m`, if the value is less than `10m`, it will be overridden to `10m`. | diff --git a/config/metasrv.example.toml b/config/metasrv.example.toml index 514993a4cae..b923c5a549c 100644 --- a/config/metasrv.example.toml +++ b/config/metasrv.example.toml @@ -328,8 +328,8 @@ ttl = "90d" ## `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`, `repartition`, -## `repartition_group`, `wal_prune`, `batch_gc`, `reconcile_table`, -## `reconcile_logical_tables`. +## `repartition_group`, `wal_prune`, `batch_gc`, `reconcile_catalog`, +## `reconcile_database`, `reconcile_table`, `reconcile_logical_tables`. ## When omitted, all current and future event types are recorded. ## Set to an empty array to disable event recording. #+ event_types = ["region_migration"] diff --git a/config/standalone.example.toml b/config/standalone.example.toml index b05cf0177d0..0bef0793247 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -973,8 +973,8 @@ 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`, `admin_function`, `reconcile_table`, -## `reconcile_logical_tables`. +## `create_view`, `drop_view`, `admin_function`, `reconcile_catalog`, +## `reconcile_database`, `reconcile_table`, `reconcile_logical_tables`. ## When omitted, all current and future event types are recorded. ## Set to an empty array to disable event recording. #+ event_types = ["create_database"] diff --git a/src/catalog/src/system_schema/information_schema/columns.rs b/src/catalog/src/system_schema/information_schema/columns.rs index 4e68b093cd4..e92efd3bebf 100644 --- a/src/catalog/src/system_schema/information_schema/columns.rs +++ b/src/catalog/src/system_schema/information_schema/columns.rs @@ -31,7 +31,7 @@ use datatypes::scalars::ScalarVectorBuilder; use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; use datatypes::value::Value; use datatypes::vectors::{ - ConstantVector, Int64Vector, Int64VectorBuilder, StringVector, StringVectorBuilder, VectorRef, + Int64Vector, Int64VectorBuilder, StringVector, StringVectorBuilder, VectorRef, }; use futures::TryStreamExt; use snafu::{OptionExt, ResultExt}; @@ -408,18 +408,9 @@ impl InformationSchemaColumnsBuilder { fn finish(&mut self) -> Result { let rows_num = self.collation_names.len(); - let privileges = Arc::new(ConstantVector::new( - Arc::new(StringVector::from(vec![DEFAULT_PRIVILEGES])), - rows_num, - )); - let empty_string = Arc::new(ConstantVector::new( - Arc::new(StringVector::from(vec![EMPTY_STR])), - rows_num, - )); - let srs_ids = Arc::new(ConstantVector::new( - Arc::new(Int64Vector::from(vec![None])), - rows_num, - )); + let privileges = Arc::new(StringVector::from(vec![DEFAULT_PRIVILEGES; rows_num])); + let empty_string = Arc::new(StringVector::from(vec![EMPTY_STR; rows_num])); + let srs_ids = Arc::new(Int64Vector::from(vec![None; rows_num])); let columns: Vec = vec![ Arc::new(self.catalog_names.finish()), diff --git a/src/catalog/src/system_schema/information_schema/key_column_usage.rs b/src/catalog/src/system_schema/information_schema/key_column_usage.rs index 99b0d7f6b73..4f8f116afaf 100644 --- a/src/catalog/src/system_schema/information_schema/key_column_usage.rs +++ b/src/catalog/src/system_schema/information_schema/key_column_usage.rs @@ -26,7 +26,7 @@ use datafusion::physical_plan::streaming::PartitionStream as DfPartitionStream; use datatypes::prelude::{ConcreteDataType, MutableVector, ScalarVectorBuilder, VectorRef}; use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; use datatypes::value::Value; -use datatypes::vectors::{ConstantVector, StringVector, StringVectorBuilder, UInt32VectorBuilder}; +use datatypes::vectors::{StringVector, StringVectorBuilder, UInt32VectorBuilder}; use futures_util::TryStreamExt; use snafu::{OptionExt, ResultExt}; use store_api::storage::{ScanRequest, TableId}; @@ -326,10 +326,7 @@ impl InformationSchemaKeyColumnUsageBuilder { fn finish(&mut self) -> Result { let rows_num = self.table_catalog.len(); - let null_string_vector = Arc::new(ConstantVector::new( - Arc::new(StringVector::from(vec![None as Option<&str>])), - rows_num, - )); + let null_string_vector = Arc::new(StringVector::from(vec![None as Option<&str>; rows_num])); let columns: Vec = vec![ Arc::new(self.constraint_catalog.finish()), Arc::new(self.constraint_schema.finish()), diff --git a/src/catalog/src/system_schema/information_schema/partitions.rs b/src/catalog/src/system_schema/information_schema/partitions.rs index b9396fe5543..d2cce93150a 100644 --- a/src/catalog/src/system_schema/information_schema/partitions.rs +++ b/src/catalog/src/system_schema/information_schema/partitions.rs @@ -29,8 +29,8 @@ use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; use datatypes::timestamp::TimestampSecond; use datatypes::value::Value; use datatypes::vectors::{ - ConstantVector, Int64Vector, Int64VectorBuilder, MutableVector, StringVector, - StringVectorBuilder, TimestampSecondVector, TimestampSecondVectorBuilder, UInt64VectorBuilder, + Int64Vector, Int64VectorBuilder, MutableVector, StringVector, StringVectorBuilder, + TimestampSecondVector, TimestampSecondVectorBuilder, UInt64VectorBuilder, }; use futures::{StreamExt, TryStreamExt}; use partition::manager::PartitionInfo; @@ -361,22 +361,11 @@ impl InformationSchemaPartitionsBuilder { fn finish(&mut self) -> Result { let rows_num = self.catalog_names.len(); - let null_string_vector = Arc::new(ConstantVector::new( - Arc::new(StringVector::from(vec![None as Option<&str>])), - rows_num, - )); - let null_i64_vector = Arc::new(ConstantVector::new( - Arc::new(Int64Vector::from(vec![None])), - rows_num, - )); - let null_timestamp_second_vector = Arc::new(ConstantVector::new( - Arc::new(TimestampSecondVector::from(vec![None])), - rows_num, - )); - let partition_methods = Arc::new(ConstantVector::new( - Arc::new(StringVector::from(vec![Some("RANGE")])), - rows_num, - )); + let null_string_vector = Arc::new(StringVector::from(vec![None as Option<&str>; rows_num])); + let null_i64_vector = Arc::new(Int64Vector::from(vec![None; rows_num])); + let null_timestamp_second_vector = + Arc::new(TimestampSecondVector::from(vec![None; rows_num])); + let partition_methods = Arc::new(StringVector::from(vec![Some("RANGE"); rows_num])); let columns: Vec = vec![ Arc::new(self.catalog_names.finish()), diff --git a/src/catalog/src/system_schema/information_schema/table_constraints.rs b/src/catalog/src/system_schema/information_schema/table_constraints.rs index 8c975207f21..2208d303c2e 100644 --- a/src/catalog/src/system_schema/information_schema/table_constraints.rs +++ b/src/catalog/src/system_schema/information_schema/table_constraints.rs @@ -27,7 +27,7 @@ use datatypes::prelude::{ConcreteDataType, MutableVector}; use datatypes::scalars::ScalarVectorBuilder; use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; use datatypes::value::Value; -use datatypes::vectors::{ConstantVector, StringVector, StringVectorBuilder, VectorRef}; +use datatypes::vectors::{StringVector, StringVectorBuilder, VectorRef}; use futures::TryStreamExt; use snafu::{OptionExt, ResultExt}; use store_api::storage::{ScanRequest, TableId}; @@ -242,14 +242,8 @@ impl InformationSchemaTableConstraintsBuilder { fn finish(&mut self) -> Result { let rows_num = self.constraint_names.len(); - let constraint_catalogs = Arc::new(ConstantVector::new( - Arc::new(StringVector::from(vec!["def"])), - rows_num, - )); - let enforceds = Arc::new(ConstantVector::new( - Arc::new(StringVector::from(vec!["YES"])), - rows_num, - )); + let constraint_catalogs = Arc::new(StringVector::from(vec!["def"; rows_num])); + let enforceds = Arc::new(StringVector::from(vec!["YES"; rows_num])); let columns: Vec = vec![ constraint_catalogs, diff --git a/src/client/src/database.rs b/src/client/src/database.rs index 881730cfaaa..e97fff6db00 100644 --- a/src/client/src/database.rs +++ b/src/client/src/database.rs @@ -825,7 +825,17 @@ impl Database { /// Ingest a stream of [RecordBatch]es that belong to a table, using Arrow Flight's "`DoPut`" /// method. The return value is also a stream, produces [DoPutResponse]s. pub async fn do_put(&self, stream: FlightDataStream) -> Result { + self.do_put_with_hints(stream, &[]).await + } + + /// Ingest a stream of [RecordBatch]es using Arrow Flight's `DoPut` with request hints. + pub async fn do_put_with_hints( + &self, + stream: FlightDataStream, + hints: &[(&str, &str)], + ) -> Result { let mut request = tonic::Request::new(stream); + Self::put_hints(request.metadata_mut(), hints)?; if let Some(AuthHeader { auth_scheme: Some(AuthScheme::Basic(Basic { username, password })), diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml index 8cb8cff4300..722e48ae90a 100644 --- a/src/cmd/Cargo.toml +++ b/src/cmd/Cargo.toml @@ -119,12 +119,12 @@ tracing-appender.workspace = true [target.'cfg(unix)'.dependencies] daemonize = "0.5" -pprof = { version = "0.14", features = [ +pprof = { version = "0.15", features = [ "flamegraph", ] } [target.'cfg(not(windows))'.dependencies] -tikv-jemallocator = "0.6" +tikv-jemallocator = "0.7" [dev-dependencies] api.workspace = true diff --git a/src/cmd/src/bin/query_perf_fixture/case.rs b/src/cmd/src/bin/query_perf_fixture/case.rs index a76c49d43d0..0f5239a6b51 100644 --- a/src/cmd/src/bin/query_perf_fixture/case.rs +++ b/src/cmd/src/bin/query_perf_fixture/case.rs @@ -105,6 +105,10 @@ pub(super) struct PromRemoteWritePlan { #[serde(default = "default_visibility_timeout_seconds")] pub(super) visibility_timeout_seconds: u64, #[serde(default)] + pub(super) base_setup_sql: Vec, + #[serde(default)] + pub(super) candidate_setup_sql: Vec, + #[serde(default)] pub(super) prom_store: PromStoreConfig, #[serde(default)] pub(super) value: ValueConfig, @@ -200,6 +204,7 @@ pub(super) enum ValuePattern { QuantizedSignal, SignalWithSporadicStalls, MixedSignalRepeated, + BoundedMixed, } impl std::fmt::Display for ValuePattern { @@ -436,6 +441,63 @@ pub(super) fn default_parallelism() -> u64 { 1 } +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn remote_write_setup_sql_defaults_to_empty() { + let case: CaseFile = toml::from_str( + r#" +[scenario] +kind = "prom_remote_write_then_query" + +[scenario.remote_write] +metric = "metric" +"#, + ) + .unwrap(); + let Scenario::PromRemoteWriteThenQuery(scenario) = case.scenario else { + panic!("expected prom_remote_write_then_query scenario"); + }; + assert!(scenario.remote_write.base_setup_sql.is_empty()); + assert!(scenario.remote_write.candidate_setup_sql.is_empty()); + } + + #[test] + fn remote_write_setup_sql_roundtrips() { + let case: CaseFile = toml::from_str( + r#" +[scenario] +kind = "prom_remote_write_then_query" + +[scenario.remote_write] +metric = "metric" +base_setup_sql = ["CREATE TABLE base_table"] +candidate_setup_sql = ["CREATE TABLE candidate_table", "ALTER TABLE candidate_table SET 'x'='y'"] +"#, + ) + .unwrap(); + + let roundtrip: CaseFile = + serde_json::from_str(&serde_json::to_string(&case).unwrap()).unwrap(); + let Scenario::PromRemoteWriteThenQuery(scenario) = roundtrip.scenario else { + panic!("expected prom_remote_write_then_query scenario"); + }; + assert_eq!( + scenario.remote_write.base_setup_sql, + ["CREATE TABLE base_table"] + ); + assert_eq!( + scenario.remote_write.candidate_setup_sql, + [ + "CREATE TABLE candidate_table", + "ALTER TABLE candidate_table SET 'x'='y'" + ] + ); + } +} + impl Scenario { pub(super) fn kind(&self) -> &'static str { match self { diff --git a/src/cmd/src/bin/query_perf_fixture/direct_sst.rs b/src/cmd/src/bin/query_perf_fixture/direct_sst.rs index 746ecbd76e4..27a6b4109cd 100644 --- a/src/cmd/src/bin/query_perf_fixture/direct_sst.rs +++ b/src/cmd/src/bin/query_perf_fixture/direct_sst.rs @@ -469,6 +469,7 @@ pub(super) async fn run_direct_sst(args: DirectArgs) { write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE, row_group_size: scenario.layout.row_group_size, max_file_size: None, + float_field_encoding: Default::default(), }; let infos = match format { FormatType::Flat => writer.write_all_flat(source, Some(sequence), &opts).await, diff --git a/src/cmd/src/bin/query_perf_fixture/prom_remote_write.rs b/src/cmd/src/bin/query_perf_fixture/prom_remote_write.rs index 1f36b0c2d23..a59178f4a79 100644 --- a/src/cmd/src/bin/query_perf_fixture/prom_remote_write.rs +++ b/src/cmd/src/bin/query_perf_fixture/prom_remote_write.rs @@ -197,8 +197,35 @@ fn deterministic_prom_value(args: &PromRemoteWriteArgs, series_idx: u64, sample_ args.value_base + (series_idx % 97) as f64 + local as f64 * args.value_step } } + ValuePattern::BoundedMixed => bounded_mixed_value(args, series_idx, local), } } + +fn bounded_mixed_value(args: &PromRemoteWriteArgs, series_idx: u64, local: u64) -> f64 { + const SPANS: [i64; 4] = [10, 1_000, 100_000, 10_000_000]; + let series = splitmix64(series_idx ^ args.value_seed); + let span = SPANS[(series >> 62) as usize] * (1 + ((series >> 32) % 10) as i64); + let baseline = (series % args.value_cardinality.max(1)) as f64 * span as f64 / 10.0; + let run_length = args.value_run_length.max(1); + let group = local / run_length; + let phase = local % run_length; + let anchor = bounded_mixed_anchor(series, group, span); + let next_anchor = bounded_mixed_anchor(series, group + 1, span); + let signal = baseline + anchor + (next_anchor - anchor) * phase as f64 / run_length as f64; + let value = (args.value_base + args.value_step * signal).round(); + if series_idx % args.value_mixed_every.max(1) == args.value_mixed_every.max(1) - 1 { + let numerator = splitmix64(series ^ local) % 999 + 1; + value + numerator as f64 / 1_000.0 + } else { + value + } +} + +fn bounded_mixed_anchor(series: u64, group: u64, span: i64) -> f64 { + let limit = span - 1; + (splitmix64(series ^ group) % (2 * limit as u64 + 1)) as i64 as f64 - limit as f64 +} + fn splitmix64(mut value: u64) -> u64 { value = value.wrapping_add(0x9e3779b97f4a7c15); let mut mixed = value; @@ -307,7 +334,7 @@ mod tests { for chunk_idx in 0..10 { let mut chunk = args(ValuePattern::SeededRandom); chunk.samples_per_series = 1_440; - chunk.value_sample_offset = chunk_idx * 1_440; + chunk.value_sample_offset = chunk_idx * 1_439; chunk.value_total_samples_per_series = Some(14_400); for sample_idx in 0..chunk.samples_per_series { assert_eq!( @@ -319,6 +346,118 @@ mod tests { } } + #[test] + fn bounded_mixed_has_bounded_95_5_per_series_values() { + let mut args = args(ValuePattern::BoundedMixed); + args.series_count = 1_000; + args.samples_per_series = 4_320; + args.value_cardinality = 1_000; + args.value_seed = 42; + args.value_run_length = 60; + args.value_mixed_every = 20; + + let mut integral_series = 0; + let mut fractional_series = 0; + let mut spans = HashSet::new(); + let mut ranges = HashSet::new(); + for series_idx in 0..args.series_count { + let series = splitmix64(series_idx ^ args.value_seed); + let span = [10_i64, 1_000, 100_000, 10_000_000][(series >> 62) as usize] + * (1 + ((series >> 32) % 10) as i64); + let baseline = (series % args.value_cardinality.max(1)) as f64 * span as f64 / 10.0; + let values = (0..args.samples_per_series) + .map(|sample_idx| deterministic_prom_value(&args, series_idx, sample_idx)) + .collect::>(); + + assert!(values.iter().all(|value| value.is_finite())); + assert!( + values + .iter() + .all(|value| *value >= baseline - span as f64 - 1.0 + && *value <= baseline + span as f64 + 1.0) + ); + assert!( + values.windows(2).any(|pair| pair[1] > pair[0]), + "series {series_idx} never rises" + ); + assert!( + values.windows(2).any(|pair| pair[1] < pair[0]), + "series {series_idx} never falls" + ); + ranges.insert(( + values + .iter() + .copied() + .fold(f64::INFINITY, f64::min) + .to_bits(), + values + .iter() + .copied() + .fold(f64::NEG_INFINITY, f64::max) + .to_bits(), + )); + spans.insert(span); + if series_idx % args.value_mixed_every == args.value_mixed_every - 1 { + fractional_series += 1; + assert!(values.iter().all(|value| value.fract() != 0.0)); + } else { + integral_series += 1; + assert!(values.iter().all(|value| value.fract() == 0.0)); + } + } + + assert_eq!((integral_series, fractional_series), (950, 50)); + assert!(spans.len() > 4); + assert!(ranges.len() > 1); + } + + #[test] + fn bounded_mixed_is_seeded_and_chunk_bit_identical() { + let mut monolithic = args(ValuePattern::BoundedMixed); + monolithic.samples_per_series = 14_400; + monolithic.value_total_samples_per_series = Some(14_400); + monolithic.value_cardinality = 1_000; + monolithic.value_seed = 42; + monolithic.value_run_length = 60; + monolithic.value_mixed_every = 20; + + let original = (0..600) + .map(|sample_idx| deterministic_prom_value(&monolithic, 19, sample_idx)) + .collect::>(); + monolithic.value_seed = 43; + assert_ne!( + original, + (0..600) + .map(|sample_idx| deterministic_prom_value(&monolithic, 19, sample_idx)) + .collect::>() + ); + + monolithic.value_seed = 42; + for series_idx in [0, 19] { + for chunk_idx in 0..10 { + let mut chunk = args(ValuePattern::BoundedMixed); + chunk.samples_per_series = 1_440; + chunk.value_sample_offset = chunk_idx * 1_439; + chunk.value_total_samples_per_series = Some(14_400); + chunk.value_cardinality = 1_000; + chunk.value_seed = 42; + chunk.value_run_length = 60; + chunk.value_mixed_every = 20; + for sample_idx in 0..chunk.samples_per_series { + assert_eq!( + deterministic_prom_value(&chunk, series_idx, sample_idx).to_bits(), + deterministic_prom_value( + &monolithic, + series_idx, + chunk.value_sample_offset + sample_idx + ) + .to_bits() + ); + } + } + } + } + #[test] fn values_continue_across_series_boundaries() { let mut args = args(ValuePattern::Unique); diff --git a/src/cmd/src/bin/query_regression_runner/model.rs b/src/cmd/src/bin/query_regression_runner/model.rs index 6831a5256a4..5fafdb5635b 100644 --- a/src/cmd/src/bin/query_regression_runner/model.rs +++ b/src/cmd/src/bin/query_regression_runner/model.rs @@ -116,6 +116,10 @@ pub(super) struct RemoteWrite { pub(super) sample_chunk_size: Option, pub(super) flush_every_sample_chunks: u64, pub(super) visibility_timeout_seconds: u64, + #[serde(default)] + pub(super) base_setup_sql: Vec, + #[serde(default)] + pub(super) candidate_setup_sql: Vec, pub(super) prom_store: PromStore, pub(super) value: RemoteValue, pub(super) storage: Option, diff --git a/src/cmd/src/bin/query_regression_runner/remote/ingest.rs b/src/cmd/src/bin/query_regression_runner/remote/ingest.rs index f563402afcd..271e27fac42 100644 --- a/src/cmd/src/bin/query_regression_runner/remote/ingest.rs +++ b/src/cmd/src/bin/query_regression_runner/remote/ingest.rs @@ -57,6 +57,7 @@ pub(super) async fn run_prepare_remote(args: PrepareRemoteArgs) -> Result<()> { args.base_http_port, &args.fixture_generator, &remote, + &remote.base_setup_sql, &client, ) .await?; @@ -65,6 +66,7 @@ pub(super) async fn run_prepare_remote(args: PrepareRemoteArgs) -> Result<()> { args.candidate_http_port, &args.fixture_generator, &remote, + &remote.candidate_setup_sql, &client, ) .await?; @@ -88,6 +90,7 @@ async fn prepare_remote_target( port: u16, generator: &Path, remote: &RemoteWrite, + setup_sql: &[String], client: &Client, ) -> Result { let create_database = http_post_sql( @@ -107,6 +110,7 @@ async fn prepare_remote_target( ) .into()); } + let setup = run_setup_sql(client, port, &remote.database, name, setup_sql).await?; let (remote_write, flushes) = ingest_remote_write(generator, port, remote, client).await?; let expected_rows = remote .series_count @@ -124,6 +128,7 @@ async fn prepare_remote_target( Ok(json!({ "name": name, "create_database": create_database, + "setup_sql": setup, "remote_write": remote_write, "flushes": flushes, "visibility": visibility, @@ -131,6 +136,28 @@ async fn prepare_remote_target( })) } +async fn run_setup_sql( + client: &Client, + port: u16, + database: &str, + target: &str, + setup_sql: &[String], +) -> Result> { + let mut results = Vec::with_capacity(setup_sql.len()); + for (index, statement) in setup_sql.iter().enumerate() { + let result = http_post_sql(client, port, statement, database).await; + if !result["ok"].as_bool().unwrap_or(false) { + return Err(format!( + "setup SQL statement {} failed for {target}: {result}", + index + 1 + ) + .into()); + } + results.push(result); + } + Ok(results) +} + #[derive(Clone, Debug, PartialEq, Eq)] struct SampleChunk { index: u64, @@ -434,9 +461,110 @@ fn json_number(value: f64) -> String { #[cfg(test)] mod tests { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::{TcpListener, TcpStream}; + use super::*; use crate::query_regression_runner::model::{PromStore, RemoteValue}; + async fn read_http_request(stream: &mut TcpStream) -> Vec { + let mut request = Vec::new(); + let mut chunk = [0; 1024]; + loop { + let read = stream.read(&mut chunk).await.unwrap(); + assert_ne!( + read, 0, + "HTTP client closed request before sending its body" + ); + request.extend_from_slice(&chunk[..read]); + let Some(header_end) = request.windows(4).position(|window| window == b"\r\n\r\n") + else { + continue; + }; + let headers = std::str::from_utf8(&request[..header_end]).unwrap(); + let content_length = headers + .lines() + .find_map(|line| line.strip_prefix("content-length: ")) + .or_else(|| { + headers + .lines() + .find_map(|line| line.strip_prefix("Content-Length: ")) + }) + .unwrap() + .parse::() + .unwrap(); + if request.len() >= header_end + 4 + content_length { + return request; + } + } + } + + async fn spawn_sql_server( + responses: Vec, + ) -> (u16, tokio::task::JoinHandle>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let server = tokio::spawn(async move { + let mut requests = Vec::with_capacity(responses.len()); + for body in responses { + let (mut stream, _) = listener.accept().await.unwrap(); + let request = read_http_request(&mut stream).await; + requests.push(String::from_utf8(request).unwrap()); + let response = format!( + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\ + Connection: close\r\nContent-Length: {}\r\n\r\n{body}", + body.len() + ); + stream.write_all(response.as_bytes()).await.unwrap(); + } + requests + }); + (port, server) + } + + #[tokio::test] + async fn setup_sql_runs_in_order_with_configured_database() { + let (port, server) = spawn_sql_server(vec!["{\"code\":0}".to_string(); 2]).await; + let client = Client::new(); + let setup = vec![ + "CREATE TABLE physical (ts TIMESTAMP TIME INDEX) ENGINE=metric".to_string(), + "ALTER TABLE physical SET 'compaction.twcs.trigger_file_num'='100'".to_string(), + ]; + + let results = run_setup_sql(&client, port, "perf", "candidate", &setup) + .await + .unwrap(); + + assert_eq!(results.len(), 2); + let requests = server.await.unwrap(); + assert_eq!(requests.len(), 2); + assert!(requests[0].contains("db=perf")); + assert!(requests[0].contains("sql=CREATE")); + assert!(requests[1].contains("db=perf")); + assert!(requests[1].contains("sql=ALTER")); + } + + #[tokio::test] + async fn setup_sql_stops_at_the_failing_statement() { + let (port, server) = + spawn_sql_server(vec!["{\"code\":0}".to_string(), "{\"code\":1}".to_string()]).await; + let client = Client::new(); + let setup = vec![ + "CREATE TABLE physical (ts TIMESTAMP TIME INDEX) ENGINE=metric".to_string(), + "invalid SQL".to_string(), + "ALTER TABLE physical SET 'compaction.twcs.trigger_file_num'='100'".to_string(), + ]; + + let error = run_setup_sql(&client, port, "perf", "base", &setup) + .await + .unwrap_err(); + + assert!(error.to_string().contains("statement 2 failed for base")); + let requests = server.await.unwrap(); + assert_eq!(requests.len(), 2); + assert!(requests[1].contains("sql=invalid+SQL")); + } + #[test] fn schedules_remote_sample_chunks_and_flushes() { let remote = RemoteWrite { @@ -452,6 +580,8 @@ mod tests { sample_chunk_size: Some(2), flush_every_sample_chunks: 2, visibility_timeout_seconds: 30, + base_setup_sql: Vec::new(), + candidate_setup_sql: Vec::new(), prom_store: PromStore { pending_rows_flush_interval: "1s".to_string(), max_batch_rows: 1, diff --git a/src/cmd/src/bin/query_regression_runner/remote/read_bench.rs b/src/cmd/src/bin/query_regression_runner/remote/read_bench.rs index 6b20a31c181..b9a05f6a69a 100644 --- a/src/cmd/src/bin/query_regression_runner/remote/read_bench.rs +++ b/src/cmd/src/bin/query_regression_runner/remote/read_bench.rs @@ -306,13 +306,25 @@ fn run_bench_command(command: Vec, mut run: Value) -> Result { } fn parse_average_duration(stdout: &str) -> Option { - Regex::new(r"(?i)Average duration[^0-9]*([0-9.]+)\s*ms") - .ok()? - .captures(stdout)? - .get(1)? - .as_str() - .parse() - .ok() + let captures = Regex::new( + r"(?i)Average:\s+(?:\x1b\[[0-9;]*m)*\d+(?:\x1b\[[0-9;]*m)*\s+rows(?:,\s+(?:\x1b\[[0-9;]*m)*\d+(?:\x1b\[[0-9;]*m)*\s+record batches)?\s+in\s+([0-9]+(?:\.[0-9]+)?)\s*(ns|µs|us|ms|s)\s+over\s+\d+\s+iterations|Average duration[^0-9]*([0-9]+(?:\.[0-9]+)?)\s*ms", + ) + .ok()? + .captures(stdout)?; + let (value, unit) = match (captures.get(1), captures.get(2), captures.get(3)) { + (Some(value), Some(unit), _) => (value.as_str(), unit.as_str()), + (_, _, Some(value)) => (value.as_str(), "ms"), + _ => return None, + }; + let value = value.parse::().ok()?; + let milliseconds = match unit.to_ascii_lowercase().as_str() { + "ns" => value / 1_000_000.0, + "µs" | "us" => value / 1_000.0, + "ms" => value, + "s" => value * 1_000.0, + _ => return None, + }; + milliseconds.is_finite().then_some(milliseconds) } fn bench_median(runs: &[Value]) -> Option { @@ -328,12 +340,51 @@ mod tests { use super::*; #[test] - fn parses_average_duration_and_groups_scan_paths() { - assert_eq!( - parse_average_duration("work\nAverage duration: 12.5 ms\n"), - Some(12.5) - ); - assert_eq!(parse_average_duration("completed"), None); + fn parses_average_durations() { + let cases = [ + ( + "parquetbench nanoseconds", + "ℹ Average: 42 rows, 2 record batches in 500ns over 4 iterations", + Some(0.0005), + ), + ( + "scanbench colored microseconds", + "\x1b[1;32mSummary\x1b[0m Average: \x1b[36m12\x1b[0m rows in 250µs over 2 iterations", + Some(0.25), + ), + ( + "scanbench seconds", + "Summary Average: 12 rows in 2.5s over 2 iterations", + Some(2500.0), + ), + ( + "legacy milliseconds", + "Average duration: 12.5 ms", + Some(12.5), + ), + ("absent", "Benchmark completed!", None), + ( + "malformed", + "Summary Average: 12 rows in unknown over 2 iterations", + None, + ), + ]; + + for (name, stdout, expected) in cases { + let actual = parse_average_duration(stdout); + match (actual, expected) { + (Some(actual), Some(expected)) => assert!( + (actual - expected).abs() < f64::EPSILON, + "{name}: expected {expected}, got {actual}" + ), + (None, None) => {} + (actual, expected) => panic!("{name}: expected {expected:?}, got {actual:?}"), + } + } + } + + #[test] + fn groups_scan_paths() { let targets = vec![ BenchTarget { relative_path: "data/a.parquet".to_string(), diff --git a/src/common/datasource/src/object_store.rs b/src/common/datasource/src/object_store.rs index 97bf779413e..37dd11aafdb 100644 --- a/src/common/datasource/src/object_store.rs +++ b/src/common/datasource/src/object_store.rs @@ -50,6 +50,22 @@ pub const AZBLOB_SCHEMA: &str = "AZBLOB"; pub struct BuiltBackend { pub object_store: ObjectStore, pub object_path: Option, + local_root: Option, +} + +impl BuiltBackend { + /// Checks an explicit input using the same file-type filter as directory listing. + pub async fn is_file(&self, path: &str) -> object_store::Result { + // Stat first preserves errors for dangling or inaccessible symlink targets. + if self.object_store.stat(path).await?.mode() != object_store::EntryMode::FILE { + return Ok(false); + } + if let Some(root) = &self.local_root { + root.is_file(path).await + } else { + Ok(true) + } + } } /// Controls whether SQL paths may access the local filesystem. @@ -379,6 +395,7 @@ async fn build_backend_inner( return Ok(BuiltBackend { object_store: build_fs_backend(&root)?, object_path, + local_root: Some(root), }); } @@ -418,6 +435,7 @@ async fn build_backend_inner( Ok(BuiltBackend { object_store, object_path, + local_root: None, }) } diff --git a/src/common/function/src/aggrs/vector/avg.rs b/src/common/function/src/aggrs/vector/avg.rs index ddf1823d283..822a7c16efe 100644 --- a/src/common/function/src/aggrs/vector/avg.rs +++ b/src/common/function/src/aggrs/vector/avg.rs @@ -183,8 +183,6 @@ mod tests { use std::sync::Arc; use arrow::array::StringArray; - use datatypes::scalars::ScalarVector; - use datatypes::vectors::{ConstantVector, StringVector, Vector}; use super::*; @@ -252,14 +250,12 @@ mod tests { vec_avg.evaluate().unwrap() ); - // test update with constant vector + // test update with repeated values let mut vec_avg = VectorAvg::default(); - let v: Vec = vec![ - Arc::new(ConstantVector::new( - Arc::new(StringVector::from_vec(vec!["[1.0,2.0,3.0]".to_string()])), - 4, - )) - .to_arrow_array(), + let v = vec![ + ScalarValue::Utf8(Some("[1.0,2.0,3.0]".to_string())) + .to_array_of_size(4) + .unwrap(), ]; vec_avg.update_batch(&v).unwrap(); assert_eq!( diff --git a/src/common/function/src/aggrs/vector/product.rs b/src/common/function/src/aggrs/vector/product.rs index d8f201966ab..11cfa4cb1a8 100644 --- a/src/common/function/src/aggrs/vector/product.rs +++ b/src/common/function/src/aggrs/vector/product.rs @@ -159,9 +159,6 @@ impl Accumulator for VectorProduct { mod tests { use std::sync::Arc; - use datatypes::scalars::ScalarVector; - use datatypes::vectors::{ConstantVector, StringVector, Vector}; - use super::*; #[test] @@ -213,14 +210,12 @@ mod tests { vec_product.update_batch(&v).unwrap(); assert_eq!(ScalarValue::Binary(None), vec_product.evaluate().unwrap()); - // test update with constant vector + // test update with repeated values let mut vec_product = VectorProduct::default(); - let v: Vec = vec![ - Arc::new(ConstantVector::new( - Arc::new(StringVector::from_vec(vec!["[1.0,2.0,3.0]".to_string()])), - 4, - )) - .to_arrow_array(), + let v = vec![ + ScalarValue::Utf8(Some("[1.0,2.0,3.0]".to_string())) + .to_array_of_size(4) + .unwrap(), ]; vec_product.update_batch(&v).unwrap(); diff --git a/src/common/function/src/aggrs/vector/sum.rs b/src/common/function/src/aggrs/vector/sum.rs index 3f875a8412c..0ee8f0df7bf 100644 --- a/src/common/function/src/aggrs/vector/sum.rs +++ b/src/common/function/src/aggrs/vector/sum.rs @@ -171,8 +171,6 @@ mod tests { use std::sync::Arc; use arrow::array::StringArray; - use datatypes::scalars::ScalarVector; - use datatypes::vectors::{ConstantVector, StringVector, Vector}; use super::*; @@ -225,14 +223,12 @@ mod tests { vec_sum.update_batch(&v).unwrap(); assert_eq!(ScalarValue::Binary(None), vec_sum.evaluate().unwrap()); - // test update with constant vector + // test update with repeated values let mut vec_sum = VectorSum::default(); - let v: Vec = vec![ - Arc::new(ConstantVector::new( - Arc::new(StringVector::from_vec(vec!["[1.0,2.0,3.0]".to_string()])), - 4, - )) - .to_arrow_array(), + let v = vec![ + ScalarValue::Utf8(Some("[1.0,2.0,3.0]".to_string())) + .to_array_of_size(4) + .unwrap(), ]; vec_sum.update_batch(&v).unwrap(); assert_eq!( diff --git a/src/common/function/src/scalars/expression.rs b/src/common/function/src/scalars/expression.rs index 63ed40fc8f7..777c91e7c0c 100644 --- a/src/common/function/src/scalars/expression.rs +++ b/src/common/function/src/scalars/expression.rs @@ -12,13 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. -mod binary; mod ctx; mod if_func; mod is_null; mod unary; -pub use binary::scalar_binary_op; pub use ctx::EvalContext; pub use unary::scalar_unary_op; diff --git a/src/common/function/src/scalars/expression/binary.rs b/src/common/function/src/scalars/expression/binary.rs deleted file mode 100644 index 4d9f6268a89..00000000000 --- a/src/common/function/src/scalars/expression/binary.rs +++ /dev/null @@ -1,90 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::iter; - -use common_query::error::Result; -use datatypes::prelude::*; -use datatypes::vectors::{ConstantVector, Helper}; - -use crate::scalars::expression::ctx::EvalContext; - -pub fn scalar_binary_op( - l: &VectorRef, - r: &VectorRef, - f: F, - ctx: &mut EvalContext, -) -> Result<::VectorType> -where - F: Fn(Option>, Option>, &mut EvalContext) -> Option, -{ - debug_assert!( - l.len() == r.len(), - "Size of vectors must match to apply binary expression" - ); - - let result = match (l.is_const(), r.is_const()) { - (false, true) => { - let left: &::VectorType = unsafe { Helper::static_cast(l) }; - let right: &ConstantVector = unsafe { Helper::static_cast(r) }; - let right: &::VectorType = unsafe { Helper::static_cast(right.inner()) }; - let b = right.get_data(0); - - let it = left.iter_data().map(|a| f(a, b.clone(), ctx)); - ::VectorType::from_owned_iterator(it) - } - - (false, false) => { - let left: &::VectorType = unsafe { Helper::static_cast(l) }; - let right: &::VectorType = unsafe { Helper::static_cast(r) }; - - let it = left - .iter_data() - .zip(right.iter_data()) - .map(|(a, b)| f(a, b, ctx)); - ::VectorType::from_owned_iterator(it) - } - - (true, false) => { - let left: &ConstantVector = unsafe { Helper::static_cast(l) }; - let left: &::VectorType = unsafe { Helper::static_cast(left.inner()) }; - let a = left.get_data(0); - - let right: &::VectorType = unsafe { Helper::static_cast(r) }; - let it = right.iter_data().map(|b| f(a.clone(), b, ctx)); - ::VectorType::from_owned_iterator(it) - } - - (true, true) => { - let left: &ConstantVector = unsafe { Helper::static_cast(l) }; - let left: &::VectorType = unsafe { Helper::static_cast(left.inner()) }; - let a = left.get_data(0); - - let right: &ConstantVector = unsafe { Helper::static_cast(r) }; - let right: &::VectorType = unsafe { Helper::static_cast(right.inner()) }; - let b = right.get_data(0); - - let it = iter::repeat(a) - .zip(iter::repeat(b)) - .map(|(a, b)| f(a, b, ctx)) - .take(left.len()); - ::VectorType::from_owned_iterator(it) - } - }; - - if let Some(error) = ctx.error.take() { - return Err(error); - } - Ok(result) -} diff --git a/src/common/function/src/scalars/test.rs b/src/common/function/src/scalars/test.rs index 8768d546a32..e7f59a65925 100644 --- a/src/common/function/src/scalars/test.rs +++ b/src/common/function/src/scalars/test.rs @@ -14,13 +14,43 @@ use std::fmt; +use common_query::error::Result; use datafusion::logical_expr::ColumnarValue; use datafusion_expr::{ScalarFunctionArgs, Signature, Volatility}; use datatypes::arrow::datatypes::DataType; +use datatypes::prelude::{Scalar, ScalarVector, VectorRef}; use datatypes::vectors::{Helper, Vector}; use crate::function::{Function, extract_args}; -use crate::scalars::expression::{EvalContext, scalar_binary_op}; +use crate::scalars::expression::EvalContext; + +fn scalar_binary_op( + l: &VectorRef, + r: &VectorRef, + f: F, + ctx: &mut EvalContext, +) -> Result<::VectorType> +where + F: Fn(Option>, Option>, &mut EvalContext) -> Option, +{ + debug_assert!( + l.len() == r.len(), + "Size of vectors must match to apply binary expression" + ); + + let left: &::VectorType = unsafe { Helper::static_cast(l) }; + let right: &::VectorType = unsafe { Helper::static_cast(r) }; + let result = ::VectorType::from_owned_iterator( + left.iter_data() + .zip(right.iter_data()) + .map(|(a, b)| f(a, b, ctx)), + ); + + if let Some(error) = ctx.error.take() { + return Err(error); + } + Ok(result) +} #[derive(Clone)] pub(crate) struct TestAndFunction { diff --git a/src/common/function/src/state.rs b/src/common/function/src/state.rs index 2b9f3a517a5..f60b37dd36d 100644 --- a/src/common/function/src/state.rs +++ b/src/common/function/src/state.rs @@ -12,12 +12,16 @@ // See the License for the specific language governing permissions and // limitations under the License. +use common_base::Plugins; + use crate::handlers::{FlowServiceHandlerRef, ProcedureServiceHandlerRef, TableMutationHandlerRef}; /// Shared state for SQL functions. /// The handlers in state may be `None` in cli command-line or test cases. #[derive(Clone, Default)] pub struct FunctionState { + /// Process plugins available to dynamically registered functions. + pub plugins: Plugins, // The table mutation handler pub table_mutation_handler: Option, // The procedure service handler @@ -212,6 +216,7 @@ impl FunctionState { } Self { + plugins: Plugins::default(), table_mutation_handler: Some(Arc::new(MockTableMutationHandler)), procedure_service_handler: Some(Arc::new(MockProcedureServiceHandler)), flow_service_handler: Some(Arc::new(MockFlowServiceHandler)), diff --git a/src/common/grpc-expr/src/alter.rs b/src/common/grpc-expr/src/alter.rs index 85de11d5cdb..d5214ae5ddb 100644 --- a/src/common/grpc-expr/src/alter.rs +++ b/src/common/grpc-expr/src/alter.rs @@ -49,33 +49,14 @@ const LOCATION_TYPE_AFTER: i32 = LocationType::After as i32; /// mixed batch — annotation alters skip region dispatch, so they cannot share /// a statement with options that regions must see. fn annotation_family_of_keys<'a>( - mut keys: impl Iterator, + keys: impl Iterator, ) -> Result> { - let Some(first) = keys.next() else { - return Ok(None); - }; - let family = AnnotationFamily::of_key(first); - let mut count = 1usize; - for key in keys { - count += 1; - let this = AnnotationFamily::of_key(key); - if this != family { - return error::InvalidTableOptionRequestSnafu { - err_msg: family.or(this).unwrap().mixed_batch_error(), - } - .fail(); + table::requests::validate_annotation_keys(keys).map_err(|err| { + error::InvalidTableOptionRequestSnafu { + err_msg: err.to_string(), } - } - if let Some(family) = family - && family.requires_single_key() - && count > 1 - { - return error::InvalidTableOptionRequestSnafu { - err_msg: family.mixed_batch_error(), - } - .fail(); - } - Ok(family) + .build() + }) } /// Returns the annotation family when `kind` is a SET/UNSET whose keys all @@ -583,6 +564,131 @@ mod tests { assert_eq!("mem_usage".to_string(), drop_names.pop().unwrap()); } + #[test] + fn test_repartition_hints_together() { + let keys = [ + REPARTITION_COLUMN_HINT_KEY, + table::requests::REPARTITION_PARTITION_NUM_HINT_KEY, + ]; + let options = vec![ + (keys[0].to_string(), "host".to_string()), + (keys[1].to_string(), "10".to_string()), + ]; + let set = Kind::SetTableOptions(SetTableOptions { + table_options: options + .iter() + .map(|(key, value)| PbOption { + key: key.clone(), + value: value.clone(), + }) + .collect(), + }); + let unset = Kind::UnsetTableOptions(UnsetTableOptions { + keys: keys.map(str::to_string).to_vec(), + }); + for kind in [set, unset] { + assert_eq!( + annotation_alter_family(&kind).unwrap(), + Some(AnnotationFamily::RepartitionHint) + ); + let request = alter_expr_to_request( + 1, + AlterTableExpr { + kind: Some(kind), + ..Default::default() + }, + None, + ) + .unwrap(); + match request.alter_kind { + AlterKind::SetAnnotations { + family, + options: actual, + } => { + assert_eq!(family, AnnotationFamily::RepartitionHint); + assert_eq!(actual, options); + } + AlterKind::UnsetAnnotations { + family, + keys: actual, + } => { + assert_eq!(family, AnnotationFamily::RepartitionHint); + assert_eq!(actual, keys); + } + other => panic!("unexpected alter kind: {other:?}"), + } + } + } + + #[test] + fn test_repartition_partition_num_hint_expr() { + let key = table::requests::REPARTITION_PARTITION_NUM_HINT_KEY; + let set = Kind::SetTableOptions(SetTableOptions { + table_options: vec![PbOption { + key: key.to_string(), + value: "8".to_string(), + }], + }); + let unset = Kind::UnsetTableOptions(UnsetTableOptions { + keys: vec![key.to_string()], + }); + for kind in [set, unset] { + assert_eq!( + annotation_alter_family(&kind).unwrap(), + Some(AnnotationFamily::RepartitionHint) + ); + let request = alter_expr_to_request( + 1, + AlterTableExpr { + kind: Some(kind), + ..Default::default() + }, + None, + ) + .unwrap(); + match request.alter_kind { + AlterKind::SetAnnotations { family, options } => { + assert_eq!(family, AnnotationFamily::RepartitionHint); + assert_eq!(options, vec![(key.to_string(), "8".to_string())]); + } + AlterKind::UnsetAnnotations { family, keys } => { + assert_eq!(family, AnnotationFamily::RepartitionHint); + assert_eq!(keys, vec![key.to_string()]); + } + other => panic!("unexpected alter kind: {other:?}"), + } + } + for other in [key, table::requests::TTL_KEY] { + for kind in [ + Kind::SetTableOptions(SetTableOptions { + table_options: [key, other] + .into_iter() + .map(|key| PbOption { + key: key.to_string(), + value: "8".to_string(), + }) + .collect(), + }), + Kind::UnsetTableOptions(UnsetTableOptions { + keys: vec![key.to_string(), other.to_string()], + }), + ] { + assert!(annotation_alter_family(&kind).is_err()); + assert!( + alter_expr_to_request( + 1, + AlterTableExpr { + kind: Some(kind), + ..Default::default() + }, + None + ) + .is_err() + ); + } + } + } + #[test] fn test_set_repartition_column_hint_expr() { let expr = AlterTableExpr { @@ -633,7 +739,7 @@ mod tests { let err = alter_expr_to_request(1, expr, None).unwrap_err(); assert!( err.to_string() - .contains("repartition.column.hint must be altered separately") + .contains("repartition hints must be altered separately") ); // Duplicate hint entries are not a meaningful batch either. @@ -655,10 +761,7 @@ mod tests { })), }; let err = alter_expr_to_request(1, dup, None).unwrap_err(); - assert!( - err.to_string() - .contains("repartition.column.hint must be altered separately") - ); + assert!(err.to_string().contains("duplicate repartition hint keys")); } #[test] diff --git a/src/common/macro/src/admin_fn.rs b/src/common/macro/src/admin_fn.rs index c580389f3c6..053a2ad3f46 100644 --- a/src/common/macro/src/admin_fn.rs +++ b/src/common/macro/src/admin_fn.rs @@ -304,7 +304,7 @@ fn build_struct( .and_then(|cv| match cv { common_query::prelude::ColumnarValue::Vector(v) => Ok(v), common_query::prelude::ColumnarValue::Scalar(s) => { - datatypes::vectors::Helper::try_from_scalar_value(s, args.number_rows) + datatypes::vectors::Helper::try_from_scalar_value(s, args.number_rows, None) .context(common_query::error::FromScalarValueSnafu) } }) diff --git a/src/common/macro/src/lib.rs b/src/common/macro/src/lib.rs index 5a58c4d7fbe..a397292787b 100644 --- a/src/common/macro/src/lib.rs +++ b/src/common/macro/src/lib.rs @@ -75,6 +75,10 @@ pub fn as_aggr_func_creator(args: TokenStream, input: TokenStream) -> TokenStrea /// - `name`: The name of the generated `ScalarUDF` struct. /// - `ret`: The return type of the generated UDF function. /// - `display_name`: The display name of the generated UDF function. +/// - `evaluator`: Optional path to a specialized evaluator with the calling convention +/// `fn(&[ColumnarValue], &str) -> Result`. When supplied, +/// the generated UDF `calc` delegates directly to it; without it, the default expansion is +/// unchanged. #[proc_macro_attribute] pub fn range_fn(args: TokenStream, input: TokenStream) -> TokenStream { process_range_fn(args, input) diff --git a/src/common/macro/src/range_fn.rs b/src/common/macro/src/range_fn.rs index 5ee21ca4b04..7d282db0fac 100644 --- a/src/common/macro/src/range_fn.rs +++ b/src/common/macro/src/range_fn.rs @@ -16,7 +16,7 @@ use proc_macro::TokenStream; use quote::quote; use syn::spanned::Spanned; use syn::{ - Attribute, Ident, ItemFn, Signature, Type, TypeReference, Visibility, parse_macro_input, + Attribute, Ident, ItemFn, Path, Signature, Type, TypeReference, Visibility, parse_macro_input, }; use crate::utils::extract_input_types; @@ -34,6 +34,7 @@ pub(crate) fn process_range_fn(args: TokenStream, input: TokenStream) -> TokenSt let mut name: Option = None; let mut display_name: Option = None; let mut ret: Option = None; + let mut evaluator: Option = None; let parser = syn::meta::parser(|meta| { if meta.path.is_ident("name") { @@ -45,6 +46,9 @@ pub(crate) fn process_range_fn(args: TokenStream, input: TokenStream) -> TokenSt } else if meta.path.is_ident("ret") { ret = Some(meta.value()?.parse()?); Ok(()) + } else if meta.path.is_ident("evaluator") { + evaluator = Some(meta.value()?.parse()?); + Ok(()) } else { Err(meta.error("unsupported property")) } @@ -103,10 +107,19 @@ pub(crate) fn process_range_fn(args: TokenStream, input: TokenStream) -> TokenSt arg_types, fn_name.clone(), ret.expect("ret required"), + evaluator.clone(), ); - // preserve this fn, but remove its `pub` modifier - let input_fn_code: TokenStream = quote! { - #sig { #block } + // Preserve this fn, but remove its `pub` modifier. Specialized evaluators + // do not call it in production, while tests keep it as the slice oracle. + let input_fn_code: TokenStream = if evaluator.is_some() { + quote! { + #[cfg(test)] + #sig { #block } + } + } else { + quote! { + #sig { #block } + } } .into(); @@ -161,7 +174,19 @@ fn build_calc_fn( param_types: Vec, fn_name: Ident, ret_type: Ident, + evaluator: Option, ) -> TokenStream { + if let Some(evaluator) = evaluator { + return quote! { + impl #name { + fn calc(input: &[ColumnarValue]) -> Result { + #evaluator(input, Self::name()) + } + } + } + .into(); + } + let param_names = param_types .iter() .enumerate() diff --git a/src/common/mem-prof/Cargo.toml b/src/common/mem-prof/Cargo.toml index ee4753ee885..10b67cf7fa7 100644 --- a/src/common/mem-prof/Cargo.toml +++ b/src/common/mem-prof/Cargo.toml @@ -16,7 +16,7 @@ tempfile = "3.4" tokio.workspace = true [target.'cfg(not(windows))'.dependencies] -tikv-jemalloc-ctl = { version = "0.6", features = ["use_std", "stats"] } +tikv-jemalloc-ctl = { version = "0.7", features = ["use_std", "stats"] } jemalloc-pprof-utils = { version = "0.8", package = "pprof_util", features = [ "flamegraph", "symbolize", @@ -25,4 +25,4 @@ jemalloc-pprof-mappings = { version = "0.7", package = "mappings" } # for get th [target.'cfg(not(windows))'.dependencies.tikv-jemalloc-sys] features = ["stats", "profiling", "unprefixed_malloc_on_supported_platforms"] -version = "0.6" +version = "0.7" diff --git a/src/common/meta/src/ddl/tests/alter_table.rs b/src/common/meta/src/ddl/tests/alter_table.rs index 30910eb0ba5..f0ceaf592e1 100644 --- a/src/common/meta/src/ddl/tests/alter_table.rs +++ b/src/common/meta/src/ddl/tests/alter_table.rs @@ -1123,58 +1123,85 @@ async fn test_on_submit_alter_request_with_exist_poison() { } #[tokio::test] -async fn test_semantic_annotation_alter_is_metadata_only() { - let (tx, mut rx) = mpsc::channel(8); - let node_manager = Arc::new(MockDatanodeManager::new(DatanodeWatcher::new(tx))); - let ddl_context = new_ddl_context(node_manager); - let table_id = 1024; - let table_name = "foo"; - let task = test_create_table_task(table_name, table_id); - ddl_context - .table_metadata_manager - .create_table_metadata( - task.table_info.clone(), - prepare_table_route(table_id), - HashMap::new(), - ) - .await - .unwrap(); +async fn test_annotation_alter_is_metadata_only() { + for (key, value) in [ + ("greptime.semantic.signal_type", "metric"), + (table::requests::REPARTITION_PARTITION_NUM_HINT_KEY, "8"), + ] { + let (tx, mut rx) = mpsc::channel(8); + let node_manager = Arc::new(MockDatanodeManager::new(DatanodeWatcher::new(tx))); + let ddl_context = new_ddl_context(node_manager); + let table_id = 1024; + let table_name = "foo"; + let task = test_create_table_task(table_name, table_id); + ddl_context + .table_metadata_manager + .create_table_metadata( + task.table_info.clone(), + prepare_table_route(table_id), + HashMap::new(), + ) + .await + .unwrap(); - let alter_table_task = AlterTableTask { - alter_table: AlterTableExpr { - catalog_name: DEFAULT_CATALOG_NAME.to_string(), - schema_name: DEFAULT_SCHEMA_NAME.to_string(), - table_name: table_name.to_string(), - kind: Some(Kind::SetTableOptions(SetTableOptions { - table_options: vec![api::v1::Option { - key: "greptime.semantic.signal_type".to_string(), - value: "metric".to_string(), - }], - })), - }, - }; - let mut procedure = - AlterTableProcedure::new(table_id, alter_table_task, ddl_context.clone()).unwrap(); - execute_procedure_until_done(&mut procedure).await; + let alter_table_task = AlterTableTask { + alter_table: AlterTableExpr { + catalog_name: DEFAULT_CATALOG_NAME.to_string(), + schema_name: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + kind: Some(Kind::SetTableOptions(SetTableOptions { + table_options: vec![api::v1::Option { + key: key.to_string(), + value: value.to_string(), + }], + })), + }, + }; + let mut procedure = + AlterTableProcedure::new(table_id, alter_table_task, ddl_context.clone()).unwrap(); + execute_procedure_until_done(&mut procedure).await; - // Metadata-only: no region request reaches any datanode. - rx.try_recv().unwrap_err(); + // Metadata-only: no region request reaches any datanode. + rx.try_recv().unwrap_err(); - let table_info = ddl_context - .table_metadata_manager - .table_info_manager() - .get(table_id) - .await - .unwrap() - .unwrap() - .into_inner() - .table_info; - assert_eq!( - table_info - .meta - .options - .extra_options - .get("greptime.semantic.signal_type"), - Some(&"metric".to_string()) - ); + let table_info = ddl_context + .table_metadata_manager + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap() + .into_inner() + .table_info; + assert_eq!( + table_info.meta.options.extra_options.get(key), + Some(&value.to_string()) + ); + for _ in 0..2 { + let task = AlterTableTask { + alter_table: AlterTableExpr { + catalog_name: DEFAULT_CATALOG_NAME.to_string(), + schema_name: DEFAULT_SCHEMA_NAME.to_string(), + table_name: table_name.to_string(), + kind: Some(Kind::UnsetTableOptions(api::v1::UnsetTableOptions { + keys: vec![key.to_string()], + })), + }, + }; + let mut procedure = + AlterTableProcedure::new(table_id, task, ddl_context.clone()).unwrap(); + execute_procedure_until_done(&mut procedure).await; + rx.try_recv().unwrap_err(); + let info = ddl_context + .table_metadata_manager + .table_info_manager() + .get(table_id) + .await + .unwrap() + .unwrap() + .into_inner() + .table_info; + assert!(!info.meta.options.extra_options.contains_key(key)); + } + } } diff --git a/src/common/meta/src/ddl_manager.rs b/src/common/meta/src/ddl_manager.rs index d88f08918e9..30911b7b272 100644 --- a/src/common/meta/src/ddl_manager.rs +++ b/src/common/meta/src/ddl_manager.rs @@ -2036,20 +2036,25 @@ mod tests { // The repartition hint drives physical repartitioning; on a logical // route it stays rejected by the guard. - let err = ddl_manager - .submit_ddl_task( - ExecutorContext { - query_context: Some(QueryContext::default()), - ..Default::default() - }, - SubmitDdlTaskRequest::new(DdlTask::new_alter_table(set_options_expr( - "logical", - &[("repartition.column.hint", "host")], - ))), - ) - .await - .unwrap_err(); - let msg = common_error::ext::ErrorExt::output_msg(&err); - assert!(msg.contains("non-physical TableRouteValue"), "{msg}"); + for (key, value) in [ + (table::requests::REPARTITION_COLUMN_HINT_KEY, "host"), + (table::requests::REPARTITION_PARTITION_NUM_HINT_KEY, "8"), + ] { + let err = ddl_manager + .submit_ddl_task( + ExecutorContext { + query_context: Some(QueryContext::default()), + ..Default::default() + }, + SubmitDdlTaskRequest::new(DdlTask::new_alter_table(set_options_expr( + "logical", + &[(key, value)], + ))), + ) + .await + .unwrap_err(); + let msg = common_error::ext::ErrorExt::output_msg(&err); + assert!(msg.contains("non-physical TableRouteValue"), "{msg}"); + } } } diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs index 3bd1475cb40..3f4868ffba0 100644 --- a/src/common/meta/src/error.rs +++ b/src/common/meta/src/error.rs @@ -269,6 +269,20 @@ pub enum Error { error: prost::DecodeError, }, + #[snafu(display("Failed to decode packed file references"))] + DecodePackedFileRefs { + #[snafu(implicit)] + location: Location, + #[snafu(source)] + error: base64::DecodeError, + }, + + #[snafu(display("Invalid packed file references framing"))] + InvalidPackedFileRefs { + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to encode object into json"))] EncodeJson { #[snafu(implicit)] @@ -1212,6 +1226,8 @@ impl ErrorExt for Error { | PayloadNotExist { .. } | ConvertRawKey { .. } | DecodeProto { .. } + | DecodePackedFileRefs { .. } + | InvalidPackedFileRefs { .. } | BuildTableMeta { .. } | TableRouteNotFound { .. } | TableRepartNotFound { .. } diff --git a/src/common/meta/src/heartbeat/handler.rs b/src/common/meta/src/heartbeat/handler.rs index ecc735083b5..fb5f16653a9 100644 --- a/src/common/meta/src/heartbeat/handler.rs +++ b/src/common/meta/src/heartbeat/handler.rs @@ -94,7 +94,21 @@ impl HeartbeatResponseHandlerExecutor for HandlerGroupExecutor { Ok(HandleControl::Done) => break, Ok(HandleControl::Continue) => {} Err(e) => { - error!(e;"Error while handling: {:?}", ctx.response); + let mailbox_message_id = ctx + .response + .mailbox_message + .as_ref() + .map(|message| message.id); + let json_payload_len = + ctx.response.mailbox_message.as_ref().and_then(|message| { + message.payload.as_ref().map(|payload| match payload { + api::v1::meta::mailbox_message::Payload::Json(json) => json.len(), + }) + }); + error!( + %e; + "Error while handling heartbeat response: mailbox_message_id={mailbox_message_id:?}, json_payload_len={json_payload_len:?}" + ); break; } } @@ -102,3 +116,86 @@ impl HeartbeatResponseHandlerExecutor for HandlerGroupExecutor { Ok(()) } } + +#[cfg(test)] +mod error_log_tests { + use std::fmt::Debug; + use std::sync::{Arc, Mutex}; + + use api::v1::meta::mailbox_message::Payload; + use api::v1::meta::{HeartbeatResponse, MailboxMessage}; + use common_telemetry::tracing::field::{Field, Visit}; + use common_telemetry::tracing::{Event, Subscriber}; + use common_telemetry::tracing_subscriber::layer::{Context, SubscriberExt}; + use common_telemetry::{tracing, tracing_subscriber}; + + use super::parse_mailbox_message::ParseMailboxMessageHandler; + use super::*; + use crate::heartbeat::mailbox::HeartbeatMailbox; + + #[derive(Clone, Default)] + struct LogCapture(Arc>>); + + impl tracing_subscriber::Layer for LogCapture + where + S: Subscriber, + { + fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) { + let mut visitor = FieldVisitor::default(); + event.record(&mut visitor); + self.0.lock().unwrap().push(visitor.fields.join(", ")); + } + } + + #[derive(Default)] + struct FieldVisitor { + fields: Vec, + } + + impl Visit for FieldVisitor { + fn record_debug(&mut self, field: &Field, value: &dyn Debug) { + self.fields.push(format!("{}={value:?}", field.name())); + } + + fn record_str(&mut self, field: &Field, value: &str) { + self.fields.push(format!("{}={value:?}", field.name())); + } + } + + #[tokio::test(flavor = "current_thread")] + async fn test_handler_error_log_excludes_mailbox_payload() { + let payload_sentinel = "MALFORMED_PACKED_PAYLOAD_MUST_NOT_BE_LOGGED".repeat(256); + let payload = format!( + r#"{{"PackedGcRegions":{{"regions":[],"packed_file_refs_manifest":"{}"}}"#, + payload_sentinel + ); + let payload_len = payload.len(); + let capture = LogCapture::default(); + let subscriber = tracing_subscriber::registry().with(capture.clone()); + let _guard = tracing::subscriber::set_default(subscriber); + let (mailbox_tx, _) = tokio::sync::mpsc::channel(1); + let ctx = HeartbeatResponseHandlerContext::new( + Arc::new(HeartbeatMailbox::new(mailbox_tx)), + HeartbeatResponse { + mailbox_message: Some(MailboxMessage { + id: 42, + subject: "unsafe subject".to_string(), + to: "unsafe recipient".to_string(), + from: "unsafe sender".to_string(), + payload: Some(Payload::Json(payload)), + ..Default::default() + }), + ..Default::default() + }, + ); + let executor = HandlerGroupExecutor::new(vec![Arc::new(ParseMailboxMessageHandler)]); + + executor.handle(ctx).await.unwrap(); + + let logs = capture.0.lock().unwrap().join("\n"); + assert!(logs.contains("Error while handling heartbeat response")); + assert!(logs.contains("mailbox_message_id=Some(42)")); + assert!(logs.contains(&format!("json_payload_len=Some({payload_len})"))); + assert!(!logs.contains(&payload_sentinel)); + } +} diff --git a/src/common/meta/src/instruction.rs b/src/common/meta/src/instruction.rs index ba4f2e32ab7..60ebe14e37c 100644 --- a/src/common/meta/src/instruction.rs +++ b/src/common/meta/src/instruction.rs @@ -16,16 +16,19 @@ use std::collections::{HashMap, HashSet}; use std::fmt::{Display, Formatter}; use std::time::Duration; +use base64::Engine; use common_error::ext::{ErrorExt, RetryHint}; use common_error::status_code::StatusCode; use serde::{Deserialize, Deserializer, Serialize, Serializer}; +use snafu::ResultExt as _; use store_api::region_engine::SyncRegionFromRequest; use store_api::region_request::{RegionFlushReason, RegionRequirements}; -use store_api::storage::{FileRefsManifest, GcReport, RegionId, RegionNumber}; +use store_api::storage::{FileId, FileRef, FileRefsManifest, GcReport, RegionId, RegionNumber}; use strum::Display; use table::metadata::TableId; use table::table_name::TableName; +use crate::error::{DecodePackedFileRefsSnafu, InvalidPackedFileRefsSnafu}; use crate::flow_name::FlowName; use crate::key::schema_name::SchemaName; use crate::key::{FlowId, FlowPartitionId}; @@ -579,6 +582,24 @@ impl Display for GetFileRefs { } } +/// Instruction to get file references using the packed manifest format. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct GetPackedFileRefs { + /// List of region IDs to get file references from active FileHandles (in-memory). + pub query_regions: Vec, + /// Mapping from the src region IDs (whose file references to look for) to + /// the dst region IDs (where to read the manifests). + /// Key: The source region IDs (where files originally came from). + /// Value: The set of destination region IDs (whose manifests need to be read). + pub related_regions: HashMap>, +} + +impl Display for GetPackedFileRefs { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "GetPackedFileRefs(region_ids={:?})", self.query_regions) + } +} + /// Instruction to trigger garbage collection for a region. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct GcRegions { @@ -590,6 +611,147 @@ pub struct GcRegions { pub full_file_listing: bool, } +/// Instruction to trigger garbage collection with a packed file-reference manifest. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct PackedGcRegions { + /// The region ID to perform GC on. + pub regions: Vec, + /// Packed manifest used by the packed-only GC protocol. + pub packed_file_refs_manifest: PackedFileRefsManifest, + /// Whether to perform a full file listing to find orphan files. + pub full_file_listing: bool, +} + +impl Display for PackedGcRegions { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + let file_refs_count = self.packed_file_refs_manifest.file_refs.len(); + write!( + f, + "PackedGcRegions(regions={:?}, file_refs_count={}, full_file_listing={})", + self.regions, file_refs_count, self.full_file_listing + ) + } +} + +/// A packed, JSON-compatible encoding of a file reference manifest. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)] +pub struct PackedFileRefsManifest { + /// Required packed file references for each region. + pub file_refs: HashMap, + #[serde(default)] + pub manifest_version: HashMap, + #[serde(default)] + pub cross_region_refs: HashMap>, +} + +/// Packed file references for one region, encoded with STANDARD Base64. +/// +/// `files` is required and contains 16-byte UUID records for references whose +/// `index_version` is `None`. `indexed` is required and contains 24-byte records +/// for references whose `index_version` is `Some`: a 16-byte UUID followed by +/// an 8-byte big-endian `u64` version. Decoding fails when either field is not +/// valid STANDARD Base64 or its decoded bytes do not align to the corresponding +/// record width. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)] +pub struct PackedRegionFileRefs { + /// Required STANDARD Base64-encoded 16-byte UUID records. + pub files: String, + /// Required STANDARD Base64-encoded 16-byte UUID and 8-byte version records. + pub indexed: String, +} + +impl PackedFileRefsManifest { + /// Packs a file reference manifest using required [`PackedRegionFileRefs`] + /// STANDARD Base64 fields. + /// + /// References whose `index_version` is `None` are encoded as 16-byte UUID + /// `files` records; references whose `index_version` is `Some` are encoded + /// as 24-byte `indexed` records: a 16-byte UUID followed by an 8-byte + /// big-endian `u64` version. + pub fn from_manifest(manifest: &FileRefsManifest) -> Self { + Self { + file_refs: manifest + .file_refs + .iter() + .map(|(region, refs)| (*region, PackedRegionFileRefs::from_refs(refs))) + .collect(), + manifest_version: manifest.manifest_version.clone(), + cross_region_refs: manifest.cross_region_refs.clone(), + } + } + + /// Decodes this packed representation into a file reference manifest. + /// + /// Returns a typed error for invalid STANDARD Base64 or malformed record + /// framing without retaining encoded payloads. + pub fn into_manifest(self) -> crate::error::Result { + let mut file_refs = HashMap::new(); + for (region, encoded) in self.file_refs { + file_refs.insert(region, encoded.into_refs(region)?); + } + Ok(FileRefsManifest { + file_refs, + manifest_version: self.manifest_version, + cross_region_refs: self.cross_region_refs, + }) + } +} + +impl PackedRegionFileRefs { + /// Packs file references as the documented STANDARD Base64 record streams. + pub fn from_refs(refs: &HashSet) -> Self { + let engine = base64::engine::general_purpose::STANDARD; + let mut files = Vec::new(); + let mut indexed = Vec::new(); + for file_ref in refs { + match file_ref.index_version { + None => files.extend_from_slice(file_ref.file_id.as_bytes()), + Some(version) => { + indexed.extend_from_slice(file_ref.file_id.as_bytes()); + indexed.extend_from_slice(&version.to_be_bytes()); + } + } + } + Self { + files: engine.encode(files), + indexed: engine.encode(indexed), + } + } + + /// Decodes the documented STANDARD Base64 record streams for `region`. + fn into_refs(self, region: RegionId) -> crate::error::Result> { + let engine = base64::engine::general_purpose::STANDARD; + let files = engine + .decode(self.files) + .context(DecodePackedFileRefsSnafu)?; + let indexed = engine + .decode(self.indexed) + .context(DecodePackedFileRefsSnafu)?; + if files.len() % 16 != 0 || indexed.len() % 24 != 0 { + return InvalidPackedFileRefsSnafu.fail(); + } + + let mut refs = HashSet::new(); + for bytes in files.chunks_exact(16) { + let mut id = [0; 16]; + id.copy_from_slice(bytes); + refs.insert(FileRef::new(region, FileId::from_bytes(id), None)); + } + for bytes in indexed.chunks_exact(24) { + let mut id = [0; 16]; + id.copy_from_slice(&bytes[..16]); + let mut version = [0; 8]; + version.copy_from_slice(&bytes[16..]); + refs.insert(FileRef::new( + region, + FileId::from_bytes(id), + Some(u64::from_be_bytes(version)), + )); + } + Ok(refs) + } +} + impl Display for GcRegions { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { write!( @@ -625,6 +787,26 @@ impl Display for GetFileRefsReply { } } +/// Reply for GetPackedFileRefs instruction. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct GetPackedFileRefsReply { + pub packed_file_refs_manifest: PackedFileRefsManifest, + pub success: bool, + pub error: Option, +} + +impl Display for GetPackedFileRefsReply { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "GetPackedFileRefsReply(success={}, file_refs_count={}, error={:?})", + self.success, + self.packed_file_refs_manifest.file_refs.len(), + self.error + ) + } +} + /// Reply for GC instruction. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct GcRegionsReply { @@ -827,8 +1009,12 @@ pub enum Instruction { FlushRegions(FlushRegions), /// Gets file references for regions. GetFileRefs(GetFileRefs), + /// Gets file references using the packed manifest format. + GetPackedFileRefs(GetPackedFileRefs), /// Triggers garbage collection for a region. GcRegions(GcRegions), + /// Triggers garbage collection with a packed file-reference manifest. + PackedGcRegions(PackedGcRegions), /// Temporary suspend serving reads or writes Suspend, /// Makes regions enter staging state. @@ -1111,6 +1297,7 @@ pub enum InstructionReply { DowngradeRegions(DowngradeRegionsReply), FlushRegions(FlushRegionReply), GetFileRefs(GetFileRefsReply), + GetPackedFileRefs(GetPackedFileRefsReply), GcRegions(GcRegionsReply), EnterStagingRegions(EnterStagingRegionsReply), SyncRegions(SyncRegionsReply), @@ -1131,6 +1318,9 @@ impl Display for InstructionReply { } Self::FlushRegions(reply) => write!(f, "InstructionReply::FlushRegions({})", reply), Self::GetFileRefs(reply) => write!(f, "InstructionReply::GetFileRefs({})", reply), + Self::GetPackedFileRefs(reply) => { + write!(f, "InstructionReply::GetPackedFileRefs({})", reply) + } Self::GcRegions(reply) => write!(f, "InstructionReply::GcRegion({})", reply), Self::EnterStagingRegions(reply) => { write!( @@ -1221,6 +1411,7 @@ impl InstructionReply { #[cfg(test)] mod tests { use std::collections::HashSet; + use std::error::Error as _; use common_error::mock::MockError; use common_wal::options::WalOptions; @@ -1771,6 +1962,161 @@ mod tests { } } + #[test] + fn test_legacy_gc_file_refs_compatibility() { + #[derive(Debug, Deserialize)] + struct LegacyGetFileRefs { + query_regions: Vec, + related_regions: HashMap>, + } + + #[derive(Debug, Deserialize)] + struct LegacyGetFileRefsReply { + file_refs_manifest: FileRefsManifest, + success: bool, + error: Option, + } + + #[derive(Debug, Deserialize)] + enum LegacyInstruction { + GetFileRefs(LegacyGetFileRefs), + } + + let get_file_refs = Instruction::GetFileRefs(GetFileRefs { + query_regions: vec![RegionId::new(7, 3)], + related_regions: HashMap::new(), + }); + let get_file_refs_json = serde_json::to_string(&get_file_refs).unwrap(); + let legacy_get_file_refs: LegacyInstruction = + serde_json::from_str(&get_file_refs_json).unwrap(); + let LegacyInstruction::GetFileRefs(legacy_get_file_refs) = legacy_get_file_refs; + assert_eq!(legacy_get_file_refs.query_regions.len(), 1); + assert!(legacy_get_file_refs.related_regions.is_empty()); + + let reply = GetFileRefsReply { + file_refs_manifest: FileRefsManifest::default(), + success: true, + error: None, + }; + let reply_json = serde_json::to_string(&reply).unwrap(); + let legacy_reply: LegacyGetFileRefsReply = serde_json::from_str(&reply_json).unwrap(); + let current_reply: GetFileRefsReply = serde_json::from_str(&reply_json).unwrap(); + assert!(legacy_reply.success); + assert!(legacy_reply.error.is_none()); + assert!(legacy_reply.file_refs_manifest.file_refs.is_empty()); + assert_eq!(current_reply, reply); + + let packed = Instruction::PackedGcRegions(PackedGcRegions { + regions: vec![], + packed_file_refs_manifest: PackedFileRefsManifest::default(), + full_file_listing: false, + }); + let packed_json = serde_json::to_string(&packed).unwrap(); + assert!(serde_json::from_str::(&packed_json).is_err()); + } + + #[test] + fn test_packed_gc_regions_round_trip_is_distinct() { + let instruction = Instruction::PackedGcRegions(PackedGcRegions { + regions: vec![RegionId::new(7, 3)], + packed_file_refs_manifest: PackedFileRefsManifest::default(), + full_file_listing: true, + }); + let serialized = serde_json::to_string(&instruction).unwrap(); + assert!(serialized.contains("PackedGcRegions")); + assert_eq!( + serde_json::from_str::(&serialized).unwrap(), + instruction + ); + } + + #[test] + fn test_packed_file_refs_manifest_round_trip_and_validation() { + let region = RegionId::new(7, 3); + let file_id = FileId::parse_str("00000000-0000-0000-0000-000000000001").unwrap(); + let mut manifest = FileRefsManifest::default(); + manifest.file_refs.insert( + region, + HashSet::from([ + FileRef::new(region, file_id, None), + FileRef::new(region, file_id, Some(0)), + FileRef::new(region, file_id, Some(u64::MAX)), + ]), + ); + manifest.manifest_version.insert(region, 42); + manifest + .cross_region_refs + .insert(region, HashSet::from([RegionId::new(7, 4)])); + let packed = PackedFileRefsManifest::from_manifest(&manifest); + assert_eq!(packed.clone().into_manifest().unwrap(), manifest); + let mut missing_files = serde_json::to_value(&packed).unwrap(); + missing_files + .get_mut("file_refs") + .and_then(serde_json::Value::as_object_mut) + .and_then(|file_refs| file_refs.values_mut().next()) + .and_then(serde_json::Value::as_object_mut) + .unwrap() + .remove("files"); + assert!(serde_json::from_value::(missing_files).is_err()); + + let mut missing_indexed = serde_json::to_value(&packed).unwrap(); + missing_indexed + .get_mut("file_refs") + .and_then(serde_json::Value::as_object_mut) + .and_then(|file_refs| file_refs.values_mut().next()) + .and_then(serde_json::Value::as_object_mut) + .unwrap() + .remove("indexed"); + assert!(serde_json::from_value::(missing_indexed).is_err()); + + let mut malformed = packed.clone(); + malformed.file_refs.get_mut(®ion).unwrap().indexed = "!".to_string(); + let err = malformed.into_manifest().unwrap_err(); + assert!(matches!( + &err, + crate::error::Error::DecodePackedFileRefs { .. } + )); + assert!( + err.source() + .is_some_and(|source| source.is::()) + ); + assert_eq!(err.status_code(), StatusCode::Unexpected); + assert_eq!(err.retry_hint(), RetryHint::NonRetryable); + + let mut malformed = packed; + malformed.file_refs.get_mut(®ion).unwrap().indexed = + base64::engine::general_purpose::STANDARD.encode([0; 1]); + let err = malformed.into_manifest().unwrap_err(); + assert!(matches!( + &err, + crate::error::Error::InvalidPackedFileRefs { .. } + )); + assert!(err.source().is_none()); + assert_eq!(err.status_code(), StatusCode::Unexpected); + assert_eq!(err.retry_hint(), RetryHint::NonRetryable); + } + + #[test] + fn test_packed_file_refs_is_smaller_than_legacy() { + let region = RegionId::new(7, 3); + let mut manifest = FileRefsManifest::default(); + manifest.file_refs.insert( + region, + (0..500) + .map(|i| { + let id = + FileId::parse_str(&format!("00000000-0000-0000-0000-{i:012}")).unwrap(); + FileRef::new(region, id, None) + }) + .collect(), + ); + let legacy = serde_json::to_vec(&manifest).unwrap().len(); + let packed = serde_json::to_vec(&PackedFileRefsManifest::from_manifest(&manifest)) + .unwrap() + .len(); + assert!(packed * 2 < legacy, "packed={packed}, legacy={legacy}"); + } + #[test] fn test_serialize_get_file_refs_instruction_reply() { let mut manifest = FileRefsManifest::default(); diff --git a/src/common/meta/src/reconciliation/event.rs b/src/common/meta/src/reconciliation/event.rs index cff4a0db19d..06629bc48c9 100644 --- a/src/common/meta/src/reconciliation/event.rs +++ b/src/common/meta/src/reconciliation/event.rs @@ -28,6 +28,10 @@ use store_api::storage::TableId; use crate::reconciliation::ResolveStrategy; +/// Stable event type stored for catalog reconciliation procedures. +pub(crate) const RECONCILE_CATALOG_EVENT_TYPE: &str = "reconcile_catalog"; +/// Stable event type stored for database reconciliation procedures. +pub(crate) const RECONCILE_DATABASE_EVENT_TYPE: &str = "reconcile_database"; /// Stable event type stored for logical table reconciliation procedures. pub(crate) const RECONCILE_LOGICAL_TABLES_EVENT_TYPE: &str = "reconcile_logical_tables"; /// Stable event type stored for physical table reconciliation procedures. @@ -45,6 +49,23 @@ pub(crate) struct ReconciliationLocator { } impl ReconciliationLocator { + /// Creates a locator for a catalog. + pub(crate) fn catalog(catalog_name: &str) -> Self { + Self { + catalog_name: Some(catalog_name.to_string()), + ..Default::default() + } + } + + /// Creates a locator for a database. + pub(crate) fn database(catalog_name: &str, schema_name: &str) -> Self { + Self { + catalog_name: Some(catalog_name.to_string()), + schema_name: Some(schema_name.to_string()), + ..Default::default() + } + } + /// Creates a locator for a physical table with its fully qualified name and ID. pub(crate) fn physical_table( catalog_name: &str, @@ -101,6 +122,229 @@ impl ReconciliationLocator { } } +#[derive(Debug, Serialize)] +#[serde(untagged)] +enum ReconcileCatalogPayload { + Submitted(CatalogSubmittedPayload), + Result(CatalogResultPayload), +} + +#[derive(Debug, Serialize)] +struct CatalogSubmittedPayload { + version: u8, + resolve_strategy: &'static str, + fail_fast: bool, + parallelism: usize, +} + +#[derive(Debug, Serialize)] +struct CatalogResultPayload { + version: u8, + complete: bool, + processed_database_count: usize, + succeeded_database_count: usize, + failed_database_count: usize, +} + +/// Event representation for catalog reconciliation. +#[derive(Debug)] +pub(crate) struct ReconcileCatalogEvent { + locator: ReconciliationLocator, + payload: Option, +} + +impl ReconcileCatalogEvent { + /// Builds the bounded intent event emitted when catalog reconciliation is submitted. + pub(crate) fn submitted( + locator: ReconciliationLocator, + resolve_strategy: ResolveStrategy, + fail_fast: bool, + parallelism: usize, + ) -> Self { + Self { + locator, + payload: Some(ReconcileCatalogPayload::Submitted( + CatalogSubmittedPayload { + version: PAYLOAD_VERSION, + resolve_strategy: resolve_strategy_name(resolve_strategy), + fail_fast, + parallelism, + }, + )), + } + } + + /// Builds a terminal event from the existing volatile reconciliation metrics. + /// + /// Metrics are best-effort observations from the current process and reset on recovery. + pub(crate) fn result( + locator: ReconciliationLocator, + complete: bool, + succeeded_database_count: usize, + failed_database_count: usize, + ) -> Self { + Self { + locator, + payload: Some(ReconcileCatalogPayload::Result(CatalogResultPayload { + version: PAYLOAD_VERSION, + complete, + processed_database_count: succeeded_database_count + failed_database_count, + succeeded_database_count, + failed_database_count, + })), + } + } + + /// Builds a catalog lifecycle event whose reconciliation payload is null. + pub(crate) fn lifecycle(locator: ReconciliationLocator) -> Self { + Self { + locator, + payload: None, + } + } +} + +impl Event for ReconcileCatalogEvent { + fn event_type(&self) -> &str { + RECONCILE_CATALOG_EVENT_TYPE + } + + fn json_payload(&self) -> Result { + match &self.payload { + Some(payload) => serde_json::to_value(payload).context(SerializeEventSnafu), + None => Ok(serde_json::Value::Null), + } + } + + fn extra_schema(&self) -> Vec { + ReconciliationLocator::schema() + } + + fn extra_rows(&self) -> Result> { + Ok(vec![self.locator.row()]) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +#[derive(Debug, Serialize)] +#[serde(untagged)] +enum ReconcileDatabasePayload { + Submitted(DatabaseSubmittedPayload), + Result(DatabaseResultPayload), +} + +#[derive(Debug, Serialize)] +struct DatabaseSubmittedPayload { + version: u8, + resolve_strategy: &'static str, + fail_fast: bool, + parallelism: usize, + is_subprocedure: bool, +} + +#[derive(Debug, Serialize)] +struct DatabaseResultPayload { + version: u8, + complete: bool, + processed_table_count: usize, + succeeded_table_count: usize, + failed_table_count: usize, + succeeded_subprocedure_count: usize, + failed_subprocedure_count: usize, +} + +/// Event representation for database reconciliation. +#[derive(Debug)] +pub(crate) struct ReconcileDatabaseEvent { + locator: ReconciliationLocator, + payload: Option, +} + +impl ReconcileDatabaseEvent { + /// Builds the bounded intent event emitted when database reconciliation is submitted. + pub(crate) fn submitted( + locator: ReconciliationLocator, + resolve_strategy: ResolveStrategy, + fail_fast: bool, + parallelism: usize, + is_subprocedure: bool, + ) -> Self { + Self { + locator, + payload: Some(ReconcileDatabasePayload::Submitted( + DatabaseSubmittedPayload { + version: PAYLOAD_VERSION, + resolve_strategy: resolve_strategy_name(resolve_strategy), + fail_fast, + parallelism, + is_subprocedure, + }, + )), + } + } + + /// Builds a terminal event from the existing volatile reconciliation metrics. + /// + /// Metrics are best-effort observations from the current process and reset on recovery. + pub(crate) fn result( + locator: ReconciliationLocator, + complete: bool, + succeeded_table_count: usize, + failed_table_count: usize, + succeeded_subprocedure_count: usize, + failed_subprocedure_count: usize, + ) -> Self { + Self { + locator, + payload: Some(ReconcileDatabasePayload::Result(DatabaseResultPayload { + version: PAYLOAD_VERSION, + complete, + processed_table_count: succeeded_table_count + failed_table_count, + succeeded_table_count, + failed_table_count, + succeeded_subprocedure_count, + failed_subprocedure_count, + })), + } + } + + /// Builds a database lifecycle event whose reconciliation payload is null. + pub(crate) fn lifecycle(locator: ReconciliationLocator) -> Self { + Self { + locator, + payload: None, + } + } +} + +impl Event for ReconcileDatabaseEvent { + fn event_type(&self) -> &str { + RECONCILE_DATABASE_EVENT_TYPE + } + + fn json_payload(&self) -> Result { + match &self.payload { + Some(payload) => serde_json::to_value(payload).context(SerializeEventSnafu), + None => Ok(serde_json::Value::Null), + } + } + + fn extra_schema(&self) -> Vec { + ReconciliationLocator::schema() + } + + fn extra_rows(&self) -> Result> { + Ok(vec![self.locator.row()]) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + #[derive(Debug, Serialize)] #[serde(untagged)] enum ReconcileTablePayload { @@ -355,12 +599,14 @@ mod tests { #[test] fn reconciliation_events_use_the_shared_locator_contract() { + let catalog = ReconcileCatalogEvent::lifecycle(ReconciliationLocator::catalog("greptime")); + assert_eq!(catalog.event_type(), RECONCILE_CATALOG_EVENT_TYPE); let table = ReconcileTableEvent::table_lifecycle(ReconciliationLocator::physical_table( "greptime", "public", "metrics", 42, )); assert_eq!(table.event_type(), RECONCILE_TABLE_EVENT_TYPE); assert_eq!( - table + catalog .extra_schema() .into_iter() .map(|column| { @@ -375,22 +621,22 @@ mod tests { ( "catalog_name".to_string(), ColumnDataType::String, - SemanticType::Field + SemanticType::Field, ), ( "schema_name".to_string(), ColumnDataType::String, - SemanticType::Field + SemanticType::Field, ), ( "table_name".to_string(), ColumnDataType::String, - SemanticType::Field + SemanticType::Field, ), ( "table_id".to_string(), ColumnDataType::Uint32, - SemanticType::Field + SemanticType::Field, ), ( "physical_table_id".to_string(), @@ -399,6 +645,40 @@ mod tests { ), ] ); + assert_eq!( + catalog.extra_rows().unwrap(), + vec![Row { + values: vec![ + ValueData::StringValue("greptime".to_string()).into(), + Value::default(), + Value::default(), + Value::default(), + Value::default(), + ], + }] + ); + assert_eq!(catalog.json_payload().unwrap(), serde_json::Value::Null); + + let database = ReconcileDatabaseEvent::lifecycle(ReconciliationLocator::database( + "greptime", "public", + )); + assert_eq!(database.event_type(), RECONCILE_DATABASE_EVENT_TYPE); + assert_eq!(database.extra_schema(), catalog.extra_schema()); + assert_eq!( + database.extra_rows().unwrap(), + vec![Row { + values: vec![ + ValueData::StringValue("greptime".to_string()).into(), + ValueData::StringValue("public".to_string()).into(), + Value::default(), + Value::default(), + Value::default(), + ], + }] + ); + assert_eq!(database.json_payload().unwrap(), serde_json::Value::Null); + + assert_eq!(table.extra_schema(), catalog.extra_schema()); assert_eq!( table.extra_rows().unwrap(), vec![Row { @@ -458,6 +738,22 @@ mod tests { (ResolveStrategy::UseMetasrv, "use_metasrv"), (ResolveStrategy::AbortOnConflict, "abort_on_conflict"), ] { + let catalog = ReconcileCatalogEvent::submitted( + ReconciliationLocator::catalog("greptime"), + strategy, + true, + 16, + ); + assert_eq!( + catalog.json_payload().unwrap(), + json!({ + "version": 1, + "resolve_strategy": expected, + "fail_fast": true, + "parallelism": 16, + }) + ); + let table = ReconcileTableEvent::table_submitted( ReconciliationLocator::physical_table("greptime", "public", "metrics", 42), strategy, @@ -473,6 +769,24 @@ mod tests { ); } + let database = ReconcileDatabaseEvent::submitted( + ReconciliationLocator::database("greptime", "public"), + ResolveStrategy::UseMetasrv, + false, + 64, + true, + ); + assert_eq!( + database.json_payload().unwrap(), + json!({ + "version": 1, + "resolve_strategy": "use_metasrv", + "fail_fast": false, + "parallelism": 64, + "is_subprocedure": true, + }) + ); + let logical_tables = ReconcileLogicalTablesEvent::submitted( vec![ ReconciliationLocator::logical_table("greptime", "public", "cpu", 43, 42), @@ -492,6 +806,40 @@ mod tests { #[test] fn terminal_payloads_distinguish_complete_and_partial_results() { + let catalog = + ReconcileCatalogEvent::result(ReconciliationLocator::catalog("greptime"), true, 3, 1); + assert_eq!( + catalog.json_payload().unwrap(), + json!({ + "version": 1, + "complete": true, + "processed_database_count": 4, + "succeeded_database_count": 3, + "failed_database_count": 1, + }) + ); + + let database = ReconcileDatabaseEvent::result( + ReconciliationLocator::database("greptime", "public"), + false, + 5, + 2, + 4, + 1, + ); + assert_eq!( + database.json_payload().unwrap(), + json!({ + "version": 1, + "complete": false, + "processed_table_count": 7, + "succeeded_table_count": 5, + "failed_table_count": 2, + "succeeded_subprocedure_count": 4, + "failed_subprocedure_count": 1, + }) + ); + let table = ReconcileTableEvent::table_result( ReconciliationLocator::physical_table("greptime", "public", "metrics", 42), false, diff --git a/src/common/meta/src/reconciliation/reconcile_catalog.rs b/src/common/meta/src/reconciliation/reconcile_catalog.rs index d3c2e526dd3..e28e81d323d 100644 --- a/src/common/meta/src/reconciliation/reconcile_catalog.rs +++ b/src/common/meta/src/reconciliation/reconcile_catalog.rs @@ -18,8 +18,8 @@ use std::time::Instant; use common_procedure::error::FromJsonSnafu; use common_procedure::{ - Context as ProcedureContext, Error as ProcedureError, LockKey, Procedure, - Result as ProcedureResult, Status, + Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, + Procedure, Result as ProcedureResult, Status, }; use futures::stream::BoxStream; use serde::{Deserialize, Serialize}; @@ -31,6 +31,9 @@ use crate::key::TableMetadataManagerRef; use crate::lock_key::CatalogLock; use crate::metrics; use crate::node_manager::NodeManagerRef; +use crate::reconciliation::event::{ + RECONCILE_CATALOG_EVENT_TYPE, ReconcileCatalogEvent, ReconciliationLocator, +}; use crate::reconciliation::reconcile_catalog::start::ReconcileCatalogStart; use crate::reconciliation::reconcile_table::resolve_column_metadata::ResolveStrategy; use crate::reconciliation::utils::{ @@ -216,6 +219,43 @@ impl Procedure for ReconcileCatalogProcedure { LockKey::new(vec![CatalogLock::Write(catalog).into()]) } + + fn event(&self, ctx: &EventContext<'_>) -> Option> { + if !ctx.event_type_filter.allows(RECONCILE_CATALOG_EVENT_TYPE) { + return None; + } + + let persistent_ctx = &self.context.persistent_ctx; + let locator = ReconciliationLocator::catalog(&persistent_ctx.catalog); + let event = match ctx.trigger { + EventTrigger::Submitted => ReconcileCatalogEvent::submitted( + locator, + persistent_ctx.resolve_strategy, + persistent_ctx.fast_fail, + persistent_ctx.parallelism, + ), + EventTrigger::Succeeded => self.result_event(locator, true), + EventTrigger::Failed | EventTrigger::Poisoned => self.result_event(locator, false), + _ => ReconcileCatalogEvent::lifecycle(locator), + }; + Some(Box::new(event)) + } +} + +impl ReconcileCatalogProcedure { + fn result_event( + &self, + locator: ReconciliationLocator, + complete: bool, + ) -> ReconcileCatalogEvent { + let metrics = &self.context.volatile_ctx.metrics; + ReconcileCatalogEvent::result( + locator, + complete, + metrics.succeeded_databases, + metrics.failed_databases, + ) + } } #[async_trait::async_trait] @@ -235,3 +275,217 @@ pub(crate) trait State: Sync + Send + Debug { fn as_any(&self) -> &dyn Any; } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_event_recorder::{EventTypeFilter, EventTypeFilterRef}; + use common_procedure::{ + ChildSubmissionOutcome, EventContext, EventTrigger, Procedure, ProcedureId, ProcedureState, + RetryPhase, + }; + use serde_json::{Value, json}; + + use super::*; + use crate::reconciliation::event::RECONCILE_DATABASE_EVENT_TYPE; + use crate::test_util::{MockDatanodeManager, new_ddl_context}; + + struct CatalogEventHarness { + procedure_id: ProcedureId, + lifecycle_state: ProcedureState, + event_type_filter: EventTypeFilterRef, + } + + impl CatalogEventHarness { + fn all() -> Self { + Self { + procedure_id: ProcedureId::random(), + lifecycle_state: ProcedureState::Running, + event_type_filter: Arc::new(EventTypeFilter::All), + } + } + + fn selected(event_types: impl IntoIterator) -> Self { + Self { + event_type_filter: Arc::new(EventTypeFilter::Only( + event_types.into_iter().map(str::to_string).collect(), + )), + ..Self::all() + } + } + + fn event( + &self, + procedure: &dyn Procedure, + trigger: EventTrigger, + ) -> Option> { + procedure.event(&EventContext { + procedure_id: self.procedure_id, + lifecycle_state: &self.lifecycle_state, + trigger, + event_type_filter: self.event_type_filter.clone(), + event_context: None, + }) + } + } + + #[test] + fn catalog_submitted_event_reports_intent() { + let submitted = CatalogEventHarness::all() + .event(&test_procedure(), EventTrigger::Submitted) + .unwrap(); + assert_eq!(submitted.event_type(), RECONCILE_CATALOG_EVENT_TYPE); + assert_eq!( + submitted.json_payload().unwrap(), + json!({ + "version": 1, + "resolve_strategy": "use_latest", + "fail_fast": false, + "parallelism": 8, + }) + ); + } + + #[test] + fn catalog_non_terminal_lifecycle_events_have_null_payloads() { + let events = CatalogEventHarness::all(); + let mut procedure = test_procedure(); + procedure.context.volatile_ctx.metrics = populated_metrics(); + + for trigger in [ + EventTrigger::Recovered, + EventTrigger::ChildSubmitted { + procedure_id: ProcedureId::random(), + outcome: ChildSubmissionOutcome::Accepted, + }, + EventTrigger::Retrying { + phase: RetryPhase::Execute, + attempt: 1, + }, + EventTrigger::RollingBack, + ] { + assert_eq!( + events + .event(&procedure, trigger) + .unwrap() + .json_payload() + .unwrap(), + Value::Null + ); + } + } + + #[test] + fn catalog_terminal_events_report_existing_metrics() { + let events = CatalogEventHarness::all(); + let mut procedure = test_procedure(); + procedure.context.volatile_ctx.metrics = populated_metrics(); + + for (trigger, complete) in [ + (EventTrigger::Succeeded, true), + (EventTrigger::Failed, false), + (EventTrigger::Poisoned, false), + ] { + assert_eq!( + events + .event(&procedure, trigger) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": complete, + "processed_database_count": 4, + "succeeded_database_count": 3, + "failed_database_count": 1, + }) + ); + } + } + + #[test] + fn catalog_event_filtering_uses_the_catalog_event_type() { + let procedure = test_procedure(); + assert!( + CatalogEventHarness::selected([RECONCILE_CATALOG_EVENT_TYPE]) + .event(&procedure, EventTrigger::Submitted) + .is_some() + ); + assert!( + CatalogEventHarness::selected([RECONCILE_DATABASE_EVENT_TYPE]) + .event(&procedure, EventTrigger::Submitted) + .is_none() + ); + assert!( + CatalogEventHarness::selected([]) + .event(&procedure, EventTrigger::Submitted) + .is_none() + ); + } + + #[test] + fn catalog_recovery_preserves_locator_and_resets_metrics() { + let events = CatalogEventHarness::all(); + let mut procedure = test_procedure(); + procedure.context.volatile_ctx.metrics = populated_metrics(); + let original_dump = procedure.dump().unwrap(); + procedure.context.volatile_ctx.metrics = ReconcileCatalogMetrics::default(); + assert_eq!(procedure.dump().unwrap(), original_dump); + + let loaded = ReconcileCatalogProcedure::from_json(test_context(), &original_dump).unwrap(); + assert_eq!(loaded.dump().unwrap(), original_dump); + assert_eq!( + events + .event(&loaded, EventTrigger::Recovered) + .unwrap() + .extra_rows() + .unwrap(), + events + .event(&procedure, EventTrigger::Submitted) + .unwrap() + .extra_rows() + .unwrap(), + ); + assert_eq!( + events + .event(&loaded, EventTrigger::Succeeded) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": true, + "processed_database_count": 0, + "succeeded_database_count": 0, + "failed_database_count": 0, + }) + ); + } + + fn populated_metrics() -> ReconcileCatalogMetrics { + ReconcileCatalogMetrics { + succeeded_databases: 3, + failed_databases: 1, + } + } + + fn test_procedure() -> ReconcileCatalogProcedure { + ReconcileCatalogProcedure::new( + test_context(), + "greptime".to_string(), + false, + ResolveStrategy::UseLatest, + 8, + ) + } + + fn test_context() -> Context { + let ddl_context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + Context { + node_manager: ddl_context.node_manager, + table_metadata_manager: ddl_context.table_metadata_manager, + cache_invalidator: ddl_context.cache_invalidator, + } + } +} diff --git a/src/common/meta/src/reconciliation/reconcile_database.rs b/src/common/meta/src/reconciliation/reconcile_database.rs index 7d4262877c1..f5e2fe9f9a0 100644 --- a/src/common/meta/src/reconciliation/reconcile_database.rs +++ b/src/common/meta/src/reconciliation/reconcile_database.rs @@ -25,8 +25,8 @@ use std::time::Instant; use async_trait::async_trait; use common_procedure::error::{FromJsonSnafu, ToJsonSnafu}; use common_procedure::{ - Context as ProcedureContext, Error as ProcedureError, LockKey, Procedure, - Result as ProcedureResult, Status, + Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey, + Procedure, Result as ProcedureResult, Status, }; use futures::stream::BoxStream; use serde::{Deserialize, Serialize}; @@ -41,6 +41,9 @@ use crate::key::table_name::TableNameValue; use crate::lock_key::{CatalogLock, SchemaLock}; use crate::metrics; use crate::node_manager::NodeManagerRef; +use crate::reconciliation::event::{ + RECONCILE_DATABASE_EVENT_TYPE, ReconcileDatabaseEvent, ReconciliationLocator, +}; use crate::reconciliation::reconcile_database::start::ReconcileDatabaseStart; use crate::reconciliation::reconcile_table::resolve_column_metadata::ResolveStrategy; use crate::reconciliation::utils::{ @@ -264,6 +267,47 @@ impl Procedure for ReconcileDatabaseProcedure { SchemaLock::write(catalog, schema).into(), ]) } + + fn event(&self, ctx: &EventContext<'_>) -> Option> { + if !ctx.event_type_filter.allows(RECONCILE_DATABASE_EVENT_TYPE) { + return None; + } + + let persistent_ctx = &self.context.persistent_ctx; + let locator = + ReconciliationLocator::database(&persistent_ctx.catalog, &persistent_ctx.schema); + let event = match ctx.trigger { + EventTrigger::Submitted => ReconcileDatabaseEvent::submitted( + locator, + persistent_ctx.resolve_strategy, + persistent_ctx.fail_fast, + persistent_ctx.parallelism, + persistent_ctx.is_subprocedure, + ), + EventTrigger::Succeeded => self.result_event(locator, true), + EventTrigger::Failed | EventTrigger::Poisoned => self.result_event(locator, false), + _ => ReconcileDatabaseEvent::lifecycle(locator), + }; + Some(Box::new(event)) + } +} + +impl ReconcileDatabaseProcedure { + fn result_event( + &self, + locator: ReconciliationLocator, + complete: bool, + ) -> ReconcileDatabaseEvent { + let metrics = self.context.metrics(); + ReconcileDatabaseEvent::result( + locator, + complete, + metrics.succeeded_tables, + metrics.failed_tables, + metrics.succeeded_procedures, + metrics.failed_procedures, + ) + } } #[async_trait::async_trait] @@ -283,3 +327,229 @@ pub(crate) trait State: Sync + Send + Debug { fn as_any(&self) -> &dyn Any; } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use common_event_recorder::{EventTypeFilter, EventTypeFilterRef}; + use common_procedure::{ + ChildSubmissionOutcome, EventContext, EventTrigger, Procedure, ProcedureId, ProcedureState, + RetryPhase, + }; + use serde_json::{Value, json}; + + use super::*; + use crate::reconciliation::event::RECONCILE_CATALOG_EVENT_TYPE; + use crate::test_util::{MockDatanodeManager, new_ddl_context}; + + struct DatabaseEventHarness { + procedure_id: ProcedureId, + lifecycle_state: ProcedureState, + event_type_filter: EventTypeFilterRef, + } + + impl DatabaseEventHarness { + fn all() -> Self { + Self { + procedure_id: ProcedureId::random(), + lifecycle_state: ProcedureState::Running, + event_type_filter: Arc::new(EventTypeFilter::All), + } + } + + fn selected(event_types: impl IntoIterator) -> Self { + Self { + event_type_filter: Arc::new(EventTypeFilter::Only( + event_types.into_iter().map(str::to_string).collect(), + )), + ..Self::all() + } + } + + fn event( + &self, + procedure: &dyn Procedure, + trigger: EventTrigger, + ) -> Option> { + procedure.event(&EventContext { + procedure_id: self.procedure_id, + lifecycle_state: &self.lifecycle_state, + trigger, + event_type_filter: self.event_type_filter.clone(), + event_context: None, + }) + } + } + + #[test] + fn database_submitted_events_cover_root_and_child_intent() { + let events = DatabaseEventHarness::all(); + let root = test_procedure(false); + let child = test_procedure(true); + for (procedure, is_subprocedure) in [(&root, false), (&child, true)] { + let submitted = events.event(procedure, EventTrigger::Submitted).unwrap(); + assert_eq!(submitted.event_type(), RECONCILE_DATABASE_EVENT_TYPE); + assert_eq!( + submitted.json_payload().unwrap(), + json!({ + "version": 1, + "resolve_strategy": "use_metasrv", + "fail_fast": false, + "parallelism": 8, + "is_subprocedure": is_subprocedure, + }) + ); + } + } + + #[test] + fn database_non_terminal_lifecycle_events_have_null_payloads() { + let events = DatabaseEventHarness::all(); + let mut procedure = test_procedure(true); + procedure.context.volatile_ctx.metrics = populated_metrics(); + + for trigger in [ + EventTrigger::Recovered, + EventTrigger::ChildSubmitted { + procedure_id: ProcedureId::random(), + outcome: ChildSubmissionOutcome::Accepted, + }, + EventTrigger::Retrying { + phase: RetryPhase::Execute, + attempt: 2, + }, + EventTrigger::RollingBack, + ] { + assert_eq!( + events + .event(&procedure, trigger) + .unwrap() + .json_payload() + .unwrap(), + Value::Null + ); + } + } + + #[test] + fn database_terminal_events_report_existing_metrics() { + let events = DatabaseEventHarness::all(); + let mut procedure = test_procedure(true); + procedure.context.volatile_ctx.metrics = populated_metrics(); + + for (trigger, complete) in [ + (EventTrigger::Succeeded, true), + (EventTrigger::Failed, false), + (EventTrigger::Poisoned, false), + ] { + assert_eq!( + events + .event(&procedure, trigger) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": complete, + "processed_table_count": 8, + "succeeded_table_count": 5, + "failed_table_count": 3, + "succeeded_subprocedure_count": 3, + "failed_subprocedure_count": 2, + }) + ); + } + } + + #[test] + fn database_event_filtering_uses_the_database_event_type() { + let procedure = test_procedure(false); + assert!( + DatabaseEventHarness::selected([RECONCILE_DATABASE_EVENT_TYPE]) + .event(&procedure, EventTrigger::Submitted) + .is_some() + ); + assert!( + DatabaseEventHarness::selected([RECONCILE_CATALOG_EVENT_TYPE]) + .event(&procedure, EventTrigger::Submitted) + .is_none() + ); + assert!( + DatabaseEventHarness::selected([]) + .event(&procedure, EventTrigger::Submitted) + .is_none() + ); + } + + #[test] + fn database_recovery_preserves_locator_and_resets_metrics() { + let events = DatabaseEventHarness::all(); + let mut procedure = test_procedure(false); + procedure.context.volatile_ctx.metrics = populated_metrics(); + let original_dump = procedure.dump().unwrap(); + procedure.context.volatile_ctx.metrics = ReconcileDatabaseMetrics::default(); + assert_eq!(procedure.dump().unwrap(), original_dump); + + let loaded = ReconcileDatabaseProcedure::from_json(test_context(), &original_dump).unwrap(); + assert_eq!(loaded.dump().unwrap(), original_dump); + assert_eq!( + events + .event(&loaded, EventTrigger::Recovered) + .unwrap() + .extra_rows() + .unwrap(), + events + .event(&procedure, EventTrigger::Submitted) + .unwrap() + .extra_rows() + .unwrap(), + ); + assert_eq!( + events + .event(&loaded, EventTrigger::Succeeded) + .unwrap() + .json_payload() + .unwrap(), + json!({ + "version": 1, + "complete": true, + "processed_table_count": 0, + "succeeded_table_count": 0, + "failed_table_count": 0, + "succeeded_subprocedure_count": 0, + "failed_subprocedure_count": 0, + }) + ); + } + + fn populated_metrics() -> ReconcileDatabaseMetrics { + ReconcileDatabaseMetrics { + succeeded_tables: 5, + failed_tables: 3, + succeeded_procedures: 3, + failed_procedures: 2, + } + } + + fn test_procedure(is_subprocedure: bool) -> ReconcileDatabaseProcedure { + ReconcileDatabaseProcedure::new( + test_context(), + "greptime".to_string(), + "public".to_string(), + false, + 8, + ResolveStrategy::UseMetasrv, + is_subprocedure, + ) + } + + fn test_context() -> Context { + let ddl_context = new_ddl_context(Arc::new(MockDatanodeManager::new(()))); + Context { + node_manager: ddl_context.node_manager, + table_metadata_manager: ddl_context.table_metadata_manager, + cache_invalidator: ddl_context.cache_invalidator, + } + } +} diff --git a/src/common/pprof/Cargo.toml b/src/common/pprof/Cargo.toml index 6d96b26d9fd..e3a5bce78ea 100644 --- a/src/common/pprof/Cargo.toml +++ b/src/common/pprof/Cargo.toml @@ -11,11 +11,17 @@ snafu.workspace = true tokio.workspace = true [target.'cfg(unix)'.dependencies] -pprof = { version = "0.14", features = [ +pprof = { version = "0.15", features = [ "flamegraph", "prost-codec", "protobuf", ] } +# framehop-unwinder is only available on x86_64/aarch64 Linux/macOS. +# Enabling it on other targets (e.g. riscv64) would disable the +# backtrace-rs fallback without providing a TraceImpl, breaking the build. +[target.'cfg(all(unix, any(target_arch = "x86_64", target_arch = "aarch64"), any(target_os = "linux", target_os = "macos")))'.dependencies] +pprof = { version = "0.15", features = ["framehop-unwinder"] } + [lints] workspace = true diff --git a/src/common/telemetry/src/logging/file_retention.rs b/src/common/telemetry/src/logging/file_retention.rs index 49cc976db36..6ea9846be1a 100644 --- a/src/common/telemetry/src/logging/file_retention.rs +++ b/src/common/telemetry/src/logging/file_retention.rs @@ -453,7 +453,9 @@ pub(crate) fn build_file_appender( #[cfg(test)] mod tests { - use std::fs::{self, File}; + #[cfg(unix)] + use std::fs; + use std::fs::File; use std::io::Write; #[cfg(unix)] use std::os::unix::fs::symlink; @@ -464,8 +466,11 @@ mod tests { use tempfile::TempDir; use super::FileIndex; + #[cfg(unix)] use crate::logging::LoggingOptions; - use crate::logging::file_retention::{DirectoryRetention, LogFileKind, build_file_appender}; + #[cfg(unix)] + use crate::logging::file_retention::build_file_appender; + use crate::logging::file_retention::{DirectoryRetention, LogFileKind}; fn write_file(path: &Path, contents: &[u8]) { let mut file = File::create(path).unwrap(); diff --git a/src/datanode/src/heartbeat.rs b/src/datanode/src/heartbeat.rs index 0e0c51624e3..f8d92212728 100644 --- a/src/datanode/src/heartbeat.rs +++ b/src/datanode/src/heartbeat.rs @@ -18,6 +18,7 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use api::v1::meta::heartbeat_request::NodeWorkloads; +use api::v1::meta::mailbox_message::Payload; use api::v1::meta::{DatanodeWorkloads, HeartbeatRequest, NodeInfo, Peer, RegionRole, RegionStat}; use common_base::Plugins; use common_meta::cache_invalidator::CacheInvalidatorRef; @@ -146,7 +147,15 @@ impl HeartbeatTask { None }) { if let Some(msg) = res.mailbox_message.as_ref() { - info!("Received mailbox message: {msg:?}, meta_client id: {client_id:?}"); + let payload_len = msg + .payload + .as_ref() + .map(|Payload::Json(payload)| payload.len()) + .unwrap_or_default(); + info!( + message_id = msg.id, + payload_len, client_id, "Received mailbox message" + ); } if let Some(lease) = res.region_lease.as_ref() { metrics::LAST_RECEIVED_HEARTBEAT_ELAPSED @@ -193,7 +202,8 @@ impl HeartbeatTask { ctx: HeartbeatResponseHandlerContext, handler_executor: HeartbeatResponseHandlerExecutorRef, ) -> Result<()> { - trace!("Heartbeat response: {:?}", ctx.response); + let mailbox_message_id = ctx.response.mailbox_message.as_ref().map(|msg| msg.id); + trace!(?mailbox_message_id, "Handling heartbeat response"); handler_executor .handle(ctx) .await @@ -379,7 +389,8 @@ impl HeartbeatTask { .set(last_sent.elapsed().as_millis() as i64); // Resets the timer. last_sent = Instant::now(); - debug!("Sending heartbeat request: {:?}", req); + let mailbox_message_id = req.mailbox_message.as_ref().map(|msg| msg.id); + debug!(?mailbox_message_id, "Sending heartbeat request"); if let Err(e) = tx.send(req).await { error!(e; "Failed to send heartbeat to metasrv"); match Self::create_streams( diff --git a/src/datanode/src/heartbeat/handler.rs b/src/datanode/src/heartbeat/handler.rs index 71797592c3b..2d18d56004b 100644 --- a/src/datanode/src/heartbeat/handler.rs +++ b/src/datanode/src/heartbeat/handler.rs @@ -41,9 +41,9 @@ use crate::heartbeat::handler::apply_staging_manifest::ApplyStagingManifestsHand use crate::heartbeat::handler::close_region::CloseRegionsHandler; use crate::heartbeat::handler::downgrade_region::DowngradeRegionsHandler; use crate::heartbeat::handler::enter_staging::EnterStagingRegionsHandler; -use crate::heartbeat::handler::file_ref::GetFileRefsHandler; +use crate::heartbeat::handler::file_ref::{GetFileRefsHandler, GetPackedFileRefsHandler}; use crate::heartbeat::handler::flush_region::FlushRegionsHandler; -use crate::heartbeat::handler::gc_worker::GcRegionsHandler; +use crate::heartbeat::handler::gc_worker::{GcRegionsHandler, PackedGcRegionsHandler}; use crate::heartbeat::handler::open_region::OpenRegionsHandler; use crate::heartbeat::handler::remap_manifest::RemapManifestHandler; use crate::heartbeat::handler::sync_region::SyncRegionHandler; @@ -135,7 +135,11 @@ impl RegionHeartbeatResponseHandler { .into(), ))), Instruction::GetFileRefs(_) => Ok(Some(Box::new(GetFileRefsHandler.into()))), + Instruction::GetPackedFileRefs(_) => { + Ok(Some(Box::new(GetPackedFileRefsHandler.into()))) + } Instruction::GcRegions(_) => Ok(Some(Box::new(GcRegionsHandler.into()))), + Instruction::PackedGcRegions(_) => Ok(Some(Box::new(PackedGcRegionsHandler.into()))), Instruction::InvalidateCaches(_) => InvalidHeartbeatResponseSnafu.fail(), Instruction::Suspend => Ok(None), Instruction::EnterStagingRegions(_) => { @@ -159,7 +163,9 @@ pub enum InstructionHandlers { DowngradeRegions(DowngradeRegionsHandler), UpgradeRegions(UpgradeRegionsHandler), GetFileRefs(GetFileRefsHandler), + GetPackedFileRefs(GetPackedFileRefsHandler), GcRegions(GcRegionsHandler), + PackedGcRegions(PackedGcRegionsHandler), EnterStagingRegions(EnterStagingRegionsHandler), SyncRegions(SyncRegionHandler), RemapManifest(RemapManifestHandler), @@ -192,7 +198,9 @@ impl_from_handler!( DowngradeRegionsHandler => DowngradeRegions, UpgradeRegionsHandler => UpgradeRegions, GetFileRefsHandler => GetFileRefs, + GetPackedFileRefsHandler => GetPackedFileRefs, GcRegionsHandler => GcRegions, + PackedGcRegionsHandler => PackedGcRegions, EnterStagingRegionsHandler => EnterStagingRegions, SyncRegionHandler => SyncRegions, RemapManifestHandler => RemapManifest, @@ -240,7 +248,9 @@ dispatch_instr!( DowngradeRegions => DowngradeRegions, UpgradeRegions => UpgradeRegions, GetFileRefs => GetFileRefs, + GetPackedFileRefs => GetPackedFileRefs, GcRegions => GcRegions, + PackedGcRegions => PackedGcRegions, EnterStagingRegions => EnterStagingRegions, SyncRegions => SyncRegions, RemapManifest => RemapManifest, diff --git a/src/datanode/src/heartbeat/handler/file_ref.rs b/src/datanode/src/heartbeat/handler/file_ref.rs index ef9cbd4158e..554f4119b05 100644 --- a/src/datanode/src/heartbeat/handler/file_ref.rs +++ b/src/datanode/src/heartbeat/handler/file_ref.rs @@ -13,7 +13,10 @@ // limitations under the License. use common_error::ext::ErrorExt; -use common_meta::instruction::{GetFileRefs, GetFileRefsReply, InstructionError, InstructionReply}; +use common_meta::instruction::{ + GetFileRefs, GetFileRefsReply, GetPackedFileRefs, GetPackedFileRefsReply, InstructionError, + InstructionReply, PackedFileRefsManifest, +}; use store_api::storage::FileRefsManifest; use crate::heartbeat::handler::{HandlerContext, InstructionHandler}; @@ -65,3 +68,54 @@ impl InstructionHandler for GetFileRefsHandler { } } } + +pub struct GetPackedFileRefsHandler; + +#[async_trait::async_trait] +impl InstructionHandler for GetPackedFileRefsHandler { + type Instruction = GetPackedFileRefs; + + async fn handle( + &self, + ctx: &HandlerContext, + get_file_refs: Self::Instruction, + ) -> Option { + let region_server = &ctx.region_server; + let Some(mito_engine) = region_server.mito_engine() else { + return Some(InstructionReply::GetPackedFileRefs( + GetPackedFileRefsReply { + packed_file_refs_manifest: PackedFileRefsManifest::default(), + success: false, + error: Some(InstructionError::legacy_internal_retryable( + "MitoEngine not found", + )), + }, + )); + }; + match mito_engine + .get_snapshot_of_file_refs(get_file_refs.query_regions, get_file_refs.related_regions) + .await + { + Ok(all_file_refs) => Some(InstructionReply::GetPackedFileRefs( + GetPackedFileRefsReply { + packed_file_refs_manifest: PackedFileRefsManifest::from_manifest( + &all_file_refs, + ), + success: true, + error: None, + }, + )), + Err(e) => Some(InstructionReply::GetPackedFileRefs( + GetPackedFileRefsReply { + packed_file_refs_manifest: PackedFileRefsManifest::default(), + success: false, + error: Some(InstructionError::new( + e.status_code(), + format!("Failed to get file refs: {}", e.output_msg()), + e.retry_hint(), + )), + }, + )), + } + } +} diff --git a/src/datanode/src/heartbeat/handler/gc_worker.rs b/src/datanode/src/heartbeat/handler/gc_worker.rs index b32a1c29ad4..e63d21b6e5b 100644 --- a/src/datanode/src/heartbeat/handler/gc_worker.rs +++ b/src/datanode/src/heartbeat/handler/gc_worker.rs @@ -15,7 +15,7 @@ use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; -use common_meta::instruction::{GcRegions, GcRegionsReply, InstructionReply}; +use common_meta::instruction::{GcRegions, GcRegionsReply, InstructionReply, PackedGcRegions}; use common_meta::key::table_info::TableInfoManager; use common_meta::key::table_route::TableRouteManager; use common_telemetry::{debug, warn}; @@ -34,6 +34,7 @@ use crate::error::{GcMitoEngineSnafu, GetMetadataSnafu, Result, UnexpectedSnafu} use crate::heartbeat::handler::{HandlerContext, InstructionHandler}; pub struct GcRegionsHandler; +pub struct PackedGcRegionsHandler; #[async_trait::async_trait] impl InstructionHandler for GcRegionsHandler { @@ -67,7 +68,7 @@ impl InstructionHandler for GcRegionsHandler { .push(rid); } - let file_refs_manifest = gc_regions.file_refs_manifest.clone(); + let file_refs_manifest = gc_regions.file_refs_manifest; let full_file_listing = gc_regions.full_file_listing; let ctx_clone = ctx.clone(); @@ -138,6 +139,36 @@ impl InstructionHandler for GcRegionsHandler { } } +#[async_trait::async_trait] +impl InstructionHandler for PackedGcRegionsHandler { + type Instruction = PackedGcRegions; + + async fn handle( + &self, + ctx: &HandlerContext, + instruction: Self::Instruction, + ) -> Option { + let manifest = match instruction.packed_file_refs_manifest.into_manifest() { + Ok(manifest) => manifest, + Err(err) => { + return Some(InstructionReply::GcRegions(GcRegionsReply { + result: Err(common_meta::instruction::InstructionError::from_error(&err)), + })); + } + }; + GcRegionsHandler + .handle( + ctx, + GcRegions { + regions: instruction.regions, + file_refs_manifest: manifest, + full_file_listing: instruction.full_file_listing, + }, + ) + .await + } +} + impl GcRegionsHandler { /// Create a GC worker for the given table and region IDs. async fn create_gc_worker( diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 79d0dd19152..10bcf32b13d 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -2018,6 +2018,7 @@ mod tests { assert!(RegionServerInner::is_ingest_request(&RegionRequest::Put( RegionPutRequest { + skip_wal: false, rows: rows(), hint: None, partition_expr_version: None, diff --git a/src/datatypes/src/schema/constraint.rs b/src/datatypes/src/schema/constraint.rs index 60025be68e4..d946a92bec5 100644 --- a/src/datatypes/src/schema/constraint.rs +++ b/src/datatypes/src/schema/constraint.rs @@ -23,7 +23,7 @@ use crate::error::{self, Result}; use crate::types::cast; use crate::value::Value; use crate::vectors::operations::VectorOp; -use crate::vectors::{TimestampMillisecondVector, VectorRef}; +use crate::vectors::{Helper, TimestampMillisecondVector, VectorRef}; pub const CURRENT_TIMESTAMP: &str = "current_timestamp"; pub const CURRENT_TIMESTAMP_FN: &str = "current_timestamp()"; @@ -151,15 +151,20 @@ impl ColumnDefaultConstraint { ColumnDefaultConstraint::Value(v) => { ensure!(is_nullable || !v.is_null(), error::NullDefaultSnafu); - // TODO(yingwen): - // 1. For null value, we could use NullVector once it supports custom logical type. - // 2. For non null value, we could use ConstantVector, but it would cause all codes - // attempt to downcast the vector fail if they don't check whether the vector is const - // first. - let mut mutable_vector = data_type.create_mutable_vector(1); - mutable_vector.try_push_value_ref(&v.as_value_ref())?; - let base_vector = mutable_vector.to_vector(); - Ok(base_vector.replicate(&[num_rows])) + if let Ok(vector) = v.try_to_scalar_value(data_type).and_then(|scalar| { + Helper::try_from_scalar_value(scalar, num_rows, Some(data_type)) + }) { + return Ok(vector); + } + + // Some extension values, such as JSON nested in a struct, cannot safely + // round-trip through ScalarValue. Preserve their logical type with the + // type-specific vector builder instead. + let mut mutable_vector = data_type.create_mutable_vector(num_rows); + for _ in 0..num_rows { + mutable_vector.try_push_value_ref(&v.as_value_ref())?; + } + Ok(mutable_vector.to_vector()) } } } @@ -391,6 +396,82 @@ mod tests { assert_eq!(Value::Int32(10), v); } + #[test] + fn test_struct_default_null_fields_and_json() { + use crate::types::{StructField, StructType}; + use crate::value::StructValue; + + let inner_type = StructType::from([StructField::new( + "x", + ConcreteDataType::int32_datatype(), + true, + )]); + let inner = Value::Struct(StructValue::new(vec![Value::Null], inner_type)); + let json = crate::json::JsonSettings::default() + .encode(serde_json::json!({"answer": 42})) + .unwrap(); + let nested_type = StructType::from([StructField::new("nested", inner.data_type(), true)]); + let json_type = StructType::from([StructField::new("json", json.data_type(), true)]); + let values = [ + inner.clone(), + Value::Struct(StructValue::new(vec![inner], nested_type)), + Value::Struct(StructValue::new(vec![], StructType::default())), + Value::Struct(StructValue::new(vec![json], json_type)), + ]; + for value in values { + let data_type = value.data_type(); + // JSON children are read back as their underlying struct values. + let expected = serde_json::Value::try_from(value.clone()).unwrap(); + for num_rows in [1, 3] { + let vector = ColumnDefaultConstraint::Value(value.clone()) + .create_default_vector(&data_type, false, num_rows) + .unwrap(); + assert_eq!(data_type, vector.data_type()); + assert_eq!( + data_type.as_arrow_type(), + *vector.to_arrow_array().data_type() + ); + assert_eq!(num_rows, vector.len()); + assert_eq!(0, vector.null_count()); + for row in 0..num_rows { + assert_eq!( + expected, + serde_json::Value::try_from(vector.get(row)).unwrap() + ); + } + } + } + } + + #[test] + fn test_string_default_preserves_batch_schema() { + use arrow::datatypes::{Field, Schema}; + use arrow::record_batch::RecordBatch; + + for data_type in [ + ConcreteDataType::large_string_datatype(), + ConcreteDataType::utf8_view_datatype(), + ] { + let schema = Arc::new(Schema::new(vec![Field::new( + "tag", + data_type.as_arrow_type(), + true, + )])); + for value in [Value::from("greptime"), Value::Null] { + let vector = ColumnDefaultConstraint::Value(value.clone()) + .create_default_vector(&data_type, true, 3) + .unwrap(); + let batch = + RecordBatch::try_new(schema.clone(), vec![vector.to_arrow_array()]).unwrap(); + assert_eq!(3, batch.num_rows()); + assert_eq!(data_type, vector.data_type()); + for row in 0..batch.num_rows() { + assert_eq!(value, vector.get(row)); + } + } + } + } + #[test] fn test_create_default_vector_by_func() { let constraint = ColumnDefaultConstraint::Function(CURRENT_TIMESTAMP.to_string()); diff --git a/src/datatypes/src/types/cast.rs b/src/datatypes/src/types/cast.rs index d7a27a45e10..2f92713c26f 100644 --- a/src/datatypes/src/types/cast.rs +++ b/src/datatypes/src/types/cast.rs @@ -26,7 +26,7 @@ pub fn cast(src_value: Value, dest_type: &ConcreteDataType) -> Result { } let src_type = src_value.data_type(); let scalar_value = src_value.try_to_scalar_value(&src_type)?; - let new_value = Helper::try_from_scalar_value(scalar_value, 1)? + let new_value = Helper::try_from_scalar_value(scalar_value, 1, None)? .cast(dest_type)? .get(0); Ok(new_value) diff --git a/src/datatypes/src/value.rs b/src/datatypes/src/value.rs index 9ad9e1d46ef..35a6acabd33 100644 --- a/src/datatypes/src/value.rs +++ b/src/datatypes/src/value.rs @@ -1085,11 +1085,21 @@ impl StructValue { } fn try_to_scalar_value(&self, output_type: &StructType) -> Result { + let output_fields = output_type.fields(); + ensure!( + self.items.len() == output_fields.len(), + InconsistentStructFieldsAndItemsSnafu { + field_len: output_fields.len(), + item_len: self.items.len() + } + ); let arrays = self .items .iter() - .map(|value| { - let scalar_value = value.try_to_scalar_value(&value.data_type())?; + .zip(output_fields.iter()) + .map(|(value, field)| { + // Null values need the declared field type to produce a typed null array. + let scalar_value = value.try_to_scalar_value(field.data_type())?; scalar_value .to_array() .context(ConvertScalarToArrowArraySnafu) @@ -1097,7 +1107,13 @@ impl StructValue { .collect::>>>()?; let fields = output_type.as_arrow_fields(); - let struct_array = StructArray::new(fields, arrays, None); + let struct_array = + StructArray::try_new_with_length(fields, arrays, None, 1).map_err(|error| { + error::ToScalarValueSnafu { + reason: error.to_string(), + } + .build() + })?; Ok(ScalarValue::Struct(Arc::new(struct_array))) } } @@ -3082,6 +3098,66 @@ pub(crate) mod tests { } } + #[test] + fn test_struct_scalar_null_fields() { + let struct_type = StructType::from([ + StructField::new("x", ConcreteDataType::int32_datatype(), true), + StructField::new("name", ConcreteDataType::string_datatype(), true), + ]); + let value = StructValue::new(vec![Value::Null, Value::from("hello")], struct_type.clone()); + let ScalarValue::Struct(array) = value.try_to_scalar_value(&struct_type).unwrap() else { + panic!("Expected struct scalar"); + }; + assert_eq!(1, array.len()); + assert_eq!(0, array.null_count()); + assert_eq!(&struct_type.as_arrow_fields(), array.fields()); + assert_eq!( + ScalarValue::Int32(None), + ScalarValue::try_from_array(array.column(0), 0).unwrap() + ); + + let nested_type = StructType::from([StructField::new( + "nested", + ConcreteDataType::struct_datatype(struct_type), + true, + )]); + for child in [Value::Struct(value), Value::Null] { + let nested = StructValue::new(vec![child.clone()], nested_type.clone()); + let ScalarValue::Struct(array) = nested.try_to_scalar_value(&nested_type).unwrap() + else { + panic!("Expected struct scalar"); + }; + let vector = crate::vectors::Helper::try_into_vector(array.column(0).clone()).unwrap(); + assert_eq!(child, vector.get(0)); + } + + let empty_type = StructType::default(); + let empty = StructValue::new(vec![], empty_type.clone()); + let ScalarValue::Struct(array) = empty.try_to_scalar_value(&empty_type).unwrap() else { + panic!("Expected struct scalar"); + }; + assert_eq!(1, array.len()); + assert_eq!(0, array.num_columns()); + assert_eq!(0, array.null_count()); + } + + #[test] + fn test_struct_scalar_invalid_fields() { + let struct_type = StructType::from([StructField::new( + "x", + ConcreteDataType::int32_datatype(), + false, + )]); + for child in [Value::Null, Value::from("wrong type")] { + let value = StructValue::new(vec![child], struct_type.clone()); + assert!(value.try_to_scalar_value(&struct_type).is_err()); + } + let value = StructValue::new(vec![Value::Int32(1)], struct_type.clone()); + assert!(value.try_to_scalar_value(&StructType::default()).is_err()); + let empty = StructValue::new(vec![], StructType::default()); + assert!(empty.try_to_scalar_value(&struct_type).is_err()); + } + #[test] fn test_timestamp_to_scalar_value() { assert_eq!( diff --git a/src/datatypes/src/vectors.rs b/src/datatypes/src/vectors.rs index 593bec833da..38db5940df9 100644 --- a/src/datatypes/src/vectors.rs +++ b/src/datatypes/src/vectors.rs @@ -27,7 +27,6 @@ use crate::vectors::operations::VectorOp; mod binary; mod boolean; -mod constant; mod date; mod decimal; mod dictionary; @@ -48,7 +47,6 @@ mod validity; pub use binary::{BinaryVector, BinaryVectorBuilder}; pub use boolean::{BooleanVector, BooleanVectorBuilder}; -pub use constant::ConstantVector; pub use date::{DateVector, DateVectorBuilder}; pub use decimal::{Decimal128Vector, Decimal128VectorBuilder}; pub(crate) use dictionary::StringDictionaryVectorBuilder; @@ -126,11 +124,6 @@ pub trait Vector: Send + Sync + Serializable + Debug + VectorOp { /// This is `O(1)`. fn null_count(&self) -> usize; - /// Returns true when it's a ConstantColumn - fn is_const(&self) -> bool { - false - } - /// Returns whether row is null. fn is_null(&self, row: usize) -> bool; diff --git a/src/datatypes/src/vectors/binary.rs b/src/datatypes/src/vectors/binary.rs index 6d9954fb80f..77344b46512 100644 --- a/src/datatypes/src/vectors/binary.rs +++ b/src/datatypes/src/vectors/binary.rs @@ -479,7 +479,6 @@ mod tests { assert_eq!(2, v.len()); assert_eq!("BinaryVector", v.vector_type_name()); - assert!(!v.is_const()); assert!(v.validity().is_all_valid()); assert!(!v.only_null()); assert_eq!(128, v.memory_size()); diff --git a/src/datatypes/src/vectors/boolean.rs b/src/datatypes/src/vectors/boolean.rs index 4ea7a4591b9..48625da516b 100644 --- a/src/datatypes/src/vectors/boolean.rs +++ b/src/datatypes/src/vectors/boolean.rs @@ -40,10 +40,6 @@ impl BooleanVector { pub fn as_boolean_array(&self) -> &BooleanArray { &self.array } - - pub(crate) fn false_count(&self) -> usize { - self.array.false_count() - } } impl From> for BooleanVector { @@ -251,7 +247,6 @@ mod tests { let v = BooleanVector::from(bools.clone()); assert_eq!(9, v.len()); assert_eq!("BooleanVector", v.vector_type_name()); - assert!(!v.is_const()); assert!(v.validity().is_all_valid()); assert!(!v.only_null()); assert_eq!(2, v.memory_size()); diff --git a/src/datatypes/src/vectors/constant.rs b/src/datatypes/src/vectors/constant.rs deleted file mode 100644 index 3a0adc7b62a..00000000000 --- a/src/datatypes/src/vectors/constant.rs +++ /dev/null @@ -1,252 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::any::Any; -use std::fmt; -use std::sync::Arc; - -use arrow::array::{Array, ArrayRef, UInt32Array}; -use snafu::{ResultExt, ensure}; - -use crate::data_type::ConcreteDataType; -use crate::error::{self, Result, SerializeSnafu}; -use crate::serialize::Serializable; -use crate::value::{Value, ValueRef}; -use crate::vectors::{BooleanVector, Helper, UInt32Vector, Validity, Vector, VectorRef}; - -#[derive(Clone)] -pub struct ConstantVector { - length: usize, - vector: VectorRef, -} - -impl ConstantVector { - /// Create a new [ConstantVector]. - /// - /// # Panics - /// Panics if `vector.len() != 1`. - pub fn new(vector: VectorRef, length: usize) -> Self { - assert_eq!(1, vector.len()); - - // Avoid const recursion. - if vector.is_const() { - let vec: &ConstantVector = unsafe { Helper::static_cast(&vector) }; - return Self::new(vec.inner().clone(), length); - } - Self { vector, length } - } - - pub fn inner(&self) -> &VectorRef { - &self.vector - } - - /// Returns the constant value. - pub fn get_constant_ref(&self) -> ValueRef<'_> { - self.vector.get_ref(0) - } - - pub(crate) fn replicate_vector(&self, offsets: &[usize]) -> VectorRef { - assert_eq!(offsets.len(), self.len()); - - if offsets.is_empty() { - return self.slice(0, 0); - } - - Arc::new(ConstantVector::new( - self.vector.clone(), - *offsets.last().unwrap(), - )) - } - - pub(crate) fn filter_vector(&self, filter: &BooleanVector) -> Result { - let length = self.len() - filter.false_count(); - if length == self.len() { - return Ok(Arc::new(self.clone())); - } - Ok(Arc::new(ConstantVector::new(self.inner().clone(), length))) - } - - pub(crate) fn cast_vector(&self, to_type: &ConcreteDataType) -> Result { - Ok(Arc::new(ConstantVector::new( - self.inner().cast(to_type)?, - self.length, - ))) - } - - pub(crate) fn take_vector(&self, indices: &UInt32Vector) -> Result { - if indices.is_empty() { - return Ok(self.slice(0, 0)); - } - ensure!( - indices.null_count() == 0, - error::UnsupportedOperationSnafu { - op: "taking a null index", - vector_type: self.vector_type_name(), - } - ); - - let len = self.len(); - let arr = indices.to_arrow_array(); - let indices_arr = arr.as_any().downcast_ref::().unwrap(); - if !arrow::compute::min_boolean( - &arrow::compute::kernels::cmp::lt(indices_arr, &UInt32Array::new_scalar(len as u32)) - .unwrap(), - ) - .unwrap() - { - panic!( - "Array index out of bounds, cannot take index out of the length of the array: {len}" - ); - } - - Ok(Arc::new(ConstantVector::new( - self.inner().clone(), - indices.len(), - ))) - } -} - -impl Vector for ConstantVector { - fn data_type(&self) -> ConcreteDataType { - self.vector.data_type() - } - - fn vector_type_name(&self) -> String { - "ConstantVector".to_string() - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn len(&self) -> usize { - self.length - } - - fn to_arrow_array(&self) -> ArrayRef { - let v = self.vector.replicate(&[self.length]); - v.to_arrow_array() - } - - fn to_boxed_arrow_array(&self) -> Box { - let v = self.vector.replicate(&[self.length]); - v.to_boxed_arrow_array() - } - - fn is_const(&self) -> bool { - true - } - - fn validity(&self) -> Validity { - if self.vector.is_null(0) { - Validity::all_null(self.length) - } else { - Validity::all_valid(self.length) - } - } - - fn memory_size(&self) -> usize { - self.vector.memory_size() - } - - fn is_null(&self, _row: usize) -> bool { - self.vector.is_null(0) - } - - fn only_null(&self) -> bool { - self.vector.is_null(0) - } - - fn slice(&self, _offset: usize, length: usize) -> VectorRef { - Arc::new(Self { - vector: self.vector.clone(), - length, - }) - } - - fn get(&self, _index: usize) -> Value { - self.vector.get(0) - } - - fn get_ref(&self, _index: usize) -> ValueRef<'_> { - self.vector.get_ref(0) - } - - fn null_count(&self) -> usize { - if self.only_null() { self.len() } else { 0 } - } -} - -impl fmt::Debug for ConstantVector { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "ConstantVector([{:?}; {}])", self.get(0), self.len()) - } -} - -impl Serializable for ConstantVector { - fn serialize_to_json(&self) -> Result> { - std::iter::repeat_n(self.get(0), self.len()) - .map(serde_json::Value::try_from) - .collect::>() - .context(SerializeSnafu) - } -} - -#[cfg(test)] -mod tests { - use arrow::datatypes::DataType as ArrowDataType; - - use super::*; - use crate::vectors::Int32Vector; - - #[test] - fn test_constant_vector_misc() { - let a = Int32Vector::from_slice(vec![1]); - let c = ConstantVector::new(Arc::new(a), 10); - - assert_eq!("ConstantVector", c.vector_type_name()); - assert!(c.is_const()); - assert_eq!(10, c.len()); - assert!(c.validity().is_all_valid()); - assert!(!c.only_null()); - assert_eq!(4, c.memory_size()); - - for i in 0..10 { - assert!(!c.is_null(i)); - assert_eq!(Value::Int32(1), c.get(i)); - } - - let arrow_arr = c.to_arrow_array(); - assert_eq!(10, arrow_arr.len()); - assert_eq!(&ArrowDataType::Int32, arrow_arr.data_type()); - } - - #[test] - fn test_debug_null_array() { - let a = Int32Vector::from_slice(vec![1]); - let c = ConstantVector::new(Arc::new(a), 10); - - let s = format!("{c:?}"); - assert_eq!(s, "ConstantVector([Int32(1); 10])"); - } - - #[test] - fn test_serialize_json() { - let a = Int32Vector::from_slice(vec![1]); - let c = ConstantVector::new(Arc::new(a), 10); - - let s = serde_json::to_string(&c.serialize_to_json().unwrap()).unwrap(); - assert_eq!(s, "[1,1,1,1,1,1,1,1,1,1]"); - } -} diff --git a/src/datatypes/src/vectors/decimal.rs b/src/datatypes/src/vectors/decimal.rs index 46fa73297a5..34d066f0725 100644 --- a/src/datatypes/src/vectors/decimal.rs +++ b/src/datatypes/src/vectors/decimal.rs @@ -385,43 +385,6 @@ impl Decimal128VectorBuilder { vectors::impl_try_from_arrow_array_for_vector!(Decimal128Array, Decimal128Vector); -pub(crate) fn replicate_decimal128( - vector: &Decimal128Vector, - offsets: &[usize], -) -> Decimal128Vector { - assert_eq!(offsets.len(), vector.len()); - - if offsets.is_empty() { - return vector.get_slice(0, 0); - } - - // Safety: safe to unwrap because we the vector ensures precision and scale are valid. - let mut builder = Decimal128VectorBuilder::with_capacity(*offsets.last().unwrap()) - .with_precision_and_scale(vector.precision(), vector.scale()) - .unwrap(); - - let mut previous_offset = 0; - - for (offset, value) in offsets.iter().zip(vector.array.iter()) { - let repeat_times = *offset - previous_offset; - match value { - Some(data) => { - unsafe { - // Safety: std::iter::Repeat and std::iter::Take implement TrustedLen. - builder - .mutable_array - .append_trusted_len_iter(std::iter::repeat_n(data, repeat_times)); - } - } - None => { - builder.mutable_array.append_nulls(repeat_times); - } - } - previous_offset = *offset; - } - builder.finish() -} - #[cfg(test)] pub mod tests { use arrow_array::Decimal128Array; diff --git a/src/datatypes/src/vectors/dictionary.rs b/src/datatypes/src/vectors/dictionary.rs index d41daf4638f..5f117b0ae88 100644 --- a/src/datatypes/src/vectors/dictionary.rs +++ b/src/datatypes/src/vectors/dictionary.rs @@ -17,8 +17,7 @@ use std::fmt; use std::sync::Arc; use arrow::array::{ - Array, ArrayBuilder, ArrayRef, DictionaryArray, PrimitiveArray, PrimitiveBuilder, - StringDictionaryBuilder, + Array, ArrayBuilder, ArrayRef, DictionaryArray, PrimitiveArray, StringDictionaryBuilder, }; use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, UInt32Type}; use serde_json::Value as JsonValue; @@ -320,36 +319,6 @@ impl<'a, K: ArrowDictionaryKeyType> Iterator for DictionaryIter<'a, K> { } impl VectorOp for DictionaryVector { - fn replicate(&self, offsets: &[usize]) -> VectorRef { - let keys = self.array.keys(); - let mut replicated_keys = PrimitiveBuilder::new(); - - let mut previous_offset = 0; - let mut key_iter = keys.iter().chain(std::iter::repeat(None)); - for &offset in offsets { - let key = key_iter.next().unwrap(); - - // repeat this key (offset - previous_offset) times - let repeat_count = offset - previous_offset; - for _ in 0..repeat_count { - replicated_keys.append_option(key); - } - - previous_offset = offset; - } - - let new_keys = replicated_keys.finish(); - let new_array = DictionaryArray::try_new(new_keys, self.values().clone()) - .expect("Failed to create replicated dictionary array"); - - Arc::new(Self { - array: new_array, - key_type: self.key_type.clone(), - item_type: self.item_type.clone(), - item_vector: self.item_vector.clone(), - }) - } - fn filter(&self, filter: &vectors::BooleanVector) -> Result { let key_array: ArrayRef = Arc::new(self.array.keys().clone()); let key_vector = Helper::try_into_vector(&key_array)?; @@ -485,21 +454,6 @@ mod tests { assert_eq!(sliced.get(2), Value::Null); } - #[test] - fn test_replicate() { - let dict_vec = create_test_dictionary(); - - // Replicate with offsets [0, 2, 5] - should get values at these indices - let offsets = vec![0, 2, 5]; - let replicated = dict_vec.replicate(&offsets); - assert_eq!(replicated.len(), 5); - assert_eq!(replicated.get(0), Value::String("b".to_string().into())); - assert_eq!(replicated.get(1), Value::String("b".to_string().into())); - assert_eq!(replicated.get(2), Value::String("c".to_string().into())); - assert_eq!(replicated.get(3), Value::String("c".to_string().into())); - assert_eq!(replicated.get(4), Value::String("c".to_string().into())); - } - #[test] fn test_filter() { let dict_vec = create_test_dictionary(); diff --git a/src/datatypes/src/vectors/eq.rs b/src/datatypes/src/vectors/eq.rs index e8e942d3ff4..937e4587c35 100644 --- a/src/datatypes/src/vectors/eq.rs +++ b/src/datatypes/src/vectors/eq.rs @@ -18,7 +18,6 @@ use common_time::interval::IntervalUnit; use crate::data_type::DataType; use crate::types::{DurationType, TimeType, TimestampType}; -use crate::vectors::constant::ConstantVector; use crate::vectors::struct_vector::StructVector; use crate::vectors::{ BinaryVector, BooleanVector, DateVector, Decimal128Vector, DurationMicrosecondVector, @@ -58,23 +57,6 @@ fn equal(lhs: &dyn Vector, rhs: &dyn Vector) -> bool { return false; } - if lhs.is_const() || rhs.is_const() { - // Length has been checked before, so we only need to compare inner - // vector here. - return equal( - &**lhs - .as_any() - .downcast_ref::() - .unwrap() - .inner(), - &**rhs - .as_any() - .downcast_ref::() - .unwrap() - .inner(), - ); - } - use crate::data_type::ConcreteDataType::*; let lhs_type = lhs.data_type(); @@ -191,10 +173,6 @@ mod tests { Some(b"world".to_vec()), ]))); assert_vector_ref_eq(Arc::new(BooleanVector::from(vec![true, false]))); - assert_vector_ref_eq(Arc::new(ConstantVector::new( - Arc::new(BooleanVector::from(vec![true])), - 5, - ))); assert_vector_ref_eq(Arc::new(BooleanVector::from(vec![true, false]))); assert_vector_ref_eq(Arc::new(DateVector::from(vec![Some(100), Some(120)]))); assert_vector_ref_eq(Arc::new(TimestampSecondVector::from_values([100, 120]))); @@ -259,22 +237,6 @@ mod tests { ]))); } - // Regression: second arm must downcast `rhs` (was `lhs`), or same-length ConstantVectors - // with different inners compare equal. - #[test] - fn test_constant_vector_eq_compares_both_inners() { - assert_vector_ref_ne( - Arc::new(ConstantVector::new( - Arc::new(BooleanVector::from(vec![true])), - 5, - )), - Arc::new(ConstantVector::new( - Arc::new(BooleanVector::from(vec![false])), - 5, - )), - ); - } - #[test] fn test_vector_ne() { assert_vector_ref_ne( @@ -289,36 +251,6 @@ mod tests { Arc::new(Int32Vector::from_slice([1, 2, 3, 4])), Arc::new(BooleanVector::from(vec![true, true])), ); - assert_vector_ref_ne( - Arc::new(ConstantVector::new( - Arc::new(BooleanVector::from(vec![true])), - 5, - )), - Arc::new(ConstantVector::new( - Arc::new(BooleanVector::from(vec![true])), - 4, - )), - ); - assert_vector_ref_ne( - Arc::new(ConstantVector::new( - Arc::new(BooleanVector::from(vec![true])), - 5, - )), - Arc::new(ConstantVector::new( - Arc::new(BooleanVector::from(vec![false])), - 4, - )), - ); - assert_vector_ref_ne( - Arc::new(ConstantVector::new( - Arc::new(BooleanVector::from(vec![true])), - 5, - )), - Arc::new(ConstantVector::new( - Arc::new(Int32Vector::from_slice(vec![1])), - 4, - )), - ); assert_vector_ref_ne(Arc::new(NullVector::new(5)), Arc::new(NullVector::new(8))); assert_vector_ref_ne( diff --git a/src/datatypes/src/vectors/helper.rs b/src/datatypes/src/vectors/helper.rs index 685b07419f9..3d9e71853cf 100644 --- a/src/datatypes/src/vectors/helper.rs +++ b/src/datatypes/src/vectors/helper.rs @@ -30,21 +30,21 @@ use datafusion_common::ScalarValue; use snafu::{OptionExt, ResultExt}; use crate::data_type::ConcreteDataType; -use crate::error::{self, ConvertArrowArrayToScalarsSnafu, Result}; +use crate::error::{self, ConvertScalarToArrowArraySnafu, Result}; use crate::prelude::DataType; -use crate::scalars::{Scalar, ScalarVectorBuilder}; +use crate::scalars::Scalar; use crate::types::StructType; -use crate::value::{ListValue, ListValueRef, Value}; +use crate::value::Value; use crate::vectors::struct_vector::StructVector; use crate::vectors::{ - BinaryVector, BooleanVector, ConstantVector, DateVector, Decimal128Vector, DictionaryVector, + BinaryVector, BooleanVector, DateVector, Decimal128Vector, DictionaryVector, DurationMicrosecondVector, DurationMillisecondVector, DurationNanosecondVector, DurationSecondVector, Float32Vector, Float64Vector, Int8Vector, Int16Vector, Int32Vector, Int64Vector, IntervalDayTimeVector, IntervalMonthDayNanoVector, IntervalYearMonthVector, - ListVector, ListVectorBuilder, MutableVector, NullVector, StringVector, TimeMicrosecondVector, - TimeMillisecondVector, TimeNanosecondVector, TimeSecondVector, TimestampMicrosecondVector, - TimestampMillisecondVector, TimestampNanosecondVector, TimestampSecondVector, UInt8Vector, - UInt16Vector, UInt32Vector, UInt64Vector, Vector, VectorRef, + ListVector, NullVector, StringVector, TimeMicrosecondVector, TimeMillisecondVector, + TimeNanosecondVector, TimeSecondVector, TimestampMicrosecondVector, TimestampMillisecondVector, + TimestampNanosecondVector, TimestampSecondVector, UInt8Vector, UInt16Vector, UInt32Vector, + UInt64Vector, Vector, VectorRef, }; /// Helper functions for `Vector`. @@ -102,135 +102,50 @@ impl Helper { }) } - /// Try to cast an arrow scalar value into vector - pub fn try_from_scalar_value(value: ScalarValue, length: usize) -> Result { - let vector = match value { - ScalarValue::Null => ConstantVector::new(Arc::new(NullVector::new(1)), length), - ScalarValue::Boolean(v) => { - ConstantVector::new(Arc::new(BooleanVector::from(vec![v])), length) + /// Materializes an Arrow scalar as a vector of the given length. + /// + /// With a type hint, casts the array to the requested representation and returns + /// an error if the vector cannot preserve that concrete type. Without a hint, + /// normalizes the scalar to a supported GreptimeDB representation. + pub fn try_from_scalar_value( + value: ScalarValue, + length: usize, + type_hint: Option<&ConcreteDataType>, + ) -> Result { + if let Some(data_type) = type_hint { + let mut array = value + .to_array_of_size(length) + .context(ConvertScalarToArrowArraySnafu)?; + let arrow_type = data_type.as_arrow_type(); + if array.data_type() != &arrow_type { + array = compute::cast(&array, &arrow_type).context(error::ArrowComputeSnafu)?; } - ScalarValue::Float16(v) => ConstantVector::new( - Arc::new(Float32Vector::from(vec![v.map(f32::from)])), - length, - ), - ScalarValue::Float32(v) => { - ConstantVector::new(Arc::new(Float32Vector::from(vec![v])), length) + let vector = Self::try_into_vector(array)?; + if &vector.data_type() != data_type { + return error::CastTypeSnafu { + msg: format!( + "Scalar materialization produced {:?}, expected {data_type:?}", + vector.data_type() + ), + } + .fail(); } - ScalarValue::Float64(v) => { - ConstantVector::new(Arc::new(Float64Vector::from(vec![v])), length) + return Ok(vector); + } + + let value = match value { + // GreptimeDB doesn't support Float16 vectors. + ScalarValue::Float16(v) => ScalarValue::Float32(v.map(f32::from)), + ScalarValue::LargeUtf8(v) => ScalarValue::Utf8(v), + ScalarValue::LargeBinary(v) | ScalarValue::FixedSizeBinary(_, v) => { + ScalarValue::Binary(v) } - ScalarValue::Int8(v) => { - ConstantVector::new(Arc::new(Int8Vector::from(vec![v])), length) - } - ScalarValue::Int16(v) => { - ConstantVector::new(Arc::new(Int16Vector::from(vec![v])), length) - } - ScalarValue::Int32(v) => { - ConstantVector::new(Arc::new(Int32Vector::from(vec![v])), length) - } - ScalarValue::Int64(v) => { - ConstantVector::new(Arc::new(Int64Vector::from(vec![v])), length) - } - ScalarValue::UInt8(v) => { - ConstantVector::new(Arc::new(UInt8Vector::from(vec![v])), length) - } - ScalarValue::UInt16(v) => { - ConstantVector::new(Arc::new(UInt16Vector::from(vec![v])), length) - } - ScalarValue::UInt32(v) => { - ConstantVector::new(Arc::new(UInt32Vector::from(vec![v])), length) - } - ScalarValue::UInt64(v) => { - ConstantVector::new(Arc::new(UInt64Vector::from(vec![v])), length) - } - ScalarValue::Utf8(v) | ScalarValue::LargeUtf8(v) => { - ConstantVector::new(Arc::new(StringVector::from(vec![v])), length) - } - ScalarValue::Binary(v) - | ScalarValue::LargeBinary(v) - | ScalarValue::FixedSizeBinary(_, v) => { - ConstantVector::new(Arc::new(BinaryVector::from(vec![v])), length) - } - ScalarValue::List(array) => { - let item_type = Arc::new(ConcreteDataType::try_from(&array.value_type())?); - let mut builder = ListVectorBuilder::with_type_capacity(item_type.clone(), 1); - let scalar_values = ScalarValue::convert_array_to_scalar_vec(array.as_ref()) - .context(ConvertArrowArrayToScalarsSnafu)?; - let values = scalar_values - .into_iter() - .flat_map(|v| v.unwrap_or_else(|| vec![ScalarValue::Null])) - .map(ScalarValue::try_into) - .collect::>>()?; - builder.push(Some(ListValueRef::Ref { - val: &ListValue::new(values, item_type), - })); - let list_vector = builder.to_vector(); - ConstantVector::new(list_vector, length) - } - ScalarValue::Date32(v) => { - ConstantVector::new(Arc::new(DateVector::from(vec![v])), length) - } - ScalarValue::TimestampSecond(v, _) => { - // Timezone is unimplemented now. - ConstantVector::new(Arc::new(TimestampSecondVector::from(vec![v])), length) - } - ScalarValue::TimestampMillisecond(v, _) => { - // Timezone is unimplemented now. - ConstantVector::new(Arc::new(TimestampMillisecondVector::from(vec![v])), length) - } - ScalarValue::TimestampMicrosecond(v, _) => { - // Timezone is unimplemented now. - ConstantVector::new(Arc::new(TimestampMicrosecondVector::from(vec![v])), length) - } - ScalarValue::TimestampNanosecond(v, _) => { - // Timezone is unimplemented now. - ConstantVector::new(Arc::new(TimestampNanosecondVector::from(vec![v])), length) - } - ScalarValue::Time32Second(v) => { - ConstantVector::new(Arc::new(TimeSecondVector::from(vec![v])), length) - } - ScalarValue::Time32Millisecond(v) => { - ConstantVector::new(Arc::new(TimeMillisecondVector::from(vec![v])), length) - } - ScalarValue::Time64Microsecond(v) => { - ConstantVector::new(Arc::new(TimeMicrosecondVector::from(vec![v])), length) - } - ScalarValue::Time64Nanosecond(v) => { - ConstantVector::new(Arc::new(TimeNanosecondVector::from(vec![v])), length) - } - ScalarValue::IntervalYearMonth(v) => { - ConstantVector::new(Arc::new(IntervalYearMonthVector::from(vec![v])), length) - } - ScalarValue::IntervalDayTime(v) => { - ConstantVector::new(Arc::new(IntervalDayTimeVector::from(vec![v])), length) - } - ScalarValue::IntervalMonthDayNano(v) => { - ConstantVector::new(Arc::new(IntervalMonthDayNanoVector::from(vec![v])), length) - } - ScalarValue::DurationSecond(v) => { - ConstantVector::new(Arc::new(DurationSecondVector::from(vec![v])), length) - } - ScalarValue::DurationMillisecond(v) => { - ConstantVector::new(Arc::new(DurationMillisecondVector::from(vec![v])), length) - } - ScalarValue::DurationMicrosecond(v) => { - ConstantVector::new(Arc::new(DurationMicrosecondVector::from(vec![v])), length) - } - ScalarValue::DurationNanosecond(v) => { - ConstantVector::new(Arc::new(DurationNanosecondVector::from(vec![v])), length) - } - ScalarValue::Decimal128(v, p, s) => { - let vector = Decimal128Vector::from(vec![v]).with_precision_and_scale(p, s)?; - ConstantVector::new(Arc::new(vector), length) - } - ScalarValue::Struct(v) => { - let struct_type = StructType::from(v.fields()); - ConstantVector::new( - Arc::new(StructVector::try_new(struct_type, (*v).clone())?), - length, - ) - } - ScalarValue::Decimal32(_, _, _) + // Timezones are not supported by GreptimeDB vectors. + ScalarValue::TimestampSecond(v, _) => ScalarValue::TimestampSecond(v, None), + ScalarValue::TimestampMillisecond(v, _) => ScalarValue::TimestampMillisecond(v, None), + ScalarValue::TimestampMicrosecond(v, _) => ScalarValue::TimestampMicrosecond(v, None), + ScalarValue::TimestampNanosecond(v, _) => ScalarValue::TimestampNanosecond(v, None), + value @ (ScalarValue::Decimal32(_, _, _) | ScalarValue::Decimal64(_, _, _) | ScalarValue::Decimal256(_, _, _) | ScalarValue::FixedSizeList(_) @@ -241,15 +156,19 @@ impl Helper { | ScalarValue::BinaryView(_) | ScalarValue::Map(_) | ScalarValue::Date64(_) - | ScalarValue::RunEndEncoded(_, _, _) => { + | ScalarValue::RunEndEncoded(_, _, _)) => { return error::ConversionSnafu { from: format!("Unsupported scalar value: {value}"), } .fail(); } + value => value, }; - Ok(Arc::new(vector)) + let array = value + .to_array_of_size(length) + .context(ConvertScalarToArrowArraySnafu)?; + Self::try_into_vector(array) } /// Try to cast an arrow array into vector @@ -457,12 +376,15 @@ mod tests { }; use arrow::buffer::Buffer; use arrow::datatypes::{Int32Type, IntervalMonthDayNano}; - use arrow_array::{BinaryArray, DictionaryArray, FixedSizeBinaryArray, LargeStringArray}; - use arrow_schema::DataType; + use arrow_array::{ + BinaryArray, DictionaryArray, FixedSizeBinaryArray, LargeStringArray, StructArray, + }; + use arrow_schema::{DataType, Field, Fields}; use common_decimal::Decimal128; use common_time::time::Time; use common_time::timestamp::TimeUnit; use common_time::{Date, Duration}; + use datafusion_common::scalar::ScalarStructBuilder; use super::*; use crate::value::Value; @@ -496,7 +418,7 @@ mod tests { #[test] fn test_try_from_scalar_date_value() { - let vector = Helper::try_from_scalar_value(ScalarValue::Date32(Some(42)), 3).unwrap(); + let vector = Helper::try_from_scalar_value(ScalarValue::Date32(Some(42)), 3, None).unwrap(); assert_eq!(ConcreteDataType::date_datatype(), vector.data_type()); assert_eq!(3, vector.len()); for i in 0..vector.len() { @@ -507,7 +429,7 @@ mod tests { #[test] fn test_try_from_scalar_duration_value() { let vector = - Helper::try_from_scalar_value(ScalarValue::DurationSecond(Some(42)), 3).unwrap(); + Helper::try_from_scalar_value(ScalarValue::DurationSecond(Some(42)), 3, None).unwrap(); assert_eq!( ConcreteDataType::duration_second_datatype(), vector.data_type() @@ -524,7 +446,8 @@ mod tests { #[test] fn test_try_from_scalar_decimal128_value() { let vector = - Helper::try_from_scalar_value(ScalarValue::Decimal128(Some(42), 3, 1), 3).unwrap(); + Helper::try_from_scalar_value(ScalarValue::Decimal128(Some(42), 3, 1), 3, None) + .unwrap(); assert_eq!( ConcreteDataType::decimal128_datatype(3, 1), vector.data_type() @@ -542,7 +465,7 @@ mod tests { &ArrowDataType::Int32, true, )); - let vector = Helper::try_from_scalar_value(value, 3).unwrap(); + let vector = Helper::try_from_scalar_value(value, 3, None).unwrap(); assert_eq!( ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::int32_datatype())), vector.data_type() @@ -555,6 +478,101 @@ mod tests { } } + #[test] + fn test_try_from_scalar_value_materializes_values() { + let vector = Helper::try_from_scalar_value(ScalarValue::Int32(Some(42)), 4, None).unwrap(); + assert_eq!(ConcreteDataType::int32_datatype(), vector.data_type()); + assert_eq!(4, vector.len()); + assert_eq!(0, vector.null_count()); + for i in 0..vector.len() { + assert_eq!(Value::Int32(42), vector.get(i)); + } + + let empty = Helper::try_from_scalar_value(ScalarValue::Int32(Some(42)), 0, None).unwrap(); + assert_eq!(ConcreteDataType::int32_datatype(), empty.data_type()); + assert!(empty.is_empty()); + + let nulls = Helper::try_from_scalar_value(ScalarValue::Int32(None), 3, None).unwrap(); + assert_eq!(3, nulls.len()); + assert_eq!(3, nulls.null_count()); + for i in 0..nulls.len() { + assert_eq!(Value::Null, nulls.get(i)); + } + } + + #[test] + fn test_try_from_scalar_struct_value() { + let fields = Fields::from(vec![ + Field::new("id", ArrowDataType::Int32, false), + Field::new("name", ArrowDataType::Utf8, true), + ]); + let value = ScalarValue::Struct(Arc::new(StructArray::new( + fields.clone(), + vec![ + ScalarValue::Int32(Some(7)).to_array().unwrap(), + ScalarValue::Utf8(Some("greptime".to_string())) + .to_array() + .unwrap(), + ], + None, + ))); + + let vector = Helper::try_from_scalar_value(value, 3, None).unwrap(); + assert_eq!( + ConcreteDataType::struct_datatype(StructType::from(&fields)), + vector.data_type() + ); + assert_eq!(3, vector.len()); + for i in 0..vector.len() { + let Value::Struct(value) = vector.get(i) else { + panic!("expected struct value"); + }; + assert_eq!( + &[Value::Int32(7), Value::String("greptime".into())], + value.items() + ); + } + + let null = ScalarStructBuilder::new_null(fields); + let vector = Helper::try_from_scalar_value(null, 2, None).unwrap(); + assert_eq!(2, vector.len()); + assert_eq!(2, vector.null_count()); + assert_eq!(Value::Null, vector.get(0)); + assert_eq!(Value::Null, vector.get(1)); + } + + #[test] + fn test_try_from_scalar_value_normalizes_arrow_types() { + let string = Helper::try_from_scalar_value( + ScalarValue::LargeUtf8(Some("greptime".to_string())), + 2, + None, + ) + .unwrap(); + assert_eq!(ConcreteDataType::string_datatype(), string.data_type()); + assert_eq!(&ArrowDataType::Utf8, string.to_arrow_array().data_type()); + + let binary = Helper::try_from_scalar_value( + ScalarValue::FixedSizeBinary(2, Some(vec![1, 2])), + 2, + None, + ) + .unwrap(); + assert_eq!(ConcreteDataType::binary_datatype(), binary.data_type()); + assert_eq!(&ArrowDataType::Binary, binary.to_arrow_array().data_type()); + + let timestamp = Helper::try_from_scalar_value( + ScalarValue::TimestampMillisecond(Some(42), Some("UTC".into())), + 2, + None, + ) + .unwrap(); + assert_eq!( + &ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), + timestamp.to_arrow_array().data_type() + ); + } + #[test] fn test_like_utf8() { fn assert_vector(expected: Vec<&str>, actual: &VectorRef) { @@ -745,7 +763,8 @@ mod tests { #[test] fn test_try_from_scalar_time_value() { - let vector = Helper::try_from_scalar_value(ScalarValue::Time32Second(Some(42)), 3).unwrap(); + let vector = + Helper::try_from_scalar_value(ScalarValue::Time32Second(Some(42)), 3, None).unwrap(); assert_eq!(ConcreteDataType::time_second_datatype(), vector.data_type()); assert_eq!(3, vector.len()); for i in 0..vector.len() { @@ -758,6 +777,7 @@ mod tests { let vector = Helper::try_from_scalar_value( ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNano::new(1, 1, 2000))), 3, + None, ) .unwrap(); diff --git a/src/datatypes/src/vectors/null.rs b/src/datatypes/src/vectors/null.rs index e049eb205c6..2510824ce0d 100644 --- a/src/datatypes/src/vectors/null.rs +++ b/src/datatypes/src/vectors/null.rs @@ -260,16 +260,6 @@ impl ScalarVectorBuilder for NullVectorBuilder { } } -pub(crate) fn replicate_null(vector: &NullVector, offsets: &[usize]) -> VectorRef { - assert_eq!(offsets.len(), vector.len()); - - if offsets.is_empty() { - return vector.slice(0, 0); - } - - Arc::new(NullVector::new(*offsets.last().unwrap())) -} - #[cfg(test)] mod tests { use serde_json; @@ -290,7 +280,6 @@ mod tests { assert_eq!(vector2.null_count(), 16); assert_eq!("NullVector", v.vector_type_name()); - assert!(!v.is_const()); assert!(v.validity().is_all_null()); assert!(v.only_null()); diff --git a/src/datatypes/src/vectors/operations.rs b/src/datatypes/src/vectors/operations.rs index 1dc9e21a0d2..f75ff060f3e 100644 --- a/src/datatypes/src/vectors/operations.rs +++ b/src/datatypes/src/vectors/operations.rs @@ -14,14 +14,12 @@ mod cast; mod filter; -mod replicate; mod take; use std::sync::Arc; use crate::error::{self, Result}; use crate::types::LogicalPrimitiveType; -use crate::vectors::constant::ConstantVector; use crate::vectors::{ BinaryVector, BooleanVector, ConcreteDataType, Decimal128Vector, ListVector, NullVector, PrimitiveVector, StringVector, UInt32Vector, Vector, VectorRef, @@ -29,14 +27,6 @@ use crate::vectors::{ /// Vector compute operations. pub trait VectorOp { - /// Copies each element according `offsets` parameter. - /// - `i-th` element should be copied `offsets[i] - offsets[i - 1]` times - /// - `0-th` element would be copied `offsets[0]` times - /// - /// # Panics - /// Panics if `offsets.len() != self.len()`. - fn replicate(&self, offsets: &[usize]) -> VectorRef; - /// Filters the vector, returns elements matching the `filter` (i.e. where the values are true). /// /// Note that the nulls of `filter` are interpreted as `false` will lead to these elements being masked out. @@ -57,10 +47,6 @@ pub trait VectorOp { macro_rules! impl_scalar_vector_op { ($($VectorType: ident),+) => {$( impl VectorOp for $VectorType { - fn replicate(&self, offsets: &[usize]) -> VectorRef { - replicate::replicate_scalar(self, offsets) - } - fn filter(&self, filter: &BooleanVector) -> Result { filter::filter_non_constant!(self, $VectorType, filter) } @@ -92,10 +78,6 @@ macro_rules! impl_scalar_vector_op { impl_scalar_vector_op!(BinaryVector, BooleanVector, StringVector); impl VectorOp for ListVector { - fn replicate(&self, offsets: &[usize]) -> VectorRef { - replicate::replicate_list(self, offsets) - } - fn filter(&self, filter: &BooleanVector) -> Result { filter::filter_non_constant!(self, ListVector, filter) } @@ -110,10 +92,6 @@ impl VectorOp for ListVector { } impl VectorOp for Decimal128Vector { - fn replicate(&self, offsets: &[usize]) -> VectorRef { - std::sync::Arc::new(replicate::replicate_decimal128(self, offsets)) - } - fn filter(&self, filter: &BooleanVector) -> Result { filter::filter_non_constant!(self, Decimal128Vector, filter) } @@ -128,10 +106,6 @@ impl VectorOp for Decimal128Vector { } impl VectorOp for PrimitiveVector { - fn replicate(&self, offsets: &[usize]) -> VectorRef { - std::sync::Arc::new(replicate::replicate_primitive(self, offsets)) - } - fn filter(&self, filter: &BooleanVector) -> Result { filter::filter_non_constant!(self, PrimitiveVector, filter) } @@ -146,10 +120,6 @@ impl VectorOp for PrimitiveVector { } impl VectorOp for NullVector { - fn replicate(&self, offsets: &[usize]) -> VectorRef { - replicate::replicate_null(self, offsets) - } - fn filter(&self, filter: &BooleanVector) -> Result { filter::filter_non_constant!(self, NullVector, filter) } @@ -166,21 +136,3 @@ impl VectorOp for NullVector { take::take_indices!(self, NullVector, indices) } } - -impl VectorOp for ConstantVector { - fn replicate(&self, offsets: &[usize]) -> VectorRef { - self.replicate_vector(offsets) - } - - fn filter(&self, filter: &BooleanVector) -> Result { - self.filter_vector(filter) - } - - fn cast(&self, to_type: &ConcreteDataType) -> Result { - self.cast_vector(to_type) - } - - fn take(&self, indices: &UInt32Vector) -> Result { - self.take_vector(indices) - } -} diff --git a/src/datatypes/src/vectors/operations/filter.rs b/src/datatypes/src/vectors/operations/filter.rs index 7dde6c1f946..b27bbb5de48 100644 --- a/src/datatypes/src/vectors/operations/filter.rs +++ b/src/datatypes/src/vectors/operations/filter.rs @@ -39,7 +39,6 @@ mod tests { TimestampMicrosecond, TimestampMillisecond, TimestampNanosecond, TimestampSecond, }; use crate::types::WrapperType; - use crate::vectors::constant::ConstantVector; use crate::vectors::{ BooleanVector, Int32Vector, NullVector, StringVector, VectorOp, VectorRef, }; @@ -67,26 +66,6 @@ mod tests { check_filter_primitive(&[5, 7], &[5, 6, 7], &[true, false, true]); } - fn check_filter_constant(expect_length: usize, input_length: usize, filter: &[bool]) { - let v = ConstantVector::new(Arc::new(Int32Vector::from_slice([123])), input_length); - let filter = BooleanVector::from_slice(filter); - let out = v.filter(&filter).unwrap(); - - assert!(out.is_const()); - assert_eq!(expect_length, out.len()); - } - - #[test] - fn test_filter_constant() { - check_filter_constant(0, 0, &[]); - check_filter_constant(1, 1, &[true]); - check_filter_constant(0, 1, &[false]); - check_filter_constant(1, 2, &[false, true]); - check_filter_constant(2, 2, &[true, true]); - check_filter_constant(1, 4, &[false, false, false, true]); - check_filter_constant(2, 4, &[false, true, false, true]); - } - #[test] fn test_filter_scalar() { let v = StringVector::from_slice(&["0", "1", "2", "3"]); diff --git a/src/datatypes/src/vectors/operations/replicate.rs b/src/datatypes/src/vectors/operations/replicate.rs deleted file mode 100644 index dee0b3855b0..00000000000 --- a/src/datatypes/src/vectors/operations/replicate.rs +++ /dev/null @@ -1,232 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use crate::prelude::*; -pub(crate) use crate::vectors::decimal::replicate_decimal128; -pub(crate) use crate::vectors::null::replicate_null; -pub(crate) use crate::vectors::primitive::replicate_primitive; -use crate::vectors::{ListVector, ListVectorBuilder}; - -pub(crate) fn replicate_scalar(c: &C, offsets: &[usize]) -> VectorRef { - assert_eq!(offsets.len(), c.len()); - - if offsets.is_empty() { - return c.slice(0, 0); - } - let mut builder = <::Builder>::with_capacity(c.len()); - - let mut previous_offset = 0; - for (i, offset) in offsets.iter().enumerate() { - let data = c.get_data(i); - for _ in previous_offset..*offset { - builder.push(data.clone()); - } - previous_offset = *offset; - } - builder.to_vector() -} - -pub(crate) fn replicate_list(c: &ListVector, offsets: &[usize]) -> VectorRef { - assert_eq!(offsets.len(), c.len()); - - if offsets.is_empty() { - return c.slice(0, 0); - } - let mut builder = - ListVectorBuilder::with_type_capacity(c.item_type(), *offsets.last().unwrap()); - - let mut previous_offset = 0; - for (i, offset) in offsets.iter().enumerate() { - let data = c.get_data(i); - for _ in previous_offset..*offset { - builder.push(data.clone()); - } - previous_offset = *offset; - } - builder.to_vector() -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use common_time::timestamp::TimeUnit; - use common_time::{Date, Timestamp}; - use paste::paste; - - use super::*; - use crate::value::{ListValue, ListValueRef}; - use crate::vectors::constant::ConstantVector; - use crate::vectors::{ - Decimal128Vector, Int32Vector, ListVectorBuilder, NullVector, StringVector, VectorOp, - }; - - #[test] - fn test_replicate_primitive() { - let v = Int32Vector::from_iterator(0..5); - let offsets = [0, 1, 2, 3, 4]; - - let v = v.replicate(&offsets); - assert_eq!(4, v.len()); - - for i in 0..4 { - assert_eq!(Value::Int32(i as i32 + 1), v.get(i)); - } - } - - #[test] - fn test_replicate_nullable_primitive() { - let v = Int32Vector::from(vec![None, Some(1), None, Some(2)]); - let offsets = [2, 4, 6, 8]; - let v = v.replicate(&offsets); - assert_eq!(8, v.len()); - - let expect: VectorRef = Arc::new(Int32Vector::from(vec![ - None, - None, - Some(1), - Some(1), - None, - None, - Some(2), - Some(2), - ])); - assert_eq!(expect, v); - } - - #[test] - fn test_replicate_scalar() { - let v = StringVector::from_slice(&["0", "1", "2", "3"]); - let offsets = [1, 3, 5, 6]; - - let v = v.replicate(&offsets); - assert_eq!(6, v.len()); - - let expect: VectorRef = Arc::new(StringVector::from_slice(&["0", "1", "1", "2", "2", "3"])); - assert_eq!(expect, v); - } - - #[test] - fn test_replicate_list() { - let item_type = Arc::new(ConcreteDataType::int32_datatype()); - let first = ListValue::new(vec![Value::Int32(1), Value::Int32(2)], item_type.clone()); - let second = ListValue::new(vec![Value::Int32(3)], item_type.clone()); - let mut builder = ListVectorBuilder::with_type_capacity(item_type, 2); - builder.push(Some(ListValueRef::Ref { val: &first })); - builder.push(Some(ListValueRef::Ref { val: &second })); - let v = builder.finish(); - - let v = v.replicate(&[1, 3]); - let mut expect_builder = - ListVectorBuilder::with_type_capacity(Arc::new(ConcreteDataType::int32_datatype()), 3); - expect_builder.push(Some(ListValueRef::Ref { val: &first })); - expect_builder.push(Some(ListValueRef::Ref { val: &second })); - expect_builder.push(Some(ListValueRef::Ref { val: &second })); - let expect = expect_builder.to_vector(); - assert_eq!(expect, v); - } - - #[test] - fn test_replicate_constant() { - let v = Arc::new(StringVector::from_slice(&["hello"])); - let cv = ConstantVector::new(v.clone(), 2); - let offsets = [1, 4]; - - let cv = cv.replicate(&offsets); - assert_eq!(4, cv.len()); - - let expect: VectorRef = Arc::new(ConstantVector::new(v, 4)); - assert_eq!(expect, cv); - } - - #[test] - fn test_replicate_null() { - let v = NullVector::new(0); - let offsets = []; - let v = v.replicate(&offsets); - assert!(v.is_empty()); - - let v = NullVector::new(3); - let offsets = [1, 3, 5]; - - let v = v.replicate(&offsets); - assert_eq!(5, v.len()); - } - - macro_rules! impl_replicate_date_like_test { - ($VectorType: ident, $ValueType: ident, $method: ident) => {{ - use $crate::vectors::$VectorType; - - let v = $VectorType::from_iterator((0..5).map($ValueType::$method)); - let offsets = [0, 1, 2, 3, 4]; - - let v = v.replicate(&offsets); - assert_eq!(4, v.len()); - - for i in 0..4 { - assert_eq!( - Value::$ValueType($ValueType::$method((i as i32 + 1).into())), - v.get(i) - ); - } - }}; - } - - macro_rules! impl_replicate_timestamp_test { - ($unit: ident) => {{ - paste!{ - use $crate::vectors::[]; - use $crate::timestamp::[]; - let v = []::from_iterator((0..5).map([]::from)); - let offsets = [0, 1, 2, 3, 4]; - let v = v.replicate(&offsets); - assert_eq!(4, v.len()); - for i in 0..4 { - assert_eq!( - Value::Timestamp(Timestamp::new(i as i64 + 1, TimeUnit::$unit)), - v.get(i) - ); - } - } - }}; - } - - #[test] - fn test_replicate_date_like() { - impl_replicate_date_like_test!(DateVector, Date, new); - impl_replicate_timestamp_test!(Second); - impl_replicate_timestamp_test!(Millisecond); - impl_replicate_timestamp_test!(Microsecond); - impl_replicate_timestamp_test!(Nanosecond); - } - - #[test] - fn test_replicate_decimal() { - let data = vec![100]; - // create a decimal vector - let v = Decimal128Vector::from_values(data.clone()) - .with_precision_and_scale(10, 2) - .unwrap(); - let offsets = [5]; - let v = v.replicate(&offsets); - assert_eq!(5, v.len()); - - let expect: VectorRef = Arc::new( - Decimal128Vector::from_values(vec![100; 5]) - .with_precision_and_scale(10, 2) - .unwrap(), - ); - assert_eq!(expect, v); - } -} diff --git a/src/datatypes/src/vectors/operations/take.rs b/src/datatypes/src/vectors/operations/take.rs index 41947d3ef51..167f650c8ed 100644 --- a/src/datatypes/src/vectors/operations/take.rs +++ b/src/datatypes/src/vectors/operations/take.rs @@ -43,8 +43,7 @@ mod tests { use crate::types::{LogicalPrimitiveType, WrapperType}; use crate::vectors::operations::VectorOp; use crate::vectors::{ - BooleanVector, ConstantVector, Int32Vector, NullVector, PrimitiveVector, StringVector, - UInt32Vector, + BooleanVector, Int32Vector, NullVector, PrimitiveVector, StringVector, UInt32Vector, }; fn check_take_primitive( @@ -119,29 +118,6 @@ mod tests { take_time_like_test!(TimestampNanosecondVector, TimestampNanosecond, from_native); } - fn check_take_constant(expect_length: usize, input_length: usize, indices: &[u32]) { - let v = ConstantVector::new(Arc::new(Int32Vector::from_slice([111])), input_length); - let indices = UInt32Vector::from_slice(indices); - let out = v.take(&indices).unwrap(); - - assert!(out.is_const()); - assert_eq!(expect_length, out.len()); - } - - #[test] - fn test_take_constant() { - check_take_constant(2, 5, &[3, 4]); - check_take_constant(3, 10, &[1, 2, 3]); - check_take_constant(4, 10, &[1, 5, 3, 6]); - check_take_constant(5, 10, &[1, 9, 8, 7, 3]); - } - - #[test] - #[should_panic] - fn test_take_constant_out_of_index() { - check_take_constant(2, 5, &[3, 5]); - } - #[test] #[should_panic] fn test_take_out_of_index() { diff --git a/src/datatypes/src/vectors/primitive.rs b/src/datatypes/src/vectors/primitive.rs index dc85f94c4ff..92c8394e710 100644 --- a/src/datatypes/src/vectors/primitive.rs +++ b/src/datatypes/src/vectors/primitive.rs @@ -366,40 +366,6 @@ where } } -pub(crate) fn replicate_primitive( - vector: &PrimitiveVector, - offsets: &[usize], -) -> PrimitiveVector { - assert_eq!(offsets.len(), vector.len()); - - if offsets.is_empty() { - return vector.get_slice(0, 0); - } - - let mut builder = PrimitiveVectorBuilder::::with_capacity(*offsets.last().unwrap()); - - let mut previous_offset = 0; - - for (offset, value) in offsets.iter().zip(vector.array.iter()) { - let repeat_times = *offset - previous_offset; - match value { - Some(data) => { - unsafe { - // Safety: std::iter::Repeat and std::iter::Take implement TrustedLen. - builder - .mutable_array - .append_trusted_len_iter(std::iter::repeat_n(data, repeat_times)); - } - } - None => { - builder.mutable_array.append_nulls(repeat_times); - } - } - previous_offset = *offset; - } - builder.finish() -} - #[cfg(test)] mod tests { use std::vec; @@ -433,7 +399,6 @@ mod tests { fn check_vec(v: Int32Vector) { assert_eq!(4, v.len()); assert_eq!("Int32Vector", v.vector_type_name()); - assert!(!v.is_const()); assert!(v.validity().is_all_valid()); assert!(!v.only_null()); diff --git a/src/datatypes/src/vectors/string.rs b/src/datatypes/src/vectors/string.rs index 53324280c3f..0476ecc6440 100644 --- a/src/datatypes/src/vectors/string.rs +++ b/src/datatypes/src/vectors/string.rs @@ -604,7 +604,6 @@ mod tests { let v = StringVector::from(strs.clone()); assert_eq!(3, v.len()); assert_eq!("StringVector", v.vector_type_name()); - assert!(!v.is_const()); assert!(v.validity().is_all_valid()); assert!(!v.only_null()); assert_eq!(1040, v.memory_size()); diff --git a/src/datatypes/src/vectors/struct_vector.rs b/src/datatypes/src/vectors/struct_vector.rs index e6284351b4f..4798c8fddcd 100644 --- a/src/datatypes/src/vectors/struct_vector.rs +++ b/src/datatypes/src/vectors/struct_vector.rs @@ -15,7 +15,7 @@ use std::any::Any; use std::sync::Arc; -use arrow::array::{MutableArrayData, NullBufferBuilder}; +use arrow::array::NullBufferBuilder; use arrow::compute::TakeOptions; use arrow::datatypes::DataType as ArrowDataType; use arrow_array::{Array, ArrayRef, StructArray}; @@ -143,31 +143,6 @@ impl Vector for StructVector { } impl VectorOp for StructVector { - fn replicate(&self, offsets: &[usize]) -> VectorRef { - assert_eq!(offsets.len(), self.len()); - assert!(offsets.is_sorted(), "offsets must be non-decreasing"); - - let Some(&output_len) = offsets.last() else { - return self.slice(0, 0); - }; - - let source = self.array.to_data(); - let mut output = MutableArrayData::new(vec![&source], false, output_len); - let mut previous_offset = 0; - - for (index, &offset) in offsets.iter().enumerate() { - for _ in previous_offset..offset { - output.extend(0, index, index + 1); - } - previous_offset = offset; - } - - Arc::new(StructVector { - array: StructArray::from(output.freeze()), - fields: self.fields.clone(), - }) - } - fn cast(&self, _to_type: &ConcreteDataType) -> Result { UnsupportedOperationSnafu { op: "cast", @@ -466,10 +441,6 @@ impl ScalarVectorBuilder for StructVectorBuilder { #[cfg(test)] mod tests { - use arrow::array::{DictionaryArray, Int8Array, StringArray}; - use arrow::buffer::NullBuffer; - use arrow::datatypes::Int8Type; - use super::*; use crate::json::JsonSettings; use crate::schema::{ColumnDefaultConstraint, ColumnSchema}; @@ -555,7 +526,7 @@ mod tests { } #[test] - fn test_replicate_preserves_json2_identity() { + fn test_default_vector_preserves_json2_identity() { let json = JsonSettings::default() .encode(serde_json::json!({"answer": 42})) .unwrap(); @@ -576,39 +547,6 @@ mod tests { assert_eq!(replicated.len(), 2); } - #[test] - fn test_replicate_preserves_dictionary_and_nulls() { - let fields = StructType::new(Arc::new(vec![StructField::new( - "label", - ConcreteDataType::dictionary_datatype( - ConcreteDataType::int8_datatype(), - ConcreteDataType::string_datatype(), - ), - true, - )])); - let dictionary = DictionaryArray::::new( - Int8Array::from(vec![Some(0), Some(1)]), - Arc::new(StringArray::from(vec!["a", "b"])), - ); - let array = StructArray::new( - fields.as_arrow_fields(), - vec![Arc::new(dictionary)], - Some(NullBuffer::from(vec![true, false])), - ); - let vector = StructVector::try_new(fields.clone(), array).unwrap(); - - let replicated = vector.replicate(&[2, 3]); - - assert_eq!( - replicated.data_type(), - ConcreteDataType::struct_datatype(fields) - ); - assert_eq!(replicated.len(), 3); - assert_eq!(replicated.null_count(), 1); - assert_eq!(replicated.get(0), replicated.get(1)); - assert!(replicated.is_null(2)); - } - #[test] fn test_deep_nested_struct_list() { // level 1: struct diff --git a/src/flow/src/batching_mode/frontend_client.rs b/src/flow/src/batching_mode/frontend_client.rs index 10ce3b6820d..e932933ff84 100644 --- a/src/flow/src/batching_mode/frontend_client.rs +++ b/src/flow/src/batching_mode/frontend_client.rs @@ -183,13 +183,21 @@ impl DatabaseWithPeer { /// Try sending a "SELECT 1" to the database async fn try_select_one(&self) -> Result<(), Error> { // notice here use `sql` for `SELECT 1` return 1 row - let _ = self + let output = self .database .sql("SELECT 1") .await .with_context(|_| InvalidRequestSnafu { context: format!("Failed to handle `SELECT 1` request at {:?}", self.peer), })?; + + if let OutputData::Stream(stream) = output.data { + common_recordbatch::util::collect(stream) + .await + .map_err(BoxedError::new) + .context(ExternalSnafu)?; + } + Ok(()) } } @@ -698,9 +706,11 @@ mod tests { use api::v1::query_request::Query; use arrow_flight::flight_service_server::FlightServiceServer; use arrow_flight::{FlightData, Ticket}; + use common_grpc::flight::FlightEncoder; use common_query::{Output, OutputData}; use common_recordbatch::adapter::RecordBatchMetrics; use common_recordbatch::{OrderOption, RecordBatch, RecordBatchStream}; + use datatypes::arrow::datatypes::Schema as ArrowSchema; use datatypes::prelude::{ConcreteDataType, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::Int32Vector; @@ -775,6 +785,14 @@ mod tests { #[derive(Debug)] struct SlowFlight; + struct DelayedEofFlight { + schema_sent: Mutex>>, + release: Arc, + } + + #[derive(Debug)] + struct LateStreamErrorFlight; + struct WaitForConcurrentFlight { barrier: Arc, } @@ -870,6 +888,45 @@ mod tests { } } + #[async_trait::async_trait] + impl FlightCraft for DelayedEofFlight { + async fn do_get( + &self, + _request: TonicRequest, + ) -> std::result::Result>, Status> { + let schema = FlightEncoder::default().encode_schema(&ArrowSchema::empty()); + let schema_sent = self.schema_sent.lock().unwrap().take(); + let schema_stream = futures::stream::once(async move { + if let Some(schema_sent) = schema_sent { + let _ = schema_sent.send(()); + } + Ok(schema) + }); + let release = self.release.clone(); + let delayed_eof = futures::stream::unfold(release, |release| async move { + release.notified().await; + None::<(std::result::Result, _)> + }); + + Ok(TonicResponse::new(Box::pin( + schema_stream.chain(delayed_eof), + ))) + } + } + + #[async_trait::async_trait] + impl FlightCraft for LateStreamErrorFlight { + async fn do_get( + &self, + _request: TonicRequest, + ) -> std::result::Result>, Status> { + let schema = FlightEncoder::default().encode_schema(&ArrowSchema::empty()); + let stream = + futures::stream::iter([Ok(schema), Err(Status::unavailable("late stream error"))]); + Ok(TonicResponse::new(Box::pin(stream))) + } + } + #[async_trait::async_trait] impl FlightCraft for WaitForConcurrentFlight { async fn do_get( @@ -1056,6 +1113,71 @@ mod tests { assert!(format!("{err:?}").contains("Invalid value for flow.return_region_seq")); } + #[tokio::test] + async fn test_try_select_one_waits_for_stream_eof() { + let (schema_sent, schema_sent_rx) = tokio::sync::oneshot::channel(); + let release = Arc::new(tokio::sync::Notify::new()); + let (addr, server) = start_flight_server(DelayedEofFlight { + schema_sent: Mutex::new(Some(schema_sent)), + release: release.clone(), + }) + .await; + let database = Database::new( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + Client::with_urls([addr.as_str()]), + ); + let db = DatabaseWithPeer::new( + database, + Peer { + id: 1, + addr: addr.clone(), + }, + ); + let mut probe = tokio::spawn(async move { db.try_select_one().await }); + + timeout(Duration::from_secs(1), schema_sent_rx) + .await + .expect("server should send the schema") + .expect("schema signal should be sent"); + assert!( + timeout(Duration::from_millis(100), &mut probe) + .await + .is_err(), + "SELECT 1 must wait for the delayed stream tail and EOF" + ); + + release.notify_one(); + timeout(Duration::from_secs(1), &mut probe) + .await + .expect("SELECT 1 should complete after EOF") + .expect("probe task should not panic") + .expect("SELECT 1 should succeed after EOF"); + server.abort(); + } + + #[tokio::test] + async fn test_try_select_one_propagates_late_stream_error() { + let (addr, server) = start_flight_server(LateStreamErrorFlight).await; + let database = Database::new( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + Client::with_urls([addr.as_str()]), + ); + let db = DatabaseWithPeer::new( + database, + Peer { + id: 1, + addr: addr.clone(), + }, + ); + + let err = db.try_select_one().await.unwrap_err(); + server.abort(); + + assert!(format!("{err:?}").contains("late stream error")); + } + #[tokio::test] async fn test_check_all_frontends_without_auth_fails_fast_on_unauthenticated_frontend() { let (addr, server) = start_flight_server(RejectUnauthenticatedFlight).await; diff --git a/src/flow/src/expr/scalar.rs b/src/flow/src/expr/scalar.rs index 18649bff123..af16381e4cf 100644 --- a/src/flow/src/expr/scalar.rs +++ b/src/flow/src/expr/scalar.rs @@ -216,6 +216,7 @@ impl ScalarExpr { msg: "Failed to convert literal to scalar value", })?, batch.row_count(), + None, ) .context(DataTypeSnafu { msg: "Failed to convert scalar value to vector ref when parsing literal", diff --git a/src/flow/src/server.rs b/src/flow/src/server.rs index 9bda516f24c..1b26fbbf15e 100644 --- a/src/flow/src/server.rs +++ b/src/flow/src/server.rs @@ -124,6 +124,7 @@ impl flow_server::Flow for FlowService { api::v1::region::InsertRequest { region_id: insert.region_id, rows: insert.rows, + skip_wal: false, partition_expr_version: insert.partition_expr_version, } }) diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index dd89a1f3f95..f834722cca5 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -1329,6 +1329,7 @@ impl PrometheusHandler for Instance { self.check_query_target_permission(targets, &query_ctx) .await?; + let requires_output_ordering = query.requires_output_ordering(); let (query, stmt) = query.into_parts(); let QueryStatement::Promql(eval_stmt, _) = &stmt else { @@ -1342,6 +1343,12 @@ impl PrometheusHandler for Instance { .map_err(BoxedError::new) .context(ExecuteQuerySnafu)?; + let plan = if requires_output_ordering { + plan + } else { + promql::remove_output_sort(plan) + }; + interceptor.pre_execute(&query, &eval_stmt.expr, Some(&plan), query_ctx.clone())?; // Take the EvalStmt from the original QueryStatement and use it to create the CatalogQueryStatement. diff --git a/src/frontend/src/instance/grpc.rs b/src/frontend/src/instance/grpc.rs index 7d4f0ee7959..2d89b5489df 100644 --- a/src/frontend/src/instance/grpc.rs +++ b/src/frontend/src/instance/grpc.rs @@ -453,6 +453,7 @@ impl Instance { request.flight_data, request.record_batch, request.schema_bytes, + ctx.skip_wal(), ) .await .context(TableOperationSnafu)?; diff --git a/src/frontend/src/instance/promql.rs b/src/frontend/src/instance/promql.rs index 0f6121441f6..3df70224703 100644 --- a/src/frontend/src/instance/promql.rs +++ b/src/frontend/src/instance/promql.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; use std::time::SystemTime; use auth::PermissionTableTarget; @@ -21,6 +22,7 @@ use common_catalog::consts::INFORMATION_SCHEMA_NAME; use common_catalog::format_full_table_name; use common_recordbatch::util; use common_telemetry::tracing; +use datafusion_expr::LogicalPlan; use promql_parser::label::{Matcher, Matchers}; use query::promql; use query::promql::planner::PromPlanner; @@ -35,6 +37,21 @@ use crate::error::{ }; use crate::instance::Instance; +/// Strips the output sort a PromQL plan ends with, keeping the plan schema intact. +/// +/// Only the sort the caller would observe is removed: recursion stops at any other +/// node, so ordering consumed by windows, limits or PromQL extension nodes stays. +pub(super) fn remove_output_sort(plan: LogicalPlan) -> LogicalPlan { + match plan { + LogicalPlan::Sort(sort) if sort.fetch.is_none() => Arc::unwrap_or_clone(sort.input), + LogicalPlan::Projection(mut projection) => { + projection.input = Arc::new(remove_output_sort(Arc::unwrap_or_clone(projection.input))); + LogicalPlan::Projection(projection) + } + plan => plan, + } +} + impl Instance { /// Handles metric names query request, returns the names. #[tracing::instrument(skip_all)] @@ -180,3 +197,79 @@ impl Instance { Ok(results) } } + +#[cfg(test)] +mod tests { + use datafusion::arrow::array::Int64Array; + use datafusion::prelude::{SessionConfig, SessionContext}; + use datafusion_expr::{LogicalPlanBuilder, Sort, col, lit}; + + use super::*; + + #[tokio::test] + async fn remove_output_sort_keeps_schema_and_row_selection() { + let input = + LogicalPlanBuilder::values(vec![vec![lit(3_i64)], vec![lit(1_i64)], vec![lit(2_i64)]]) + .unwrap() + .build() + .unwrap(); + let sort = Sort { + expr: vec![col("column1").sort(true, false)], + input: Arc::new(input), + fetch: None, + }; + let projected = LogicalPlanBuilder::from(LogicalPlan::Sort(sort.clone())) + .project(vec![col("column1").alias("sample")]) + .unwrap() + .project(vec![col("sample")]) + .unwrap() + .build() + .unwrap(); + // The root sort is removed, but the sort feeding the limit selects the rows. + let limited = LogicalPlanBuilder::from(LogicalPlan::Sort(sort.clone())) + .limit(0, Some(2)) + .unwrap() + .sort(vec![col("column1").sort(false, false)]) + .unwrap() + .build() + .unwrap(); + let fetched = LogicalPlan::Sort(Sort { + fetch: Some(2), + ..sort.clone() + }); + + let context = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + for (name, plan, expected) in [ + ("root", LogicalPlan::Sort(sort), vec![3, 1, 2]), + ("projection", projected, vec![3, 1, 2]), + ("sort below limit", limited, vec![1, 2]), + ("fetch", fetched, vec![1, 2]), + ] { + let schema = plan.schema().clone(); + let plan = remove_output_sort(plan); + assert_eq!(plan.schema(), &schema, "{name}"); + let output = context + .execute_logical_plan(plan) + .await + .unwrap() + .collect() + .await + .unwrap(); + let values = output + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + assert_eq!(values, expected, "{name}"); + } + } +} diff --git a/src/meta-srv/src/error.rs b/src/meta-srv/src/error.rs index 53872d4817f..70997ebe7d1 100644 --- a/src/meta-srv/src/error.rs +++ b/src/meta-srv/src/error.rs @@ -368,9 +368,9 @@ pub enum Error { location: Location, }, - #[snafu(display("Failed to deserialize from json: {}", input))] + #[snafu(display("Failed to deserialize from json payload of length {}", input_len))] DeserializeFromJson { - input: String, + input_len: usize, #[snafu(source)] error: serde_json::error::Error, #[snafu(implicit)] diff --git a/src/meta-srv/src/event/gc.rs b/src/meta-srv/src/event/gc.rs index 7f679f1785b..94d671a176b 100644 --- a/src/meta-srv/src/event/gc.rs +++ b/src/meta-srv/src/event/gc.rs @@ -218,6 +218,7 @@ mod tests { use common_event_recorder::testing::assert_event_contract; use common_event_recorder::{EventTypeFilter, PersistentEventContext, TriggerReason}; use common_meta::key::TableMetadataManager; + use common_meta::key::runtime_switch::RuntimeSwitchManager; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::sequence::SequenceBuilder; use common_procedure::{ @@ -588,11 +589,13 @@ mod tests { fn batch_gc_procedure() -> BatchGcProcedure { let kv_backend = Arc::new(MemoryKvBackend::new()); let table_metadata_manager = Arc::new(TableMetadataManager::new(kv_backend.clone())); + let runtime_switch_manager = Arc::new(RuntimeSwitchManager::new(kv_backend.clone())); let mailbox_sequence = SequenceBuilder::new("test_batch_gc_event", kv_backend).build(); let mailbox = MailboxContext::new(mailbox_sequence); BatchGcProcedure::new( mailbox.mailbox().clone(), table_metadata_manager, + runtime_switch_manager, "localhost".to_string(), vec![RegionId::new(1024, 1)], true, diff --git a/src/meta-srv/src/gc/ctx.rs b/src/meta-srv/src/gc/ctx.rs index 4b46a941e25..857f5e87cdc 100644 --- a/src/meta-srv/src/gc/ctx.rs +++ b/src/meta-srv/src/gc/ctx.rs @@ -27,6 +27,7 @@ use common_meta::ddl_manager::DdlManagerRef; #[cfg(feature = "enterprise")] use common_meta::key::DroppedTableName; use common_meta::key::TableMetadataManagerRef; +use common_meta::key::runtime_switch::RuntimeSwitchManagerRef; use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; #[cfg(feature = "enterprise")] @@ -155,6 +156,8 @@ pub(crate) struct DefaultGcSchedulerCtx { pub(crate) table_metadata_manager: TableMetadataManagerRef, /// Procedure manager. pub(crate) procedure_manager: ProcedureManagerRef, + /// Runtime switch manager used by recovered and newly submitted GC procedures. + pub(crate) runtime_switch_manager: RuntimeSwitchManagerRef, /// DDL manager used to submit the existing purge procedure. #[cfg(feature = "enterprise")] pub(crate) ddl_manager: DdlManagerRef, @@ -176,6 +179,7 @@ impl DefaultGcSchedulerCtx { pub fn try_new( table_metadata_manager: TableMetadataManagerRef, procedure_manager: ProcedureManagerRef, + runtime_switch_manager: RuntimeSwitchManagerRef, #[cfg(feature = "enterprise")] ddl_manager: DdlManagerRef, meta_peer_client: MetaPeerClientRef, mailbox: MailboxRef, @@ -184,6 +188,7 @@ impl DefaultGcSchedulerCtx { Ok(Self { table_metadata_manager, procedure_manager, + runtime_switch_manager, #[cfg(feature = "enterprise")] ddl_manager, #[cfg(feature = "enterprise")] @@ -320,6 +325,7 @@ impl DefaultGcSchedulerCtx { let procedure = BatchGcProcedure::new( self.mailbox.clone(), self.table_metadata_manager.clone(), + self.runtime_switch_manager.clone(), self.server_addr.clone(), region_ids.to_vec(), full_file_listing, diff --git a/src/meta-srv/src/gc/procedure.rs b/src/meta-srv/src/gc/procedure.rs index f9f597577ae..db22f3ab00d 100644 --- a/src/meta-srv/src/gc/procedure.rs +++ b/src/meta-srv/src/gc/procedure.rs @@ -18,8 +18,14 @@ use std::sync::Arc; use std::time::Duration; use api::v1::meta::MailboxMessage; -use common_meta::instruction::{self, GcRegions, GetFileRefs, GetFileRefsReply, InstructionReply}; +use common_error::ext::PlainError; +use common_error::status_code::StatusCode; +use common_meta::instruction::{ + self, GetPackedFileRefs, GetPackedFileRefsReply, InstructionReply, PackedFileRefsManifest, + PackedGcRegions, PackedRegionFileRefs, +}; use common_meta::key::TableMetadataManagerRef; +use common_meta::key::runtime_switch::RuntimeSwitchManagerRef; use common_meta::key::table_repart::TableRepartValue; use common_meta::key::table_route::PhysicalTableRouteValue; use common_meta::lock_key::{RegionLock, TableLock}; @@ -34,6 +40,7 @@ use common_telemetry::tracing::Instrument as _; use common_telemetry::tracing_context::TracingContext; use common_telemetry::{debug, error, info, warn}; use futures::future::join_all; +use futures::stream::{FuturesUnordered, StreamExt}; use itertools::Itertools as _; use serde::{Deserialize, Serialize}; use snafu::ResultExt as _; @@ -49,14 +56,14 @@ use crate::metrics::{METRIC_META_GC_DATANODE_CALLS_TOTAL, METRIC_META_GC_FAILED_ use crate::procedure::utils::{instruction_error_result, instruction_to_error}; use crate::service::mailbox::{Channel, MailboxReceiver, MailboxRef}; -async fn send_get_file_refs_inner( +async fn send_get_packed_file_refs_inner( mailbox: &MailboxRef, server_addr: &str, peer: &Peer, - instruction: GetFileRefs, + instruction: GetPackedFileRefs, timeout: Duration, ) -> Result { - let instruction = instruction::Instruction::GetFileRefs(instruction); + let instruction = instruction::Instruction::GetPackedFileRefs(instruction); let tracing_ctx = TracingContext::from_current_span(); let msg = MailboxMessage::json_message( &format!("Get file references: {}", instruction), @@ -75,25 +82,25 @@ async fn send_get_file_refs_inner( .await } -async fn recv_get_file_refs_reply( +async fn recv_get_packed_file_refs_reply( peer: &Peer, mailbox_rx: MailboxReceiver, -) -> Result { +) -> Result { let reply = match mailbox_rx.await { Ok(reply_msg) => HeartbeatMailbox::json_reply(&reply_msg)?, Err(e) => { error!( - e; "Failed to receive reply from datanode {} for GetFileRefs instruction", + e; "Failed to receive reply from datanode {} for GetPackedFileRefs instruction", peer, ); return Err(e); } }; - let InstructionReply::GetFileRefs(reply) = reply else { + let InstructionReply::GetPackedFileRefs(reply) = reply else { return error::UnexpectedInstructionReplySnafu { - mailbox_message: format!("{:?}", reply), - reason: "Unexpected reply of the GetFileRefs instruction", + mailbox_message: "unexpected instruction reply for GetPackedFileRefs".to_string(), + reason: "Unexpected reply of the GetPackedFileRefs instruction", } .fail(); }; @@ -104,12 +111,11 @@ async fn recv_get_file_refs_reply( async fn send_gc_regions_inner( mailbox: &MailboxRef, peer: &Peer, - gc_regions: &GcRegions, + instruction: instruction::Instruction, server_addr: &str, timeout: Duration, description: &str, ) -> Result { - let instruction = instruction::Instruction::GcRegions(gc_regions.clone()); let tracing_ctx = TracingContext::from_current_span(); let msg = MailboxMessage::json_message( &format!("{}: {}", description, instruction), @@ -128,9 +134,43 @@ async fn send_gc_regions_inner( .await } +fn scoped_gc_instruction( + regions: Vec, + full_manifest: &FileRefsManifest, + full_file_listing: bool, +) -> instruction::Instruction { + let packed_file_refs_manifest = PackedFileRefsManifest { + file_refs: regions + .iter() + .filter_map(|region| { + full_manifest + .file_refs + .get(region) + .map(|refs| (*region, PackedRegionFileRefs::from_refs(refs))) + }) + .collect(), + manifest_version: regions + .iter() + .filter_map(|region| { + full_manifest + .manifest_version + .get(region) + .map(|version| (*region, *version)) + }) + .collect(), + cross_region_refs: HashMap::new(), + }; + + instruction::Instruction::PackedGcRegions(PackedGcRegions { + regions, + packed_file_refs_manifest, + full_file_listing, + }) +} + async fn recv_gc_regions_reply( peer: &Peer, - gc_regions: &GcRegions, + regions: &[RegionId], description: &str, mailbox_rx: MailboxReceiver, ) -> Result { @@ -147,7 +187,7 @@ async fn recv_gc_regions_reply( let InstructionReply::GcRegions(reply) = reply else { return error::UnexpectedInstructionReplySnafu { - mailbox_message: format!("{:?}", reply), + mailbox_message: "unexpected instruction reply for GcRegions".to_string(), reason: "Unexpected reply of the GcRegions instruction", } .fail(); @@ -158,14 +198,16 @@ async fn recv_gc_regions_reply( Ok(report) => Ok(report), Err(e) => { error!( - e; "Datanode {} reported error during GC for regions {:?}", - peer, gc_regions + e; "Datanode {} reported error during GC for {} regions", + peer, regions.len() ); instruction_error_result( &e, format!( - "Datanode {} reported error during GC for regions {:?}: {}", - peer, gc_regions, e + "Datanode {} reported error during GC for {} regions: {}", + peer, + regions.len(), + e ), ) } @@ -177,6 +219,7 @@ async fn recv_gc_regions_reply( pub struct BatchGcProcedure { mailbox: MailboxRef, table_metadata_manager: TableMetadataManagerRef, + runtime_switch_manager: RuntimeSwitchManagerRef, data: BatchGcData, } @@ -217,9 +260,11 @@ pub enum State { impl BatchGcProcedure { pub const TYPE_NAME: &'static str = "metasrv-procedure::BatchGcProcedure"; + #[allow(clippy::too_many_arguments)] pub fn new( mailbox: MailboxRef, table_metadata_manager: TableMetadataManagerRef, + runtime_switch_manager: RuntimeSwitchManagerRef, server_addr: String, regions: Vec, full_file_listing: bool, @@ -229,6 +274,7 @@ impl BatchGcProcedure { Self { mailbox, table_metadata_manager, + runtime_switch_manager, data: BatchGcData { state: State::Start, server_addr, @@ -251,6 +297,7 @@ impl BatchGcProcedure { pub fn new_update_repartition_for_test( mailbox: MailboxRef, table_metadata_manager: TableMetadataManagerRef, + runtime_switch_manager: RuntimeSwitchManagerRef, server_addr: String, regions: Vec, file_refs: FileRefsManifest, @@ -259,6 +306,7 @@ impl BatchGcProcedure { Self { mailbox, table_metadata_manager, + runtime_switch_manager, data: BatchGcData { state: State::UpdateRepartition, server_addr, @@ -286,6 +334,22 @@ impl BatchGcProcedure { }) } + async fn check_maintenance_mode(&self) -> ProcedureResult<()> { + let enabled = self + .runtime_switch_manager + .maintenance_mode() + .await + .context(error::RuntimeSwitchManagerSnafu) + .map_err(ProcedureError::retry_later)?; + if enabled { + return Err(ProcedureError::external(PlainError::new( + "maintenance mode is enabled".to_string(), + StatusCode::IllegalState, + ))); + } + Ok(()) + } + fn merge_gc_report(&mut self, report: GcReport) { let accumulated = self.data.gc_report.get_or_insert_default(); // Deleted objects are cumulative, while these sets describe the latest outcome @@ -644,10 +708,10 @@ impl BatchGcProcedure { } } - // Send GetFileRefs instructions to each datanode + // Send packed GetFileRefs instructions to each datanode let mut all_file_refs: HashMap> = HashMap::new(); let mut all_manifest_versions = HashMap::new(); - let mut all_cross_region_refs = HashMap::new(); + let mut all_cross_region_refs: HashMap> = HashMap::new(); let mut peers = HashSet::new(); peers.extend(datanode2query_regions.keys().cloned()); @@ -655,33 +719,34 @@ impl BatchGcProcedure { let mailbox = &self.mailbox; let server_addr = &self.data.server_addr; - let mut tasks = Vec::new(); - + // Each future owns both send and receive, so a completed reply is merged and + // dropped immediately rather than retained in a second join_all buffer. + let mut tasks = FuturesUnordered::new(); for peer in peers { let regions = datanode2query_regions.remove(&peer).unwrap_or_default(); let related_regions_for_peer = datanode2related_regions.remove(&peer).unwrap_or_default(); - if regions.is_empty() && related_regions_for_peer.is_empty() { continue; } - tasks.push(async move { - let instruction = GetFileRefs { - query_regions: regions.clone(), - related_regions: related_regions_for_peer.clone(), + let instruction = GetPackedFileRefs { + query_regions: regions, + related_regions: related_regions_for_peer, }; - - let reply = - send_get_file_refs_inner(mailbox, server_addr, &peer, instruction, timeout) - .await; - - (peer, regions, related_regions_for_peer, reply) + let rx = send_get_packed_file_refs_inner( + mailbox, + server_addr, + &peer, + instruction, + timeout, + ) + .await?; + let reply = recv_get_packed_file_refs_reply(&peer, rx).await?; + Ok::<_, crate::error::Error>((peer, reply)) }); } - let mut recv_tasks = Vec::new(); - // store error to make sure metrics doesn't ignore other peers let mut first_error = None; let mut record_get_file_refs_error = |e| { METRIC_META_GC_DATANODE_CALLS_TOTAL @@ -691,80 +756,56 @@ impl BatchGcProcedure { first_error = Some(e); } }; - for (peer, regions, related_regions_for_peer, reply) in join_all(tasks).await { - match reply { - Ok(mailbox_rx) => { - recv_tasks.push(async move { - let reply = recv_get_file_refs_reply(&peer, mailbox_rx).await; - (peer, regions, related_regions_for_peer, reply) - }); - } - Err(e) => record_get_file_refs_error(e), - } - } - - let replies = join_all(recv_tasks).await; - - for (peer, regions, related_regions_for_peer, reply) in replies { - let reply = match reply { + while let Some(result) = tasks.next().await { + let (peer, reply) = match result { Ok(reply) => reply, Err(e) => { record_get_file_refs_error(e); continue; } }; - debug!( - "Got file references from datanode: {:?}, query_regions: {:?}, related_regions: {:?}, reply: {:?}", - peer, regions, related_regions_for_peer, reply - ); - if !reply.success { - METRIC_META_GC_DATANODE_CALLS_TOTAL - .with_label_values(&["get_file_refs", "error"]) - .inc(); let err = if let Some(error) = &reply.error { instruction_to_error( error, - format!( - "Failed to get file references from datanode {}: {:?}", - peer, error - ), + format!("Failed to get file references from datanode {peer}"), ) } else { error::UnexpectedSnafu { - violated: format!( - "Failed to get file references from datanode {}: {:?}", - peer, reply.error - ), + violated: + "Datanode returned an unsuccessful GetPackedFileRefs reply without an error" + .to_string(), } .build() }; record_get_file_refs_error(err); continue; } + let manifest = match reply.packed_file_refs_manifest.into_manifest() { + Ok(manifest) => manifest, + Err(err) => { + record_get_file_refs_error(error::Error::Other { + source: common_error::ext::BoxedError::new(err), + location: snafu::Location::new(file!(), line!(), 0), + }); + continue; + } + }; METRIC_META_GC_DATANODE_CALLS_TOTAL .with_label_values(&["get_file_refs", "success"]) .inc(); - - // Merge the file references from this datanode - for (region_id, file_refs) in reply.file_refs_manifest.file_refs { - all_file_refs - .entry(region_id) - .or_default() - .extend(file_refs); + for (region_id, refs) in manifest.file_refs { + all_file_refs.entry(region_id).or_default().extend(refs); } - - // region manifest version should be the smallest one among all peers, so outdated region can be detected - for (region_id, version) in reply.file_refs_manifest.manifest_version { + for (region_id, version) in manifest.manifest_version { let entry = all_manifest_versions.entry(region_id).or_insert(version); *entry = (*entry).min(version); } - - for (region_id, related_region_ids) in reply.file_refs_manifest.cross_region_refs { - let entry = all_cross_region_refs + for (region_id, related) in manifest.cross_region_refs { + all_cross_region_refs .entry(region_id) - .or_insert_with(HashSet::new); - entry.extend(related_region_ids); + .or_default() + .extend(related); } } @@ -813,28 +854,23 @@ impl BatchGcProcedure { let tasks = datanode2regions .into_iter() .map(|(peer, regions_for_peer)| { - let gc_regions = GcRegions { - regions: regions_for_peer.clone(), - // file_refs_manifest could be somewhere large. But still intentionally clone per datanode here: - // this path is admin-triggered or scheduler-triggered, peer count is expected to be bounded, and - // and abnormal manifest growth should be addressed at the source - file_refs_manifest: file_refs.clone(), - full_file_listing, - }; - let region_count = gc_regions.regions.len() as u64; + let region_count = regions_for_peer.len() as u64; + let regions = regions_for_peer.clone(); + let instruction = + scoped_gc_instruction(regions_for_peer, file_refs, full_file_listing); async move { let report = send_gc_regions_inner( mailbox, &peer, - &gc_regions, + instruction, server_addr, timeout, "Batch GC", ) .await; - (peer, gc_regions, region_count, report) + (peer, regions, region_count, report) } }); @@ -851,12 +887,12 @@ impl BatchGcProcedure { first_error = Some(e); } }; - for (peer, gc_regions, region_count, report) in join_all(tasks).await { + for (peer, regions, region_count, report) in join_all(tasks).await { match report { Ok(mailbox_rx) => { recv_tasks.push(async move { let report = - recv_gc_regions_reply(&peer, &gc_regions, "Batch GC", mailbox_rx).await; + recv_gc_regions_reply(&peer, ®ions, "Batch GC", mailbox_rx).await; (peer, region_count, report) }); } @@ -887,12 +923,12 @@ impl BatchGcProcedure { if need_retry.is_empty() { info!( - "GC report from datanode {}: successfully deleted files for regions {:?}", + "GC report from datanode {}: successfully deleted files for region IDs {:?}", peer, success ); } else { warn!( - "GC report from datanode {}: successfully deleted files for regions {:?}, need retry for regions {:?}", + "GC report from datanode {}: successfully deleted files for region IDs {:?}, need retry for region IDs {:?}", peer, success, need_retry ); } @@ -921,6 +957,8 @@ impl Procedure for BatchGcProcedure { } async fn execute(&mut self, ctx: &ProcedureContext) -> ProcedureResult { + self.check_maintenance_mode().await?; + match self.data.state { State::Start => { let _regions_span = common_telemetry::tracing::debug_span!( @@ -1122,11 +1160,15 @@ mod tests { use api::v1::meta::MailboxMessage; use api::v1::meta::mailbox_message::Payload; - use common_meta::instruction::{GcRegionsReply, InstructionReply}; + use common_meta::instruction::{GcRegionsReply, Instruction, InstructionReply}; use common_meta::key::TableMetadataManager; + use common_meta::key::runtime_switch::RuntimeSwitchManager; use common_meta::kv_backend::memory::MemoryKvBackend; + use common_meta::kv_backend::test_util::MockKvBackendBuilder; use common_meta::peer::Peer; use common_meta::sequence::SequenceBuilder; + use common_procedure::Context as ProcedureContext; + use common_procedure_test::MockContextProvider; use common_time::util::current_time_millis; use store_api::storage::FileId; use tokio::sync::mpsc; @@ -1135,6 +1177,38 @@ mod tests { use crate::procedure::test_util::{MailboxContext, send_mock_reply}; use crate::service::mailbox::Channel; + #[test] + fn test_scoped_gc_instruction_selects_and_scopes_manifest() { + let region = RegionId::new(7, 3); + let other = RegionId::new(7, 4); + let mut manifest = FileRefsManifest::default(); + manifest.file_refs.insert(region, HashSet::new()); + manifest.file_refs.insert(other, HashSet::new()); + manifest.manifest_version.insert(region, 42); + + let packed = scoped_gc_instruction(vec![region], &manifest, true); + let Instruction::PackedGcRegions(packed) = packed else { + panic!("expected packed instruction"); + }; + assert_eq!(packed.regions, vec![region]); + assert_eq!( + packed.packed_file_refs_manifest.manifest_version[®ion], + 42 + ); + assert!( + !packed + .packed_file_refs_manifest + .file_refs + .contains_key(&other) + ); + assert!( + packed + .packed_file_refs_manifest + .cross_region_refs + .is_empty() + ); + } + #[test] fn test_done_with_gc_report_keeps_report() { let region_id = RegionId::new(1024, 1); @@ -1215,6 +1289,7 @@ mod tests { let kv_backend = Arc::new(MemoryKvBackend::new()); let table_metadata_manager = Arc::new(TableMetadataManager::new(kv_backend.clone())); + let runtime_switch_manager = Arc::new(RuntimeSwitchManager::new(kv_backend.clone())); let mailbox_sequence = SequenceBuilder::new("test_batch_gc_partial_report", kv_backend).build(); let mut mailbox = MailboxContext::new(mailbox_sequence); @@ -1230,6 +1305,7 @@ mod tests { let mut procedure = BatchGcProcedure::new( mailbox.mailbox().clone(), table_metadata_manager, + runtime_switch_manager, "localhost".to_string(), vec![first_region, second_region], true, @@ -1265,11 +1341,15 @@ mod tests { fn batch_gc_procedure() -> BatchGcProcedure { let kv_backend = Arc::new(MemoryKvBackend::new()); let table_metadata_manager = Arc::new(TableMetadataManager::new(kv_backend.clone())); - let mailbox_sequence = SequenceBuilder::new("test_batch_gc_procedure", kv_backend).build(); + let runtime_switch_manager = Arc::new(RuntimeSwitchManager::new(kv_backend)); + let mailbox_sequence = + SequenceBuilder::new("test_batch_gc_procedure", Arc::new(MemoryKvBackend::new())) + .build(); let mailbox = MailboxContext::new(mailbox_sequence); BatchGcProcedure::new( mailbox.mailbox().clone(), table_metadata_manager, + runtime_switch_manager, "localhost".to_string(), vec![RegionId::new(1024, 1)], true, @@ -1277,4 +1357,113 @@ mod tests { HashMap::new(), ) } + + #[tokio::test] + async fn test_maintenance_mode_gates_every_gc_state() { + let states = [ + State::Start, + State::Acquiring, + State::Gcing, + State::UpdateRepartition, + ]; + + for state in states { + let kv_backend = Arc::new(MemoryKvBackend::new()); + let table_metadata_manager = Arc::new(TableMetadataManager::new(kv_backend.clone())); + let runtime_switch_manager = Arc::new(RuntimeSwitchManager::new(kv_backend.clone())); + runtime_switch_manager.set_maintenance_mode().await.unwrap(); + let mailbox_sequence = + SequenceBuilder::new("test_batch_gc_maintenance_gate", kv_backend).build(); + let mut mailbox = MailboxContext::new(mailbox_sequence); + let (tx, mut rx) = mpsc::channel(1); + mailbox + .insert_heartbeat_response_receiver(Channel::Datanode(1), tx) + .await; + let mut procedure = BatchGcProcedure::new( + mailbox.mailbox().clone(), + table_metadata_manager, + runtime_switch_manager, + "localhost".to_string(), + vec![RegionId::new(1024, 1)], + true, + Duration::from_secs(10), + HashMap::new(), + ); + procedure.data.state = state.clone(); + let dump_before = procedure.dump().unwrap(); + + let ctx = ProcedureContext { + procedure_id: common_procedure::ProcedureId::random(), + provider: Arc::new(MockContextProvider::default()), + event_context: None, + }; + let err = procedure.execute(&ctx).await.unwrap_err(); + + assert!(!err.is_retry_later()); + assert_eq!( + common_error::ext::ErrorExt::retry_hint(&err), + common_error::ext::RetryHint::NonRetryable + ); + assert_eq!(procedure.data.state, state); + assert_eq!(procedure.dump().unwrap(), dump_before); + if matches!(state, State::Acquiring | State::Gcing) { + assert!(rx.try_recv().is_err()); + } + } + } + + #[tokio::test] + async fn test_maintenance_mode_read_error_retries_without_advancing_state() { + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let failing_once_calls = calls.clone(); + let kv_backend = Arc::new( + MockKvBackendBuilder::default() + .range_fn(Arc::new(move |_| { + if failing_once_calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 { + common_meta::error::UnexpectedSnafu { + err_msg: "maintenance read failed", + } + .fail() + } else { + Ok(common_meta::rpc::store::RangeResponse { + kvs: vec![], + more: false, + }) + } + })) + .build() + .unwrap(), + ); + let table_metadata_manager = Arc::new(TableMetadataManager::new(kv_backend.clone())); + let runtime_switch_manager = Arc::new(RuntimeSwitchManager::new(kv_backend.clone())); + let mailbox_sequence = + SequenceBuilder::new("test_batch_gc_maintenance_read_error", kv_backend).build(); + let mailbox = MailboxContext::new(mailbox_sequence); + let mut procedure = BatchGcProcedure::new( + mailbox.mailbox().clone(), + table_metadata_manager, + runtime_switch_manager, + "localhost".to_string(), + vec![RegionId::new(1024, 1)], + true, + Duration::from_secs(10), + HashMap::new(), + ); + let ctx = ProcedureContext { + procedure_id: common_procedure::ProcedureId::random(), + provider: Arc::new(MockContextProvider::default()), + event_context: None, + }; + let dump_before = procedure.dump().unwrap(); + + assert!(procedure.execute(&ctx).await.unwrap_err().is_retry_later()); + assert_eq!(procedure.data.state, State::Start); + assert_eq!(procedure.dump().unwrap(), dump_before); + + assert!(matches!( + procedure.execute(&ctx).await.unwrap(), + Status::Executing { .. } + )); + assert_eq!(procedure.data.state, State::Acquiring); + } } diff --git a/src/meta-srv/src/handler.rs b/src/meta-srv/src/handler.rs index 50b53f4970a..320f26fae7c 100644 --- a/src/meta-srv/src/handler.rs +++ b/src/meta-srv/src/handler.rs @@ -428,9 +428,11 @@ impl HeartbeatMailbox { .as_ref() .with_context(|| UnexpectedInstructionReplySnafu { mailbox_message: msg.to_string(), - reason: format!("empty payload, msg: {msg:?}"), + reason: "empty JSON payload".to_string(), })?; - serde_json::from_str(payload).context(DeserializeFromJsonSnafu { input: payload }) + serde_json::from_str(payload).context(DeserializeFromJsonSnafu { + input_len: payload.len(), + }) } /// Parses the [Instruction] from [MailboxMessage]. @@ -443,9 +445,11 @@ impl HeartbeatMailbox { .as_ref() .with_context(|| UnexpectedInstructionReplySnafu { mailbox_message: msg.to_string(), - reason: format!("empty payload, msg: {msg:?}"), + reason: "empty JSON payload".to_string(), })?; - serde_json::from_str(payload).context(DeserializeFromJsonSnafu { input: payload }) + serde_json::from_str(payload).context(DeserializeFromJsonSnafu { + input_len: payload.len(), + }) } pub fn create(pushers: Pushers, sequence: Sequence) -> MailboxRef { @@ -526,7 +530,12 @@ impl Mailbox for HeartbeatMailbox { msg.id = message_id; let pusher_id = ch.pusher_id(); - debug!("Sending mailbox message {msg:?} to {pusher_id}"); + let payload_len = msg + .payload + .as_ref() + .map(|Payload::Json(payload)| payload.len()) + .unwrap_or_default(); + debug!(message_id, payload_len, %pusher_id, "Sending mailbox message"); let (tx, rx) = oneshot::channel(); let _ = self.senders.insert(message_id, tx); @@ -548,7 +557,12 @@ impl Mailbox for HeartbeatMailbox { msg.id = message_id; let pusher_id = ch.pusher_id(); - debug!("Sending mailbox message {msg:?} to {pusher_id}"); + let payload_len = msg + .payload + .as_ref() + .map(|Payload::Json(payload)| payload.len()) + .unwrap_or_default(); + debug!(message_id, payload_len, %pusher_id, "Sending one-way mailbox message"); self.pushers.push(pusher_id, msg).await?; @@ -560,7 +574,18 @@ impl Mailbox for HeartbeatMailbox { } async fn on_recv(&self, id: MessageId, maybe_msg: Result) -> Result<()> { - debug!("Received mailbox message {maybe_msg:?}"); + let payload_len = maybe_msg + .as_ref() + .ok() + .and_then(|msg| msg.payload.as_ref()) + .map(|Payload::Json(payload)| payload.len()) + .unwrap_or_default(); + debug!( + message_id = id, + payload_len, + success = maybe_msg.is_ok(), + "Received mailbox message" + ); let _ = self.timeouts.remove(&id); @@ -568,7 +593,15 @@ impl Mailbox for HeartbeatMailbox { tx.send(maybe_msg) .map_err(|_| error::MailboxClosedSnafu { id }.build())?; } else if let Ok(finally_msg) = maybe_msg { - warn!("The response arrived too late: {finally_msg:?}"); + let payload_len = finally_msg + .payload + .as_ref() + .map(|Payload::Json(payload)| payload.len()) + .unwrap_or_default(); + warn!( + message_id = id, + payload_len, "The mailbox response arrived too late" + ); } Ok(()) diff --git a/src/meta-srv/src/key.rs b/src/meta-srv/src/key.rs index 803ee9f0c6e..abdc92ff6d4 100644 --- a/src/meta-srv/src/key.rs +++ b/src/meta-srv/src/key.rs @@ -87,7 +87,9 @@ impl FromStr for LeaseValue { type Err = error::Error; fn from_str(value: &str) -> crate::Result { - serde_json::from_str(value).context(error::DeserializeFromJsonSnafu { input: value }) + serde_json::from_str(value).context(error::DeserializeFromJsonSnafu { + input_len: value.len(), + }) } } diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs index 75b72f9a114..b40f06e3cfe 100644 --- a/src/meta-srv/src/metasrv/builder.rs +++ b/src/meta-srv/src/metasrv/builder.rs @@ -530,6 +530,7 @@ impl MetasrvBuilder { let gc_scheduler_ctx = DefaultGcSchedulerCtx::try_new( table_metadata_manager.clone(), procedure_manager.clone(), + runtime_switch_manager.clone(), #[cfg(feature = "enterprise")] ddl_manager.clone(), meta_peer_client.clone(), diff --git a/src/meta-srv/src/service/heartbeat.rs b/src/meta-srv/src/service/heartbeat.rs index 066d1560477..5cdc5b57894 100644 --- a/src/meta-srv/src/service/heartbeat.rs +++ b/src/meta-srv/src/service/heartbeat.rs @@ -206,7 +206,13 @@ where /// Handles the incoming heartbeat request, and returns whether to continue the session. async fn handle_request(&mut self, req: HeartbeatRequest, is_handshake: bool) -> bool { - debug!("Receiving heartbeat request: {:?}", req); + let mailbox_message_id = req.mailbox_message.as_ref().map(|msg| msg.id); + debug!( + sender_id = %self.sender_id, + ?mailbox_message_id, + is_handshake, + "Receiving heartbeat request" + ); let sender_id = self.sender_id.to_string(); METRIC_META_HEARTBEAT_RECV @@ -224,7 +230,12 @@ where let is_not_leader = res.as_ref().is_ok_and(|r| r.is_not_leader()); - debug!("Sending heartbeat response: {:?}", res); + debug!( + sender_id = %self.sender_id, + success = res.is_ok(), + is_not_leader, + "Sending heartbeat response" + ); if self.tx.send(res).await.is_err() { info!( diff --git a/src/metric-engine/src/engine.rs b/src/metric-engine/src/engine.rs index 6d5dace9c43..18d6ff6f197 100644 --- a/src/metric-engine/src/engine.rs +++ b/src/metric-engine/src/engine.rs @@ -636,6 +636,7 @@ mod test { .handle_request( logical_region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: Rows { schema: row_schema_with_tags(&["job"]), rows: build_rows(1, 5), diff --git a/src/metric-engine/src/engine/bulk_insert.rs b/src/metric-engine/src/engine/bulk_insert.rs index d700d858acf..cb43a71406d 100644 --- a/src/metric-engine/src/engine/bulk_insert.rs +++ b/src/metric-engine/src/engine/bulk_insert.rs @@ -124,6 +124,7 @@ impl MetricEngineInner { let aligned_schema_version = Some(self.physical_schema_version(data_region_id).await?); let request = RegionBulkInsertsRequest { + skip_wal: request.skip_wal, region_id: data_region_id, payload: modified_batch, raw_data: ArrowIpc { @@ -241,18 +242,28 @@ mod tests { use common_error::ext::ErrorExt; use common_query::prelude::{greptime_timestamp, greptime_value}; use common_recordbatch::RecordBatches; + use datafusion::parquet::basic::Encoding; use datatypes::arrow::array::{Float64Array, StringArray, TimestampMillisecondArray}; use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema, TimeUnit}; use datatypes::arrow::record_batch::RecordBatch; use mito2::config::MitoConfig; - use store_api::metric_engine_consts::PRIMARY_KEY_ENCODING; + use mito2::sst::parquet::metadata::MetadataLoader; + use mito2::sst::parquet::reader::MetadataCacheMetrics; + use store_api::metric_engine_consts::{ + METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY, PRIMARY_KEY_ENCODING, + }; + use store_api::mito_engine_options::EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING; use store_api::path_utils::table_dir; use store_api::region_engine::RegionEngine; - use store_api::region_request::{RegionBulkInsertsRequest, RegionPutRequest, RegionRequest}; + use store_api::region_request::{ + PathType, RegionBulkInsertsRequest, RegionCloseRequest, RegionFlushRequest, + RegionOpenRequest, RegionPutRequest, RegionRequest, + }; use store_api::storage::{RegionId, ScanRequest}; use super::record_batch_to_ipc; use crate::batch_modifier::{TagColumnInfo, modify_batch_sparse}; + use crate::engine::MetricEngine; use crate::error::Error; use crate::test_util::{self, TestEnv}; @@ -287,9 +298,14 @@ mod tests { .unwrap() } - fn build_bulk_request(logical_region_id: RegionId, batch: RecordBatch) -> RegionRequest { + fn build_bulk_request( + logical_region_id: RegionId, + batch: RecordBatch, + skip_wal: bool, + ) -> RegionRequest { let (schema, data_header, payload) = record_batch_to_ipc(&batch).unwrap(); RegionRequest::BulkInserts(RegionBulkInsertsRequest { + skip_wal, region_id: logical_region_id, payload: batch, raw_data: ArrowIpc { @@ -332,6 +348,7 @@ mod tests { let batch = build_logical_batch(0, 0); let request = RegionRequest::BulkInserts(RegionBulkInsertsRequest { + skip_wal: false, region_id: logical_region_id, payload: batch, raw_data: ArrowIpc::default(), @@ -348,6 +365,11 @@ mod tests { #[tokio::test] async fn test_bulk_insert_physical_region_passthrough() { + check_bulk_insert_physical_region_passthrough(false).await; + check_bulk_insert_physical_region_passthrough(true).await; + } + + async fn check_bulk_insert_physical_region_passthrough(skip_wal: bool) { // Use flat format so that BulkMemtable is used (supports write_bulk). let mito_config = MitoConfig { default_flat_format: true, @@ -355,12 +377,13 @@ mod tests { }; let env = TestEnv::with_mito_config("", mito_config, Default::default()).await; env.init_metric_region().await; + env.metric().inner.flush_task.stop().await.unwrap(); let physical_region_id = env.default_physical_region_id(); let logical_region_id = env.default_logical_region_id(); // First, do a normal logical bulk insert so we can compare results. let logical_batch = build_logical_batch(0, 3); - let logical_request = build_bulk_request(logical_region_id, logical_batch.clone()); + let logical_request = build_bulk_request(logical_region_id, logical_batch, skip_wal); let response = env .metric() .handle_request(logical_region_id, logical_request) @@ -385,7 +408,7 @@ mod tests { &non_tag_indices, ) .unwrap(); - let request = build_bulk_request(physical_region_id, physical_batch); + let request = build_bulk_request(physical_region_id, physical_batch, skip_wal); let response = env .metric() .handle_request(physical_region_id, request) @@ -401,6 +424,48 @@ mod tests { .unwrap(); let batches = RecordBatches::try_collect(stream).await.unwrap(); assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 6); + + // Closing without a flush must recover only WAL-backed data. Recreate + // the wrapper too, so its metadata cache cannot hide missing metadata. + let stat = env.mito().region_statistic(physical_region_id).unwrap(); + assert_eq!(stat.sst_num, 0); + env.metric() + .handle_request( + physical_region_id, + RegionRequest::Close(RegionCloseRequest { + flush_on_close: false, + }), + ) + .await + .unwrap(); + let reopened = MetricEngine::try_new(env.mito(), Default::default()).unwrap(); + reopened.inner.flush_task.stop().await.unwrap(); + reopened + .handle_request( + physical_region_id, + RegionRequest::Open(RegionOpenRequest { + engine: METRIC_ENGINE_NAME.to_string(), + table_dir: TestEnv::default_table_dir(), + path_type: PathType::Bare, + options: [(PHYSICAL_TABLE_METADATA_KEY.to_string(), String::new())] + .into_iter() + .collect(), + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + let stream = reopened + .scan_to_stream(logical_region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + batches.iter().map(|b| b.num_rows()).sum::(), + if skip_wal { 0 } else { 6 }, + ); } #[tokio::test] @@ -415,7 +480,7 @@ mod tests { let physical_region_id = env.default_physical_region_id(); let batch = build_logical_batch(0, 0); - let request = build_bulk_request(physical_region_id, batch); + let request = build_bulk_request(physical_region_id, batch, false); let response = env .metric() .handle_request(physical_region_id, request) @@ -449,7 +514,7 @@ mod tests { ) .unwrap(); - let request = build_bulk_request(logical_region_id, batch); + let request = build_bulk_request(logical_region_id, batch, false); let err = env .metric() .handle_request(logical_region_id, request) @@ -499,7 +564,7 @@ mod tests { ) .unwrap(); - let request = build_bulk_request(logical_region_id, batch); + let request = build_bulk_request(logical_region_id, batch, false); let response = env .metric() .handle_request(logical_region_id, request) @@ -522,7 +587,7 @@ mod tests { env.init_metric_region().await; let logical_region_id = env.default_logical_region_id(); - let request = build_bulk_request(logical_region_id, build_logical_batch(0, 3)); + let request = build_bulk_request(logical_region_id, build_logical_batch(0, 3), false); let response = env .metric() .handle_request(logical_region_id, request) @@ -530,7 +595,7 @@ mod tests { .unwrap(); assert_eq!(response.affected_rows, 3); - let request = build_bulk_request(logical_region_id, build_logical_batch(3, 5)); + let request = build_bulk_request(logical_region_id, build_logical_batch(3, 5), false); let response = env .metric() .handle_request(logical_region_id, request) @@ -550,16 +615,68 @@ mod tests { #[tokio::test] async fn test_bulk_insert_sparse_encoding() { let env = TestEnv::new().await; - env.init_metric_region().await; + let physical_region_id = env.default_physical_region_id(); + env.create_physical_region( + physical_region_id, + &TestEnv::default_table_dir(), + vec![( + EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING.to_string(), + "byte_stream_split".to_string(), + )], + ) + .await; let logical_region_id = env.default_logical_region_id(); + env.create_logical_region(physical_region_id, logical_region_id) + .await; - let request = build_bulk_request(logical_region_id, build_logical_batch(0, 4)); + let rows = 4; + let request = build_bulk_request(logical_region_id, build_logical_batch(0, rows), false); let response = env .metric() .handle_request(logical_region_id, request) .await .unwrap(); - assert_eq!(response.affected_rows, 4); + assert_eq!(response.affected_rows, rows); + + let data_region_id = crate::utils::to_data_region_id(physical_region_id); + env.mito() + .handle_request( + data_region_id, + RegionRequest::Flush(RegionFlushRequest::default()), + ) + .await + .unwrap(); + + let region = env.mito().find_region(data_region_id).unwrap(); + let entry = region + .manifest_sst_entries() + .await + .into_iter() + .find(|entry| entry.visible && entry.file_path.ends_with(".parquet")) + .unwrap(); + let mut cache_metrics = MetadataCacheMetrics::default(); + let footer = MetadataLoader::new( + region.access_layer().object_store().clone(), + &entry.file_path, + entry.file_size, + ) + .load(&mut cache_metrics) + .await + .unwrap(); + let field_column_index = footer + .file_metadata() + .schema_descr() + .columns() + .iter() + .position(|column| column.name() == greptime_value()) + .unwrap(); + assert!(!footer.row_groups().is_empty()); + assert!(footer.row_groups().iter().all(|row_group| { + row_group + .column(field_column_index) + .encodings() + .any(|encoding| encoding == Encoding::BYTE_STREAM_SPLIT) + })); let stream = env .metric() @@ -567,7 +684,7 @@ mod tests { .await .unwrap(); let batches = RecordBatches::try_collect(stream).await.unwrap(); - assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 4); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), rows); } #[tokio::test] @@ -575,7 +692,7 @@ mod tests { let env = TestEnv::new().await; let logical_region_id = init_dense_metric_region(&env).await; - let request = build_bulk_request(logical_region_id, build_logical_batch(0, 2)); + let request = build_bulk_request(logical_region_id, build_logical_batch(0, 2), false); let err = env .metric() .handle_request(logical_region_id, request) @@ -599,6 +716,7 @@ mod tests { .handle_request( logical_region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: api::v1::Rows { schema, rows }, hint: None, partition_expr_version: None, @@ -616,7 +734,7 @@ mod tests { let env_bulk = TestEnv::new().await; env_bulk.init_metric_region().await; - let request = build_bulk_request(logical_region_id, build_logical_batch(0, 5)); + let request = build_bulk_request(logical_region_id, build_logical_batch(0, 5), false); env_bulk .metric() .handle_request(logical_region_id, request) diff --git a/src/metric-engine/src/engine/flush.rs b/src/metric-engine/src/engine/flush.rs index 21e642de0b2..426aa918d04 100644 --- a/src/metric-engine/src/engine/flush.rs +++ b/src/metric-engine/src/engine/flush.rs @@ -90,6 +90,7 @@ mod tests { let schema = row_schema_with_tags(&["job"]); let rows = build_rows(1, 10); let request = RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: Rows { schema, rows }, hint: None, partition_expr_version: None, diff --git a/src/metric-engine/src/engine/put.rs b/src/metric-engine/src/engine/put.rs index 96ed2c3e60a..8858251d082 100644 --- a/src/metric-engine/src/engine/put.rs +++ b/src/metric-engine/src/engine/put.rs @@ -142,6 +142,8 @@ impl MetricEngineInner { let data_region_id = to_data_region_id(physical_region_id); let primary_key_encoding = self.get_primary_key_encoding(data_region_id)?; + // TODO(weny): Consolidate validation and merging to avoid redundant request traversals, + // while ensuring the entire batch is validated before writing. // Validate all requests self.validate_batch_requests(physical_region_id, &mut requests) .await?; @@ -176,6 +178,18 @@ impl MetricEngineInner { physical_region_id: RegionId, requests: &mut [(RegionId, RegionPutRequest)], ) -> Result<()> { + let skip_wal = requests + .first() + .is_some_and(|(_, request)| request.skip_wal); + ensure!( + requests + .iter() + .all(|(_, request)| request.skip_wal == skip_wal), + InvalidRequestSnafu { + region_id: physical_region_id, + reason: "inconsistent WAL policy in batch" + } + ); for (logical_region_id, request) in requests { self.verify_rows( *logical_region_id, @@ -194,6 +208,9 @@ impl MetricEngineInner { physical_region_id: RegionId, requests: Vec<(RegionId, RegionPutRequest)>, ) -> Result<(RegionPutRequest, AffectedRows)> { + let skip_wal = requests + .first() + .is_some_and(|(_, request)| request.skip_wal); let total_rows: usize = requests.iter().map(|(_, req)| req.rows.rows.len()).sum(); let mut modified_requests = Vec::with_capacity(requests.len()); let mut total_affected_rows: AffectedRows = 0; @@ -233,6 +250,7 @@ impl MetricEngineInner { } let merged_request = RegionPutRequest { + skip_wal, rows: Rows { schema, rows: merged_rows, @@ -256,6 +274,9 @@ impl MetricEngineInner { data_region_id: RegionId, requests: Vec<(RegionId, RegionPutRequest)>, ) -> Result<(RegionPutRequest, AffectedRows)> { + let skip_wal = requests + .first() + .is_some_and(|(_, request)| request.skip_wal); // Build union schema from all requests let merged_schema = Self::build_union_schema(requests.iter().map(|(_, req)| req.rows.schema.as_slice())); @@ -291,6 +312,7 @@ impl MetricEngineInner { }; let merged_request = RegionPutRequest { + skip_wal, rows: final_rows, hint: None, partition_expr_version: merged_version, @@ -769,25 +791,272 @@ mod tests { use common_function::utils::partition_expr_version; use common_query::prelude::{greptime_native_histogram, greptime_timestamp, greptime_value}; use common_recordbatch::RecordBatches; + use datatypes::arrow::array::{Float64Array, TimestampMillisecondArray}; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema}; use datatypes::value::Value as PartitionValue; use partition::expr::col; use store_api::metadata::ColumnMetadata; use store_api::metric_engine_consts::{ - DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, PRIMARY_KEY_ENCODING, + DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME, METRIC_ENGINE_NAME, + PHYSICAL_TABLE_METADATA_KEY, PRIMARY_KEY_ENCODING, }; use store_api::path_utils::table_dir; use store_api::region_engine::RegionEngine; use store_api::region_request::{ - EnterStagingRequest, RegionRequest, StagingPartitionDirective, + EnterStagingRequest, PathType, RegionCloseRequest, RegionOpenRequest, RegionRequest, + StagingPartitionDirective, }; use store_api::storage::ScanRequest; use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; use super::*; + use crate::engine::MetricEngine; use crate::test_util::{self, TestEnv}; + async fn scan_timestamp_values(engine: &MetricEngine, region_id: RegionId) -> Vec<(i64, f64)> { + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + let mut rows = Vec::new(); + for batch in batches.iter() { + let batch = batch.df_record_batch(); + let timestamp_index = batch.schema().index_of(greptime_timestamp()).unwrap(); + let value_index = batch.schema().index_of(greptime_value()).unwrap(); + let timestamps = batch + .column(timestamp_index) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(value_index) + .as_any() + .downcast_ref::() + .unwrap(); + rows.extend( + timestamps + .values() + .iter() + .copied() + .zip(values.values().iter().copied()), + ); + } + rows.sort_unstable_by_key(|(timestamp, _)| *timestamp); + rows + } + + #[tokio::test] + async fn test_batch_partition_versions() { + check_batch_partition_versions("sparse").await; + check_batch_partition_versions("dense").await; + } + + async fn check_batch_partition_versions(encoding: &str) { + let env = TestEnv::new().await; + let physical_region_id = env.default_physical_region_id(); + let logical_region_id = env.default_logical_region_id(); + env.create_physical_region( + physical_region_id, + &TestEnv::default_table_dir(), + vec![(PRIMARY_KEY_ENCODING.to_string(), encoding.to_string())], + ) + .await; + create_logical_region_with_tags(&env, physical_region_id, logical_region_id, &["job"]) + .await; + let build_requests = |versions: [Option; 3]| { + versions + .into_iter() + .map(|partition_expr_version| { + ( + logical_region_id, + RegionPutRequest { + skip_wal: false, + rows: Rows { + schema: test_util::row_schema_with_tags(&["job"]), + rows: test_util::build_rows(1, 1), + }, + hint: None, + partition_expr_version, + }, + ) + }) + .collect::>() + }; + // Conflicting explicit versions must fail before any data is written. + let err = env + .metric() + .inner + .put_regions_batch_single_physical( + physical_region_id, + build_requests([None, Some(10), Some(11)]), + ) + .await + .unwrap_err(); + assert!( + err.to_string() + .contains("inconsistent partition expr version") + ); + assert!( + scan_timestamp_values(&env.metric(), logical_region_id) + .await + .is_empty() + ); + + for (versions, expected) in [ + ([None, None, None], None), + ([None, Some(7), None], Some(7)), + ([Some(7), None, Some(7)], Some(7)), + ] { + let mut requests = build_requests(versions); + let engine = env.metric(); + engine + .inner + .validate_batch_requests(physical_region_id, &mut requests) + .await + .unwrap(); + let (merged, _) = match encoding { + "sparse" => engine + .inner + .merge_sparse_batch(physical_region_id, requests), + "dense" => engine + .inner + .merge_dense_batch(to_data_region_id(physical_region_id), requests), + _ => unreachable!(), + } + .unwrap(); + assert_eq!(merged.partition_expr_version, expected); + } + } + + #[tokio::test] + async fn test_put_skip_wal_batch_recovery() { + check_put_skip_wal_batch_recovery("sparse", false).await; + check_put_skip_wal_batch_recovery("sparse", true).await; + check_put_skip_wal_batch_recovery("dense", false).await; + check_put_skip_wal_batch_recovery("dense", true).await; + } + + async fn check_put_skip_wal_batch_recovery(encoding: &str, skip_wal: bool) { + let env = TestEnv::new().await; + let engine = env.metric(); + engine.inner.flush_task.stop().await.unwrap(); + let physical_region_id = env.default_physical_region_id(); + let logical_region_id = env.default_logical_region_id(); + env.create_physical_region( + physical_region_id, + &TestEnv::default_table_dir(), + vec![(PRIMARY_KEY_ENCODING.to_string(), encoding.to_string())], + ) + .await; + create_logical_region_with_tags(&env, physical_region_id, logical_region_id, &["job"]) + .await; + let metadata_before = engine.get_metadata(logical_region_id).await.unwrap(); + + let requests = [skip_wal; 3] + .into_iter() + .enumerate() + .map(|(index, skip_wal)| { + let timestamp = index as i64 + 1; + let value = timestamp as f64 * 10.0; + // Every request updates the same key at timestamp zero and + // also inserts a distinct key to verify merge order. + let rows = [0, timestamp] + .into_iter() + .map(|timestamp| Row { + values: vec![ + Value { + value_data: Some(ValueData::TimestampMillisecondValue(timestamp)), + }, + Value { + value_data: Some(ValueData::F64Value(value)), + }, + Value { + value_data: Some(ValueData::StringValue("tag_0".to_string())), + }, + ], + }) + .collect(); + ( + logical_region_id, + RegionPutRequest { + rows: Rows { + schema: test_util::row_schema_with_tags(&["job"]), + rows, + }, + hint: None, + partition_expr_version: None, + skip_wal, + }, + ) + }); + let affected_rows = engine.inner.put_regions_batch(requests).await.unwrap(); + assert_eq!(affected_rows, 6); + assert_eq!( + scan_timestamp_values(&engine, logical_region_id).await, + vec![(0, 30.0), (1, 10.0), (2, 20.0), (3, 30.0)] + ); + + // Neither data nor metadata has an SST to hide missing WAL. + for region_id in [ + to_data_region_id(physical_region_id), + crate::utils::to_metadata_region_id(physical_region_id), + ] { + let stat = env.mito().region_statistic(region_id).unwrap(); + assert!(stat.memtable_size > 0); + assert_eq!(stat.sst_num, 0); + } + engine + .handle_request( + physical_region_id, + RegionRequest::Close(RegionCloseRequest { + flush_on_close: false, + }), + ) + .await + .unwrap(); + + // Recreate the wrapper as well, discarding its metadata cache. + let reopened = MetricEngine::try_new(env.mito(), Default::default()).unwrap(); + reopened.inner.flush_task.stop().await.unwrap(); + reopened + .handle_request( + physical_region_id, + RegionRequest::Open(RegionOpenRequest { + engine: METRIC_ENGINE_NAME.to_string(), + table_dir: TestEnv::default_table_dir(), + path_type: PathType::Bare, + options: [ + (PHYSICAL_TABLE_METADATA_KEY.to_string(), String::new()), + (PRIMARY_KEY_ENCODING.to_string(), encoding.to_string()), + ] + .into_iter() + .collect(), + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + let recovered_metadata = reopened.get_metadata(logical_region_id).await.unwrap(); + assert_eq!( + metadata_before.column_metadatas, + recovered_metadata.column_metadatas + ); + let expected = if skip_wal { + vec![] + } else { + vec![(0, 30.0), (1, 10.0), (2, 20.0), (3, 30.0)] + }; + assert_eq!( + scan_timestamp_values(&reopened, logical_region_id).await, + expected, + "encoding={encoding}, skip_wal={skip_wal}" + ); + } + fn assert_merged_schema(rows: &Rows, expect_sparse: bool) { let column_names: HashSet = rows .schema @@ -874,6 +1143,45 @@ mod tests { .unwrap() } + fn check_batch_merge_wal_policy( + env: &TestEnv, + physical_region_id: RegionId, + mut requests: Vec<(RegionId, RegionPutRequest)>, + expect_sparse: bool, + skip_wal: bool, + ) { + for (_, request) in &mut requests { + request.skip_wal = skip_wal; + } + let (merged_request, affected_rows) = if expect_sparse { + let (merged_request, affected_rows) = env + .metric() + .inner + .merge_sparse_batch(physical_region_id, requests) + .unwrap(); + let hint = merged_request + .hint + .as_ref() + .expect("missing sparse write hint"); + assert_eq!( + hint.primary_key_encoding, + PrimaryKeyEncodingProto::Sparse as i32 + ); + (merged_request, affected_rows) + } else { + let (merged_request, affected_rows) = env + .metric() + .inner + .merge_dense_batch(to_data_region_id(physical_region_id), requests) + .unwrap(); + assert!(merged_request.hint.is_none()); + (merged_request, affected_rows) + }; + assert_merged_schema(&merged_request.rows, expect_sparse); + assert_eq!(merged_request.skip_wal, skip_wal); + assert_eq!(affected_rows, 5); + } + async fn run_batch_write_with_schema_variants( env: &TestEnv, physical_region_id: RegionId, @@ -921,6 +1229,7 @@ mod tests { ( logical_region_1, RegionPutRequest { + skip_wal: false, rows: Rows { schema: schema_1.clone(), rows: rows_1, @@ -932,6 +1241,7 @@ mod tests { ( logical_region_2, RegionPutRequest { + skip_wal: false, rows: Rows { schema: schema_2.clone(), rows: rows_2, @@ -943,32 +1253,41 @@ mod tests { ] }; - let merged_request = if expect_sparse { - let (merged_request, _) = env - .metric() - .inner - .merge_sparse_batch(physical_region_id, build_requests()) - .unwrap(); - let hint = merged_request - .hint - .as_ref() - .expect("missing sparse write hint"); - assert_eq!( - hint.primary_key_encoding, - PrimaryKeyEncodingProto::Sparse as i32 - ); - merged_request - } else { - let (merged_request, _) = env - .metric() - .inner - .merge_dense_batch(data_region_id, build_requests()) - .unwrap(); - assert!(merged_request.hint.is_none()); - merged_request - }; + check_batch_merge_wal_policy( + env, + physical_region_id, + build_requests(), + expect_sparse, + false, + ); + check_batch_merge_wal_policy( + env, + physical_region_id, + build_requests(), + expect_sparse, + true, + ); - assert_merged_schema(&merged_request.rows, expect_sparse); + for policies in [[false, true], [true, false]] { + let mut mixed_requests = build_requests(); + for ((_, request), skip_wal) in mixed_requests.iter_mut().zip(policies) { + request.skip_wal = skip_wal; + } + let err = env + .metric() + .inner + .put_regions_batch(mixed_requests.into_iter()) + .await + .unwrap_err(); + assert!(err.to_string().contains("inconsistent WAL policy in batch")); + for logical_region_id in [logical_region_1, logical_region_2] { + assert!( + scan_timestamp_values(&env.metric(), logical_region_id) + .await + .is_empty() + ); + } + } let affected_rows = env .metric() @@ -1095,6 +1414,7 @@ mod tests { let schema = test_util::row_schema_with_tags(&["job"]); let rows = test_util::build_rows(1, 5); let request = RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: Rows { schema, rows }, hint: None, partition_expr_version: None, @@ -1170,6 +1490,7 @@ mod tests { let schema = test_util::row_schema_with_tags(columns); let rows = test_util::build_rows(3, 100); let request = RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: Rows { schema, rows }, hint: None, partition_expr_version: None, @@ -1193,6 +1514,7 @@ mod tests { let schema = test_util::row_schema_with_tags(&["abc"]); let rows = test_util::build_rows(1, 100); let request = RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: Rows { schema, rows }, hint: None, partition_expr_version: None, @@ -1214,6 +1536,7 @@ mod tests { let schema = test_util::row_schema_with_tags(&["def"]); let rows = test_util::build_rows(1, 100); let request = RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: Rows { schema, rows }, hint: None, partition_expr_version: None, @@ -1274,6 +1597,7 @@ mod tests { ( logical_region_1, RegionPutRequest { + skip_wal: false, rows: Rows { schema: schema.clone(), rows: rows1, @@ -1285,6 +1609,7 @@ mod tests { ( logical_region_2, RegionPutRequest { + skip_wal: false, rows: Rows { schema: schema.clone(), rows: rows2, @@ -1296,6 +1621,7 @@ mod tests { ( logical_region_3, RegionPutRequest { + skip_wal: false, rows: Rows { schema: schema.clone(), rows: rows3, @@ -1347,6 +1673,7 @@ mod tests { ( logical_region_1, RegionPutRequest { + skip_wal: false, rows: Rows { schema: schema.clone(), rows: test_util::build_rows(1, 3), @@ -1358,6 +1685,7 @@ mod tests { ( nonexistent_region, RegionPutRequest { + skip_wal: false, rows: Rows { schema: schema.clone(), rows: test_util::build_rows(1, 2), @@ -1369,6 +1697,7 @@ mod tests { ( logical_region_2, RegionPutRequest { + skip_wal: false, rows: Rows { schema: schema.clone(), rows: test_util::build_rows(1, 5), @@ -1407,6 +1736,7 @@ mod tests { let requests = vec![( physical_region_id, RegionPutRequest { + skip_wal: false, rows: Rows { schema, rows: test_util::build_rows(1, 1), @@ -1441,6 +1771,7 @@ mod tests { ( logical_region_id, RegionPutRequest { + skip_wal: false, rows: Rows { schema: schema.clone(), rows: test_util::build_rows(1, 1), @@ -1452,6 +1783,7 @@ mod tests { ( physical_region_id, RegionPutRequest { + skip_wal: false, rows: Rows { schema, rows: test_util::build_rows(1, 1), @@ -1487,6 +1819,7 @@ mod tests { let requests = vec![( logical_region_id, RegionPutRequest { + skip_wal: false, rows: Rows { schema, rows: test_util::build_rows(1, 5), @@ -1565,6 +1898,7 @@ mod tests { .handle_request( logical_region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows, hint: None, partition_expr_version: Some(1), @@ -1607,6 +1941,7 @@ mod tests { .handle_request( logical_region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: rows.clone(), hint: None, partition_expr_version: Some(expected_version.wrapping_add(1)), @@ -1621,6 +1956,7 @@ mod tests { .handle_request( logical_region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: rows.clone(), hint: None, partition_expr_version: None, @@ -1635,6 +1971,7 @@ mod tests { .handle_request( logical_region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows, hint: None, partition_expr_version: Some(expected_version), @@ -1703,6 +2040,7 @@ mod tests { .handle_request( logical_region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: Rows { schema, rows }, hint: None, partition_expr_version: None, @@ -1759,6 +2097,7 @@ mod tests { .handle_request( logical_region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: Rows { schema, rows }, hint: None, partition_expr_version: None, @@ -1813,6 +2152,7 @@ mod tests { .handle_request( logical_region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: Rows { schema, rows }, hint: None, partition_expr_version: None, diff --git a/src/metric-engine/src/engine/read.rs b/src/metric-engine/src/engine/read.rs index d1afce5e22d..8d8d32fcc6c 100644 --- a/src/metric-engine/src/engine/read.rs +++ b/src/metric-engine/src/engine/read.rs @@ -500,6 +500,7 @@ mod test { let schema = test_util::row_schema_with_tags(&["job"]); let put = |rows| { RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: Rows { schema: schema.clone(), rows: test_util::build_rows(1, rows), diff --git a/src/metric-engine/src/metadata_region.rs b/src/metric-engine/src/metadata_region.rs index ec67ac70101..18fa51f3597 100644 --- a/src/metric-engine/src/metadata_region.rs +++ b/src/metric-engine/src/metadata_region.rs @@ -585,6 +585,8 @@ impl MetadataRegion { }; RegionPutRequest { + // Metadata must remain recoverable regardless of the user write policy. + skip_wal: false, rows, hint: None, partition_expr_version: None, @@ -819,6 +821,37 @@ mod test { use crate::test_util::TestEnv; use crate::utils::to_metadata_region_id; + #[test] + fn test_metadata_put_always_writes_wal() { + // Both metadata put paths use this constructor rather than forwarding + // a user insert request, so its WAL policy must always be independent. + for entries in [ + vec![], + vec![("region", "")], + vec![("region", ""), ("column", "metadata")], + ] { + let request = MetadataRegion::build_put_request_from_iter( + entries + .iter() + .map(|(key, value)| (key.to_string(), value.to_string())), + ); + assert!(!request.skip_wal); + assert!(request.hint.is_none()); + assert!(request.partition_expr_version.is_none()); + let expected_rows = entries + .into_iter() + .map(|(key, value)| { + row(vec![ + ValueData::TimestampMillisecondValue(0), + ValueData::StringValue(key.to_string()), + ValueData::StringValue(value.to_string()), + ]) + }) + .collect::>(); + assert_eq!(request.rows.rows, expected_rows); + } + } + #[test] fn test_concat_table_key() { let region_id = RegionId::new(1234, 7844); diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md index b22ee758bb9..faa581331b4 100644 --- a/src/mito2/AGENTS.md +++ b/src/mito2/AGENTS.md @@ -26,7 +26,7 @@ snapshot isolation). It implements the `RegionEngine` trait from `store-api`. | `compaction` | `src/mito2/src/compaction/` | Compaction scheduler (`scheduler.rs` + `scheduler/`), TWCS picker, strict-window manual picker, compactor, memory control | | `access_layer` | `src/mito2/src/access_layer.rs` | SST read/write over the object store | | `sst` | `src/mito2/src/sst/` | Parquet format, file metadata, index layout | -| `series_index` | `src/mito2/src/series_index/` | Series index writer/searcher, catalogs, immutable snapshots, file lifecycle, and background maintenance | +| `series_index` | `src/mito2/src/series_index/` | Series index writer/searcher, bucket planning and SST builders, catalogs, immutable snapshots, file lifecycle, and background maintenance | | `read` | `src/mito2/src/read/` | `ScanRegion`, merge, dedup, projection, streaming | | `manifest` | `src/mito2/src/manifest/` | `RegionManifestManager`, manifest actions/edits | | `cache` | `src/mito2/src/cache.rs` | Write/file/page caches | @@ -62,6 +62,8 @@ filtered `RecordBatch` stream. - **Manifest format** (`manifest/action.rs`): affects crash recovery and follower replay. Keep it backward compatible. - **SST/Parquet layout** (`sst/`): readers must stay compatible with existing files. +- **Series-index coverage** (`series_index/catalog.rs`): `SeriesIndexEntry` stores + compaction-window width and SST summaries keyed by aligned start in both catalogs and Parquet footers. - **Request types** (`request.rs`): usually tied to proto definitions consumed by `datanode`. - **WAL/memtable encoding** (`wal/`, `memtable/`): breaks replay if changed incompatibly. diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index 85808842741..14c629a01b3 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::path::Path; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -655,14 +656,17 @@ impl TempFileCleaner { } pub(crate) async fn new_fs_cache_store(root: &str) -> Result { - let atomic_write_dir = join_dir(root, ATOMIC_WRITE_DIR); - clean_dir(&atomic_write_dir).await?; + // Preserve native filesystem prefixes such as Windows UNC shares. + let atomic_write_dir = Path::new(root).join(ATOMIC_WRITE_DIR); + clean_dir(&atomic_write_dir.to_string_lossy()).await?; // Compatible code. Remove this after a major release. let old_atomic_temp_dir = join_dir(root, OLD_ATOMIC_WRITE_DIR); clean_dir(&old_atomic_temp_dir).await?; - let builder = Fs::default().root(root).atomic_write_dir(&atomic_write_dir); + let builder = Fs::default() + .root(root) + .atomic_write_dir(&atomic_write_dir.to_string_lossy()); let store = ObjectStore::new(builder).context(OpenDalSnafu)?; Ok(with_instrument_layers(store, false)) @@ -759,3 +763,37 @@ impl FilePathProvider for RegionFilePathFactory { location::sst_file_path(&self.table_dir, file_id, self.path_type) } } + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + + use super::*; + + #[tokio::test] + async fn test_new_fs_cache_store() { + let root = common_test_util::temp_dir::create_temp_dir("fs-cache-store"); + let dirs = [ + root.path().join(ATOMIC_WRITE_DIR), + PathBuf::from(join_dir( + root.path().to_str().unwrap(), + OLD_ATOMIC_WRITE_DIR, + )), + ]; + for dir in &dirs { + tokio::fs::create_dir_all(dir).await.unwrap(); + tokio::fs::write(dir.join("stale"), b"stale").await.unwrap(); + } + let store = new_fs_cache_store(root.path().to_str().unwrap()) + .await + .unwrap(); + for dir in &dirs { + assert!(!dir.join("stale").exists()); + } + store.write("index", "contents").await.unwrap(); + assert_eq!( + tokio::fs::read(root.path().join("index")).await.unwrap(), + b"contents" + ); + } +} diff --git a/src/mito2/src/compaction/compactor.rs b/src/mito2/src/compaction/compactor.rs index 8ad1cdf0b22..fb2172b4875 100644 --- a/src/mito2/src/compaction/compactor.rs +++ b/src/mito2/src/compaction/compactor.rs @@ -655,6 +655,7 @@ where write_buffer_size: compaction_region.engine_config.sst_write_buffer_size, max_file_size: picker_output.max_file_size, row_group_size: compaction_region.region_options.row_group_size(), + float_field_encoding: compaction_region.region_options.float_field_encoding, }; let merger = self.merger.clone(); let compaction_region = compaction_region.clone(); diff --git a/src/mito2/src/compaction/json2.rs b/src/mito2/src/compaction/json2.rs index 71b571c586d..8eec0e11520 100644 --- a/src/mito2/src/compaction/json2.rs +++ b/src/mito2/src/compaction/json2.rs @@ -52,6 +52,8 @@ pub(crate) type Json2RewritePlans = HashMap; #[derive(Clone)] struct Json2LeafPathStats { rows: u64, + /// Number of input schemas that explicitly represent this leaf. + sources: usize, data_type: JsonNativeType, is_type_conflicted: bool, } @@ -59,8 +61,8 @@ struct Json2LeafPathStats { /// Builds the JSON2 rewrite plans for a compaction. /// /// Type hints from current region metadata are always retained. Existing explicit dynamic paths -/// from all input SST schemas are ranked once to produce a fixed layout; paths found only in a v2 -/// remainder are deliberately not promoted. [`rewrite_json2_batch`] decodes inputs and rewrites +/// shared by every input SST schema are ranked once to produce a fixed layout; paths that may +/// reside in a v2 remainder are deliberately not promoted. [`rewrite_json2_batch`] decodes inputs and rewrites /// them according to these plans. Non-JSON2 columns are omitted from the returned map. /// /// Returns an error when a JSON2 column has invalid or missing extension metadata, or when its @@ -144,6 +146,14 @@ pub(crate) fn collect_json2_rewrite_plans( collect_json2_path_stats(field, *rows, &hint_paths, &mut stats)?; } + // A leaf absent from an input's explicit schema may have arbitrary values + // (including scalar ancestors) in its remainder. Only common explicit leaves + // are safe to promote without inspecting rows. Keep unsafe paths opaque for + // the entire output SST so existing readers never miss remainder values when + // projecting an explicit leaf. + for stat in stats.values_mut() { + stat.is_type_conflicted |= stat.sources != schemas.len(); + } let mut hints = settings.type_hints().to_vec(); hints.extend(select_dynamic_hints(settings, &hint_paths, &stats)); let target_layout = JsonSettings::try_new(hints, Some(0)).context(DataTypeMismatchSnafu)?; @@ -187,12 +197,14 @@ fn collect_json2_path_stats<'a>( path, Json2LeafPathStats { rows, + sources: 1, data_type, is_type_conflicted: false, }, ); continue; }; + stat.sources += 1; if stat.data_type != data_type { stat.is_type_conflicted = true; } else { @@ -368,6 +380,7 @@ mod tests { )?; let stat = |rows, data_type, is_type_conflicted| Json2LeafPathStats { rows, + sources: 1, data_type, is_type_conflicted, }; diff --git a/src/mito2/src/compaction/window.rs b/src/mito2/src/compaction/window.rs index a9f3f021b84..b9053466444 100644 --- a/src/mito2/src/compaction/window.rs +++ b/src/mito2/src/compaction/window.rs @@ -337,6 +337,7 @@ mod tests { primary_key_encoding: None, write_buffer_size: None, preserve_row_sequence: false, + float_field_encoding: Default::default(), }, compaction_time_window: None, } diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index 6ec790c52b0..667f156fb22 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -89,11 +89,9 @@ pub struct MitoConfig { // Background job configs: /// Max number of running background index build jobs (default: 1/8 of cpu cores). pub max_background_index_builds: usize, - // TODO: Document both series-index settings in the example configs and regenerate configuration - // docs before exposing the feature. - /// Under development; do not enable. Root directory for loading series indexes, currently stored - /// on the local filesystem. Empty disables the feature. Relative paths resolve under `data_home`. - pub experimental_series_index_root: String, + /// Under development; do not enable. Whether to enable series indexes (default false). + /// Indexes are stored on the local filesystem under `{data_home}/series_index`. + pub experimental_enable_series_index: bool, /// Interval between series-index maintenance runs (default 5 min). Zero uses the default. #[serde(with = "humantime_serde")] pub experimental_series_index_maintenance_interval: Duration, @@ -214,7 +212,7 @@ impl Default for MitoConfig { experimental_manifest_keep_removed_file_ttl: Duration::from_secs(60 * 60), compress_manifest: false, max_background_index_builds: divide_num_cpus(8), - experimental_series_index_root: String::new(), + experimental_enable_series_index: false, experimental_series_index_maintenance_interval: DEFAULT_SERIES_INDEX_MAINTENANCE_INTERVAL, max_background_flushes: divide_num_cpus(2), @@ -312,15 +310,6 @@ impl MitoConfig { DEFAULT_SERIES_INDEX_MAINTENANCE_INTERVAL; } - if !self.experimental_series_index_root.trim().is_empty() - && Path::new(&self.experimental_series_index_root).is_relative() - { - self.experimental_series_index_root = Path::new(data_home) - .join(&self.experimental_series_index_root) - .display() - .to_string(); - } - if self.global_write_buffer_reject_size <= self.global_write_buffer_size { self.global_write_buffer_reject_size = self.global_write_buffer_size * 2; warn!( @@ -432,36 +421,6 @@ mod tests { assert_eq!(ReadableSize::mb(128), config.prefilter_result_cache_size); } - #[test] - fn test_series_index_config() { - assert!( - MitoConfig::default() - .experimental_series_index_root - .is_empty() - ); - let mut config: MitoConfig = toml::from_str( - "experimental_series_index_root = 'indexes' - experimental_series_index_maintenance_interval = '30s'", - ) - .unwrap(); - config.sanitize("/data").unwrap(); - assert_eq!(config.experimental_series_index_root, "/data/indexes"); - assert_eq!( - config.experimental_series_index_maintenance_interval, - Duration::from_secs(30) - ); - let restored: MitoConfig = toml::from_str(&toml::to_string(&config).unwrap()).unwrap(); - assert_eq!(config, restored); - - let mut config: MitoConfig = - toml::from_str("experimental_series_index_maintenance_interval = '0s'").unwrap(); - config.sanitize("/data").unwrap(); - assert_eq!( - config.experimental_series_index_maintenance_interval, - MitoConfig::default().experimental_series_index_maintenance_interval - ); - } - #[test] fn test_experimental_series_scan_v2_config() { assert!(MitoConfig::default().experimental_series_scan_v2); diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs index 1c674abadbf..348ec4c1222 100644 --- a/src/mito2/src/engine.rs +++ b/src/mito2/src/engine.rs @@ -225,6 +225,7 @@ impl<'a, S: LogStore> MitoEngineBuilder<'a, S> { // so the engine (and thus the GC worker) can fire `on_region_gc`. let region_hook = self.plugins.get::(); let workers = WorkerGroup::start( + self.data_home, config.clone(), self.log_store.clone(), self.object_store_manager, @@ -1620,6 +1621,7 @@ impl MitoEngine { Ok(MitoEngine { inner: Arc::new(EngineInner { workers: WorkerGroup::start_for_test( + data_home, config.clone(), log_store, object_store_manager, diff --git a/src/mito2/src/engine/alter_test.rs b/src/mito2/src/engine/alter_test.rs index 736f32cc8ae..fd74c4b498a 100644 --- a/src/mito2/src/engine/alter_test.rs +++ b/src/mito2/src/engine/alter_test.rs @@ -3003,6 +3003,7 @@ fn build_schema_v0_bulk_request(region_id: RegionId) -> RegionBulkInsertsRequest let raw_data = encode_arrow_ipc(&payload); RegionBulkInsertsRequest { + skip_wal: false, region_id, payload, raw_data, @@ -3411,6 +3412,7 @@ async fn test_alter_time_index_widen_sparse_compaction() { }; let put_sparse = |rows| { RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows, hint: Some(WriteHint { primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse.into(), diff --git a/src/mito2/src/engine/apply_staging_manifest_test.rs b/src/mito2/src/engine/apply_staging_manifest_test.rs index 5e7f77257a3..e6ac2bc24b4 100644 --- a/src/mito2/src/engine/apply_staging_manifest_test.rs +++ b/src/mito2/src/engine/apply_staging_manifest_test.rs @@ -952,6 +952,7 @@ async fn test_apply_staging_manifest_preserves_unflushed_memtable_with_format(fl .handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: unflushed_rows, hint: None, partition_expr_version: Some(expected_version), diff --git a/src/mito2/src/engine/basic_test.rs b/src/mito2/src/engine/basic_test.rs index edaa97ff585..a9ee4054158 100644 --- a/src/mito2/src/engine/basic_test.rs +++ b/src/mito2/src/engine/basic_test.rs @@ -29,6 +29,7 @@ use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; use futures::TryStreamExt; use itertools::Itertools; +use parquet::basic::{Encoding, Type as PhysicalType}; use rstest::rstest; use rstest_reuse::{self, apply}; use store_api::metadata::ColumnMetadata; @@ -732,6 +733,7 @@ async fn test_absent_and_invalid_columns_with_format(flat_format: bool) { .handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows, hint: None, partition_expr_version: None, @@ -825,7 +827,9 @@ async fn test_engine_with_write_cache_with_format(flat_format: bool) { let engine = env.create_engine(mito_config).await; let region_id = RegionId::new(1, 1); - let request = CreateRequestBuilder::new().build(); + let request = CreateRequestBuilder::new() + .insert_option("experimental_sst_float_field_encoding", "byte_stream_split") + .build(); let column_schemas = rows_schema(&request); engine @@ -841,6 +845,38 @@ async fn test_engine_with_write_cache_with_format(flat_format: bool) { flush_region(&engine, region_id, None).await; + let region = engine.get_region(region_id).unwrap(); + let file = region + .version() + .ssts + .levels() + .iter() + .flat_map(|level| level.files.values()) + .next() + .expect("write-cache flushed SST") + .clone(); + let reader = region + .access_layer + .read_sst(file) + .build() + .await + .unwrap() + .expect("write-cache flushed SST reader"); + assert!( + reader + .parquet_metadata() + .row_groups() + .iter() + .flat_map(|row_group| row_group.columns()) + .any(|column| { + column.column_path().string() == "field_0" + && column.column_type() == PhysicalType::DOUBLE + && column + .encodings() + .any(|encoding| encoding == Encoding::BYTE_STREAM_SPLIT) + }) + ); + let request = ScanRequest::default(); let scanner = engine.scanner(region_id, request).await.unwrap(); @@ -1320,3 +1356,388 @@ async fn test_all_index_metas_list_all_types_with_format(flat_format: bool, expe assert_eq!(expect_format, debug_format); } + +#[tokio::test] +async fn test_request_skip_wal_recovery() { + check_request_skip_wal_recovery(false, false).await; + check_request_skip_wal_recovery(false, true).await; + check_request_skip_wal_recovery(true, false).await; + check_request_skip_wal_recovery(true, true).await; +} + +async fn check_request_skip_wal_recovery(flat_format: bool, skip_wal: bool) { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let table_dir = request.table_dir.clone(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let affected = engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema, + rows: build_rows_for_key("a", 0, 4, 0), + }, + hint: None, + partition_expr_version: None, + skip_wal, + }), + ) + .await + .unwrap(); + assert_eq!(affected.affected_rows, 4); + let current = engine + .get_region(region_id) + .unwrap() + .version_control + .current(); + assert_eq!(current.committed_sequence, 4); + assert_eq!(current.last_entry_id, u64::from(!skip_wal)); + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let before = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!(before.iter().map(|b| b.num_rows()).sum::(), 4); + + reopen_region(&engine, region_id, table_dir, false, HashMap::new()).await; + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let after = RecordBatches::try_collect(stream).await.unwrap(); + assert_eq!( + after.iter().map(|b| b.num_rows()).sum::(), + if skip_wal { 0 } else { 4 } + ); +} + +#[tokio::test] +async fn test_request_skip_wal_flush_watermarks() { + check_request_skip_wal_flush_watermarks(false).await; + check_request_skip_wal_flush_watermarks(true).await; +} + +async fn check_request_skip_wal_flush_watermarks(flat_format: bool) { + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + // Establish a nonzero flushed baseline, skip one write, then resume WAL. + for (round, skip_wal) in [false, true, false].into_iter().enumerate() { + let region = engine.get_region(region_id).unwrap(); + let before = region.version_control.current(); + let flushed_entry_id = engine + .region_statistic(region_id) + .unwrap() + .manifest + .data_flushed_entry_id(); + let affected = engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + rows: Rows { + schema: schema.clone(), + rows: build_rows_for_key("a", 0, 4, 0), + }, + hint: None, + partition_expr_version: None, + skip_wal, + }), + ) + .await + .unwrap(); + assert_eq!(affected.affected_rows, 4); + let after_write = region.version_control.current(); + let expected_entry_id = before.last_entry_id + u64::from(!skip_wal); + assert_eq!(after_write.last_entry_id, expected_entry_id); + assert_eq!( + after_write.committed_sequence, + before.committed_sequence + 4 + ); + assert_eq!(after_write.version.flushed_entry_id, flushed_entry_id); + assert_eq!( + engine + .region_statistic(region_id) + .unwrap() + .manifest + .data_flushed_entry_id(), + flushed_entry_id + ); + assert_eq!( + after_write.version.flushed_sequence, + before.version.flushed_sequence + ); + + flush_region(&engine, region_id, None).await; + let after_flush = region.version_control.current(); + assert_eq!(after_flush.last_entry_id, expected_entry_id); + assert_eq!(after_flush.version.flushed_sequence, (round as u64 + 1) * 4); + assert_eq!( + engine + .region_statistic(region_id) + .unwrap() + .manifest + .data_flushed_entry_id(), + expected_entry_id + ); + if skip_wal { + assert_eq!(expected_entry_id, flushed_entry_id); + } else { + assert!(expected_entry_id > flushed_entry_id); + } + } +} + +/// Returns `(committed_sequence, last_entry_id)` from one region version snapshot. +/// Returns `None` if the region is not open. +fn region_write_watermarks(engine: &MitoEngine, region_id: RegionId) -> Option<(u64, u64)> { + let region = engine.find_region(region_id)?; + let current = region.version_control.current(); + Some((current.committed_sequence, current.last_entry_id)) +} + +#[tokio::test] +async fn test_request_skip_wal_mixed_batch_recovery() { + check_request_skip_wal_mixed_batch_recovery(false, false, false).await; + check_request_skip_wal_mixed_batch_recovery(false, false, true).await; + check_request_skip_wal_mixed_batch_recovery(false, true, false).await; + check_request_skip_wal_mixed_batch_recovery(false, true, true).await; + check_request_skip_wal_mixed_batch_recovery(true, false, false).await; + check_request_skip_wal_mixed_batch_recovery(true, false, true).await; + check_request_skip_wal_mixed_batch_recovery(true, true, false).await; + check_request_skip_wal_mixed_batch_recovery(true, true, true).await; +} + +async fn check_request_skip_wal_mixed_batch_recovery( + flat_format: bool, + skip_wal: bool, + flush_before_reopen: bool, +) { + use crate::region_write_ctx::RegionWriteCtx; + use crate::request::OptionOutputTx; + use crate::test_util::LogStoreImpl; + use crate::wal::Wal; + + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + default_flat_format: flat_format, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new().build(); + let table_dir = request.table_dir.clone(); + let schema = rows_schema(&request); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + let LogStoreImpl::RaftEngine(store) = env.get_log_store().unwrap() else { + panic!("expected the default local WAL"); + }; + let wal = Wal::new(store); + // Assemble one real worker write context deterministically, instead + // of relying on concurrently submitted requests landing in one batch. + let mut ctx = RegionWriteCtx::new( + region_id, + ®ion.version_control, + region.provider.clone(), + None, + ); + let mut receivers = Vec::with_capacity(4); + for (index, (skip, timestamp)) in [(false, 0), (skip_wal, 0), (false, 0), (skip_wal, 1)] + .into_iter() + .enumerate() + { + let (tx, rx) = tokio::sync::oneshot::channel(); + receivers.push(rx); + ctx.push_mutation( + api::v1::OpType::Put as i32, + Some(Rows { + schema: schema.clone(), + // Timestamp zero ends with a WAL-backed update; timestamp one + // ends with a skipped update. Grouping mutations must not + // change either winner chosen by the original sequences. + rows: build_rows_for_key("a", timestamp, timestamp + 1, index * 10) + .into_iter() + .chain(build_rows_for_key( + "a", + index + 1, + index + 2, + index * 10 + 1, + )) + .collect(), + }), + None, + OptionOutputTx::from(tx), + None, + skip, + ); + } + let mut writer = wal.writer(); + ctx.add_wal_entry(&mut writer).unwrap(); + let response = writer.write_to_wal().await.unwrap(); + assert_eq!(response.last_entry_ids.get(®ion_id), Some(&1)); + ctx.write_memtable().await; + ctx.publish_sequence_and_entry_id(); + drop(ctx); + for rx in receivers { + assert_eq!(rx.await.unwrap().unwrap(), 2); + } + assert_eq!(region_write_watermarks(&engine, region_id), Some((8, 1))); + assert_eq!( + request_skip_wal_values(&engine, region_id).await, + vec![ + (0, 20.0), + (1000, 30.0), + (2000, 11.0), + (3000, 21.0), + (4000, 31.0) + ] + ); + + // Inspect persisted WAL mutations, not just an encoder or counter. + let mut reader = wal.wal_entry_reader(®ion.provider, region_id, None); + let entries = reader + .read(®ion.provider, 1) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].0, 1); + assert_eq!( + entries[0] + .1 + .mutations + .iter() + .map(|m| m.sequence) + .collect::>(), + if skip_wal { + vec![1, 5] + } else { + vec![1, 3, 5, 7] + } + ); + assert!(entries[0].1.bulk_entries.is_empty()); + drop(region); + if flush_before_reopen { + flush_region(&engine, region_id, None).await; + let current = engine + .get_region(region_id) + .unwrap() + .version_control + .current(); + assert_eq!( + ( + current.version.flushed_sequence, + current.version.flushed_entry_id + ), + (8, 1) + ); + } + + // A no-flush close discards memtables. Only WAL-backed rows recover + // unless an explicit flush has already persisted all requests. + reopen_region(&engine, region_id, table_dir, true, HashMap::new()).await; + let loses_skipped_rows = skip_wal && !flush_before_reopen; + let mut expected_values = if loses_skipped_rows { + vec![(0, 20.0), (1000, 1.0), (3000, 21.0)] + } else { + vec![ + (0, 20.0), + (1000, 30.0), + (2000, 11.0), + (3000, 21.0), + (4000, 31.0), + ] + }; + assert_eq!( + request_skip_wal_values(&engine, region_id).await, + expected_values + ); + let recovered_sequence = if loses_skipped_rows { 6 } else { 8 }; + assert_eq!( + region_write_watermarks(&engine, region_id), + Some((recovered_sequence, 1)) + ); + + // A subsequent default request still writes WAL, even after a + // trailing skipped request or a flush with sequence/entry-id gaps. + put_rows( + &engine, + region_id, + Rows { + schema, + rows: build_rows_for_key("a", 8, 9, 8), + }, + ) + .await; + assert_eq!( + region_write_watermarks(&engine, region_id), + Some((recovered_sequence + 1, 2)) + ); + expected_values.push((8000, 8.0)); + assert_eq!( + request_skip_wal_values(&engine, region_id).await, + expected_values + ); +} + +async fn request_skip_wal_values(engine: &MitoEngine, region_id: RegionId) -> Vec<(i64, f64)> { + use datatypes::arrow::array::{Float64Array, TimestampMillisecondArray}; + + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + let batches = RecordBatches::try_collect(stream).await.unwrap(); + let mut values = batches + .iter() + .flat_map(|batch| { + let batch = batch.df_record_batch(); + let timestamps = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + timestamps + .values() + .iter() + .copied() + .zip(values.values().iter().copied()) + }) + .collect::>(); + values.sort_unstable_by_key(|(timestamp, _)| *timestamp); + values +} diff --git a/src/mito2/src/engine/compaction_test.rs b/src/mito2/src/engine/compaction_test.rs index c67afbf8c6f..e427f3f34f5 100644 --- a/src/mito2/src/engine/compaction_test.rs +++ b/src/mito2/src/engine/compaction_test.rs @@ -26,6 +26,7 @@ use common_recordbatch::{RecordBatches, SendableRecordBatchStream}; use common_time::Timestamp; use datatypes::arrow::array::AsArray; use datatypes::arrow::datatypes::TimestampMillisecondType; +use parquet::basic::{Encoding, Type as PhysicalType}; use store_api::region_engine::{RegionEngine, RegionRole}; use store_api::region_request::AlterKind::SetRegionOptions; use store_api::region_request::{ @@ -954,6 +955,7 @@ async fn test_compaction_region_with_format(flat_format: bool) { let request = CreateRequestBuilder::new() .insert_option("compaction.type", "twcs") + .insert_option("experimental_sst_float_field_encoding", "byte_stream_split") .build(); let column_schemas = request @@ -974,6 +976,35 @@ async fn test_compaction_region_with_format(flat_format: bool) { compact(&engine, region_id).await; + let region = engine.get_region(region_id).unwrap(); + let file = region.version().ssts.levels()[1] + .files + .values() + .next() + .expect("compaction output SST") + .clone(); + let reader = region + .access_layer + .read_sst(file) + .build() + .await + .unwrap() + .expect("compaction output SST reader"); + assert!( + reader + .parquet_metadata() + .row_groups() + .iter() + .flat_map(|row_group| row_group.columns()) + .any(|column| { + column.column_path().string() == "field_0" + && column.column_type() == PhysicalType::DOUBLE + && column + .encodings() + .any(|encoding| encoding == Encoding::BYTE_STREAM_SPLIT) + }) + ); + let scanner = engine .scanner(region_id, ScanRequest::default()) .await diff --git a/src/mito2/src/engine/edit_region_test.rs b/src/mito2/src/engine/edit_region_test.rs index 992060d75d7..39e688c5c4d 100644 --- a/src/mito2/src/engine/edit_region_test.rs +++ b/src/mito2/src/engine/edit_region_test.rs @@ -286,6 +286,7 @@ async fn test_write_during_region_editing_is_queued() { .handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows, hint: None, partition_expr_version: None, @@ -400,6 +401,7 @@ async fn test_stalled_write_fails_fast_if_region_closed_during_editing() { .handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows, hint: None, partition_expr_version: None, @@ -698,6 +700,7 @@ fn build_bulk_insert_request( let (schema, record_batch) = encode_to_flight_data(payload.clone()); RegionBulkInsertsRequest { + skip_wal: false, region_id, payload, raw_data: ArrowIpc { diff --git a/src/mito2/src/engine/flush_test.rs b/src/mito2/src/engine/flush_test.rs index b71ce9f9fc0..2b1e6f0cc2d 100644 --- a/src/mito2/src/engine/flush_test.rs +++ b/src/mito2/src/engine/flush_test.rs @@ -27,6 +27,7 @@ use common_base::readable_size::ReadableSize; use common_recordbatch::RecordBatches; use common_time::util::current_time_millis; use common_wal::options::{KafkaWalOptions, WAL_OPTIONS_KEY, WalOptions}; +use parquet::basic::{Encoding, Type as PhysicalType}; use rstest::rstest; use rstest_reuse::{self, apply}; use store_api::ManifestVersion; @@ -279,7 +280,9 @@ async fn test_manual_flush_with_format(flat_format: bool) { ) .await; - let request = CreateRequestBuilder::new().build(); + let request = CreateRequestBuilder::new() + .insert_option("experimental_sst_float_field_encoding", "byte_stream_split") + .build(); let column_schemas = rows_schema(&request); engine @@ -295,6 +298,38 @@ async fn test_manual_flush_with_format(flat_format: bool) { flush_region(&engine, region_id, None).await; + let region = engine.get_region(region_id).unwrap(); + let file = region + .version() + .ssts + .levels() + .iter() + .flat_map(|level| level.files.values()) + .next() + .expect("flushed SST") + .clone(); + let reader = region + .access_layer + .read_sst(file) + .build() + .await + .unwrap() + .expect("flushed SST reader"); + assert!( + reader + .parquet_metadata() + .row_groups() + .iter() + .flat_map(|row_group| row_group.columns()) + .any(|column| { + column.column_path().string() == "field_0" + && column.column_type() == PhysicalType::DOUBLE + && column + .encodings() + .any(|encoding| encoding == Encoding::BYTE_STREAM_SPLIT) + }) + ); + let request = ScanRequest::default(); let scanner = engine.scanner(region_id, request).await.unwrap(); assert_eq!(0, scanner.num_memtables()); @@ -740,6 +775,7 @@ async fn test_region_write_buffer_does_not_stall_follower_write() { engine.handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: Rows { schema, rows: build_rows_for_key("follower", 2, 4, 0), @@ -935,6 +971,7 @@ async fn test_region_write_buffer_rejects_only_full_region_queue() { .handle_request( hot_region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: Rows { schema: stalled_schema, rows: build_rows_for_key("hot", 2, 1026, 2), diff --git a/src/mito2/src/engine/open_test.rs b/src/mito2/src/engine/open_test.rs index 55cd92e6eda..c29c9aa3915 100644 --- a/src/mito2/src/engine/open_test.rs +++ b/src/mito2/src/engine/open_test.rs @@ -478,6 +478,7 @@ async fn test_engine_open_readonly_with_format(flat_format: bool) { .handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: rows.clone(), hint: None, partition_expr_version: None, diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs index 69c08c17bb1..d63dcc77bd1 100644 --- a/src/mito2/src/engine/scan_test.rs +++ b/src/mito2/src/engine/scan_test.rs @@ -394,6 +394,167 @@ async fn test_json2_v1_region_reopen_and_compaction() -> WhateverResult<()> { Ok(()) } +#[tokio::test] +async fn test_json2_mixed_subject_compaction_preserves_values() -> WhateverResult<()> { + let request = CreateRequestBuilder::new() + .field_datatype(ConcreteDataType::json2(JsonNativeType::Object( + JsonObjectType::new(), + ))) + .insert_option("append_mode", "true") + .insert_option("memtable.type", "bulk") + .insert_option("sst_format", "flat") + .build(); + let table_dir = request.table_dir.clone(); + let schema = test_util::rows_schema(&request); + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + min_compaction_interval: std::time::Duration::from_secs(3600), + ..Default::default() + }) + .await; + let region_id = RegionId::new(1027, 0); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await?; + // #9133: the second SST stores a heterogeneous subject in its remainder. + let values = [ + json!({"subject": {"cid": "first", "uri": "at://first"}}), + json!({"subject": "did:plc:x"}), + json!({"subject": {"cid": "last", "uri": "at://last"}}), + json!({"subject": {"cid": 42}}), + ]; + for (offset, batch) in [(0, &values[..1]), (1, &values[1..])] { + test_util::put_rows( + &engine, + region_id, + Rows { + schema: schema.clone(), + rows: batch + .iter() + .enumerate() + .map(|(i, value)| { + row(vec![ + ValueData::StringValue("tag".into()), + ValueData::JsonValue(encode_json_value(JsonValue::from(value.clone()))), + ValueData::TimestampMillisecondValue((offset + i) as i64 * 1000), + ]) + }) + .collect(), + }, + ) + .await; + test_util::flush_region(&engine, region_id, None).await; + } + let old_ids = engine + .scanner(region_id, ScanRequest::default()) + .await? + .file_ids(); + assert_eq!(2, old_ids.len()); + for _ in 0..2 { + engine + .handle_request( + region_id, + RegionRequest::Compact(RegionCompactRequest { + options: compact_request::Options::StrictWindow(StrictWindow { + window_seconds: 86400, + }), + ..Default::default() + }), + ) + .await?; + let scanner = engine + .scanner( + region_id, + ScanRequest { + projection: Some(vec![1]), + ..Default::default() + }, + ) + .await?; + assert_eq!( + 1, + scanner.num_files(), + "a swallowed merge failure must not pass" + ); + assert!(scanner.file_ids().iter().all(|id| !old_ids.contains(id))); + let batches = RecordBatches::try_collect(scanner.scan().await?).await?; + let mut actual = Vec::new(); + for batch in batches.iter() { + let array = batch.column_by_name("field_0").unwrap().clone(); + for i in 0..array.len() { + actual.push(JsonArray::from(&array).try_get_value(i)?); + } + } + assert_eq!(values.as_slice(), actual.as_slice()); + reopen_region(&engine, region_id, table_dir.clone(), true, HashMap::new()).await; + let scanner = engine + .scanner( + region_id, + ScanRequest { + projection: Some(vec![1]), + json_type_hint: HashMap::from([( + "field_0".into(), + JsonNativeType::Object(JsonObjectType::from([( + "subject".into(), + JsonNativeType::String, + )])), + )]), + ..Default::default() + }, + ) + .await?; + let batches = RecordBatches::try_collect(scanner.scan().await?).await?; + let mut subjects = Vec::new(); + for batch in batches.iter() { + let array = batch.column_by_name("field_0").unwrap().clone(); + for i in 0..array.len() { + subjects.push(JsonArray::from(&array).try_get_value(i)?); + } + } + let expected = values.iter().map(|value| { + let subject = &value["subject"]; + json!({"subject": subject.as_str().map(str::to_string).unwrap_or_else(|| subject.to_string())}) + }).collect::>(); + assert_eq!( + expected, subjects, + "projection must read values spilled to remainder" + ); + let scanner = engine + .scanner( + region_id, + ScanRequest { + projection: Some(vec![1]), + json_type_hint: HashMap::from([( + "field_0".into(), + JsonNativeType::Object(JsonObjectType::from([( + "subject".into(), + JsonNativeType::Object(JsonObjectType::from([( + "cid".into(), + JsonNativeType::String, + )])), + )])), + )]), + ..Default::default() + }, + ) + .await?; + let batches = RecordBatches::try_collect(scanner.scan().await?).await?; + let mut cids = Vec::new(); + for batch in batches.iter() { + let array = batch.column_by_name("field_0").unwrap().clone(); + for i in 0..array.len() { + cids.push(JsonArray::from(&array).try_get_value(i)?["subject"]["cid"].clone()); + } + } + assert_eq!( + vec![json!("first"), json!(null), json!("last"), json!("42")], + cids + ); + } + Ok(()) +} + #[tokio::test] async fn test_flush_aligns_different_json2_layouts() -> WhateverResult<()> { let mut request = CreateRequestBuilder::new() @@ -1129,6 +1290,7 @@ async fn test_two_phase_series_scan() { }; let put = |rows| { RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows, hint: Some(WriteHint { primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse.into(), @@ -2389,6 +2551,7 @@ fn build_bulk_insert_request( let (schema, record_batch) = encode_to_flight_data(payload.clone()); RegionBulkInsertsRequest { + skip_wal: false, region_id, payload, raw_data: ArrowIpc { @@ -2401,6 +2564,139 @@ fn build_bulk_insert_request( } } +#[tokio::test] +async fn test_bulk_skip_wal_recovery() { + check_bulk_skip_wal_recovery(false, false).await; + check_bulk_skip_wal_recovery(false, true).await; + check_bulk_skip_wal_recovery(true, false).await; + check_bulk_skip_wal_recovery(true, true).await; +} + +async fn check_bulk_skip_wal_recovery(skip_wal: bool, flush: bool) { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + let request = CreateRequestBuilder::new() + .insert_option("memtable.type", "bulk") + .build(); + let table_dir = request.table_dir.clone(); + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let mut request = build_bulk_insert_request(region_id, 0, 4); + request.skip_wal = skip_wal; + let affected = engine + .handle_request(region_id, RegionRequest::BulkInserts(request)) + .await + .unwrap(); + assert_eq!(affected.affected_rows, 4); + let region = engine.get_region(region_id).unwrap(); + let current = region.version_control.current(); + assert_eq!(current.committed_sequence, 4); + assert_eq!(current.last_entry_id, u64::from(!skip_wal)); + assert_eq!(current.version.flushed_entry_id, 0); + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + assert_eq!( + RecordBatches::try_collect(stream) + .await + .unwrap() + .iter() + .map(|b| b.num_rows()) + .sum::(), + 4 + ); + if flush { + test_util::flush_region(&engine, region_id, None).await; + let current = region.version_control.current(); + assert_eq!(current.last_entry_id, u64::from(!skip_wal)); + assert_eq!(current.version.flushed_sequence, 4); + assert_eq!( + engine + .region_statistic(region_id) + .unwrap() + .manifest + .data_flushed_entry_id(), + u64::from(!skip_wal) + ); + } + reopen_region(&engine, region_id, table_dir, false, HashMap::new()).await; + let stream = engine + .scan_to_stream(region_id, ScanRequest::default()) + .await + .unwrap(); + assert_eq!( + RecordBatches::try_collect(stream) + .await + .unwrap() + .iter() + .map(|b| b.num_rows()) + .sum::(), + if skip_wal && !flush { 0 } else { 4 } + ); +} + +#[tokio::test] +async fn test_bulk_skip_wal_flush_watermarks() { + let mut env = TestEnv::new().await; + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + engine + .handle_request( + region_id, + RegionRequest::Create( + CreateRequestBuilder::new() + .insert_option("memtable.type", "bulk") + .build(), + ), + ) + .await + .unwrap(); + // Establish a persisted baseline, skip one bulk write, then resume WAL. + for (round, skip_wal) in [false, true, false].into_iter().enumerate() { + let region = engine.get_region(region_id).unwrap(); + let before = region.version_control.current(); + let mut request = build_bulk_insert_request(region_id, round * 4, (round + 1) * 4); + request.skip_wal = skip_wal; + let affected = engine + .handle_request(region_id, RegionRequest::BulkInserts(request)) + .await + .unwrap(); + assert_eq!(affected.affected_rows, 4); + let expected_entry_id = before.last_entry_id + u64::from(!skip_wal); + let after = region.version_control.current(); + assert_eq!(after.last_entry_id, expected_entry_id); + assert_eq!(after.committed_sequence, before.committed_sequence + 4); + assert_eq!( + after.version.flushed_sequence, + before.version.flushed_sequence + ); + assert_eq!( + engine + .region_statistic(region_id) + .unwrap() + .manifest + .data_flushed_entry_id(), + before.version.flushed_entry_id + ); + test_util::flush_region(&engine, region_id, None).await; + let after = region.version_control.current(); + assert_eq!(after.last_entry_id, expected_entry_id); + assert_eq!(after.version.flushed_sequence, (round as u64 + 1) * 4); + assert_eq!( + engine + .region_statistic(region_id) + .unwrap() + .manifest + .data_flushed_entry_id(), + expected_entry_id + ); + } +} + /// Bulk parts fold per-part sequences (commit-unit granularity): an exact range /// keeps or drops each whole part. The same set must hold before and after flush. #[tokio::test] @@ -3234,6 +3530,7 @@ async fn build_sparse_exact_metric_engine( }; let put = |rows| { RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows, hint: Some(WriteHint { primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse.into(), @@ -3509,6 +3806,7 @@ async fn test_two_phase_sparse_metric_exact_sequence_cache_isolation() { panic!("sparse metric scan must use SeriesScan"); }; assert_eq!("two_phase", scanner.mode()); + assert_eq!(1, scanner.properties().num_partitions()); let batches = RecordBatches::try_collect(scanner.build_stream().await.unwrap()) .await .unwrap(); @@ -3520,25 +3818,27 @@ async fn test_two_phase_sparse_metric_exact_sequence_cache_isolation() { let (cold, _) = scan(restrictive.clone()).await; // Cache insertion is performed by the range stream's async concat task. - // Wait for a completed cache read rather than assuming a fixed number of - // scheduler yields. The aggregate metric includes candidate-stage hits, so - // output-isolation assertions below remain the data-stage acceptance check. - let (warm, warm_verbose) = tokio::time::timeout(Duration::from_secs(5), async { + // Wait for the SeriesReader's data-stage entry, rather than treating a + // candidate-stage cache hit as evidence that the reader replayed its data. + let warm = tokio::time::timeout(Duration::from_secs(5), async { loop { let (rows, verbose) = scan(restrictive.clone()).await; - if verbose.contains("\"range_cache_hit\"") { - return (rows, verbose); + // The single output partition is 0; the candidate distributor uses + // synthetic partition 1. Only a partition-0 hit proves data replay. + let data_metrics = verbose + .split("\"partition\":0, \"metrics\":") + .nth(1) + .and_then(|metrics| metrics.split("\"partition\":").next()) + .expect("missing two-phase data-reader metrics"); + if data_metrics.contains("\"range_cache_hit\":") { + return rows; } tokio::task::yield_now().await; } }) .await - .expect("range cache did not complete within timeout"); + .expect("series-data range cache did not complete and replay within timeout"); assert_eq!(cold, warm); - assert!( - warm_verbose.contains("\"range_cache_hit\":"), - "warm two-phase scan did not report a range-cache hit: {warm_verbose}" - ); assert_eq!( vec![ diff --git a/src/mito2/src/engine/set_role_state_test.rs b/src/mito2/src/engine/set_role_state_test.rs index 1d0f3c188b8..53aa3511459 100644 --- a/src/mito2/src/engine/set_role_state_test.rs +++ b/src/mito2/src/engine/set_role_state_test.rs @@ -114,6 +114,7 @@ async fn test_set_role_state_gracefully_with_format(flat_format: bool) { .handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: rows.clone(), hint: None, partition_expr_version: None, @@ -208,6 +209,7 @@ async fn test_write_downgrading_region_with_format(flat_format: bool) { .handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: rows.clone(), hint: None, partition_expr_version: None, diff --git a/src/mito2/src/engine/skip_wal_test.rs b/src/mito2/src/engine/skip_wal_test.rs index a564411473c..88304777b2d 100644 --- a/src/mito2/src/engine/skip_wal_test.rs +++ b/src/mito2/src/engine/skip_wal_test.rs @@ -522,6 +522,7 @@ async fn test_close_region_skip_wal_rejects_writes_queued_after_close() { .handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: Rows { schema: rows_schema(&request), rows: build_rows(3, 4), @@ -553,6 +554,7 @@ async fn test_close_region_skip_wal_rejects_writes_queued_after_close() { .handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: Rows { schema: rows_schema(&request_cloned), rows: build_rows(4, 5), @@ -579,6 +581,7 @@ async fn test_close_region_skip_wal_rejects_writes_queued_after_close() { .handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: Rows { schema: rows_schema(&request_cloned), rows: build_rows(5, 6), diff --git a/src/mito2/src/engine/staging_test.rs b/src/mito2/src/engine/staging_test.rs index 12263976aa2..bd823708509 100644 --- a/src/mito2/src/engine/staging_test.rs +++ b/src/mito2/src/engine/staging_test.rs @@ -285,6 +285,7 @@ async fn test_staging_reject_all_writes_rejects_put() { .handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows, hint: None, partition_expr_version: None, @@ -355,6 +356,7 @@ async fn test_staging_write_partition_expr_version_with_format(flat_format: bool .handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: bad_rows, hint: None, partition_expr_version: Some(origin_version), @@ -376,6 +378,7 @@ async fn test_staging_write_partition_expr_version_with_format(flat_format: bool .handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: compat_rows, hint: None, partition_expr_version: None, @@ -393,6 +396,7 @@ async fn test_staging_write_partition_expr_version_with_format(flat_format: bool .handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: ok_rows, hint: None, partition_expr_version: Some(expected_version), @@ -440,6 +444,7 @@ async fn test_staging_write_partition_expr_version_with_format(flat_format: bool .handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: exit_rows, hint: None, partition_expr_version: Some(origin_version), @@ -457,6 +462,7 @@ async fn test_staging_write_partition_expr_version_with_format(flat_format: bool .handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: compat_rows, hint: None, partition_expr_version: None, @@ -483,6 +489,7 @@ async fn test_staging_write_partition_expr_version_with_format(flat_format: bool .handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows: commit_rows, hint: None, partition_expr_version: Some(expected_version), diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index b7c8513afc7..7b431c1b890 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -451,6 +451,7 @@ impl RegionFlushTask { let mut write_opts = WriteOptions { write_buffer_size: self.engine_config.sst_write_buffer_size, + float_field_encoding: version.options.float_field_encoding, ..Default::default() }; if let Some(row_group_size) = self.row_group_size { @@ -1728,6 +1729,7 @@ mod tests { ( SenderBulkRequest { + skip_wal: false, sender: OptionOutputTx::from(sender), region_id, request: converter.convert().unwrap(), diff --git a/src/mito2/src/memtable.rs b/src/mito2/src/memtable.rs index f6481fbcdce..bc085c3bd4e 100644 --- a/src/mito2/src/memtable.rs +++ b/src/mito2/src/memtable.rs @@ -491,6 +491,7 @@ impl MemtableBuilderProvider { BulkMemtableBuilder::new(self.write_buffer_manager.clone(), !dedup, merge_mode) .with_config(config.clone()) .with_row_group_size(options.row_group_size()) + .with_float_field_encoding(options.float_field_encoding) .with_compact_dispatcher(self.compact_dispatcher.clone()), ), Some(MemtableOptions::TimeSeries) => Arc::new(TimeSeriesMemtableBuilder::new( @@ -515,6 +516,7 @@ impl MemtableBuilderProvider { ) .with_config(self.default_bulk_memtable_config.clone()) .with_row_group_size(options.row_group_size()) + .with_float_field_encoding(options.float_field_encoding) .with_compact_dispatcher(self.compact_dispatcher.clone()); if let Some(MemtableOptions::Bulk(config)) = &options.memtable { diff --git a/src/mito2/src/memtable/bulk.rs b/src/mito2/src/memtable/bulk.rs index 8e1346b9921..b30349cc386 100644 --- a/src/mito2/src/memtable/bulk.rs +++ b/src/mito2/src/memtable/bulk.rs @@ -40,6 +40,7 @@ use rayon::prelude::*; use serde::{Deserialize, Serialize}; use serde_with::{DisplayFromStr, serde_as}; use store_api::metadata::RegionMetadataRef; +use store_api::mito_engine_options::FloatFieldEncoding; use store_api::storage::{ColumnId, FileId, RegionId, SequenceRange}; use tokio::sync::Semaphore; @@ -435,6 +436,7 @@ pub struct BulkMemtable { merge_mode: MergeMode, /// Max number of rows in a parquet row group for encoded parts. row_group_size: usize, + float_field_encoding: FloatFieldEncoding, } impl std::fmt::Debug for BulkMemtable { @@ -665,27 +667,7 @@ impl Memtable for BulkMemtable { } fn fork(&self, id: MemtableId, metadata: &RegionMetadataRef) -> MemtableRef { - Arc::new(Self { - id, - config: self.config.clone(), - parts: Arc::new(RwLock::new(BulkParts::default())), - metadata: metadata.clone(), - alloc_tracker: AllocTracker::new(self.alloc_tracker.write_buffer_manager()), - max_timestamp: AtomicI64::new(i64::MIN), - min_timestamp: AtomicI64::new(i64::MAX), - max_sequence: AtomicU64::new(0), - num_rows: AtomicUsize::new(0), - compactor: Arc::new(Mutex::new(MemtableCompactor::new( - metadata.region_id, - id, - self.config.clone(), - self.row_group_size, - ))), - compact_dispatcher: self.compact_dispatcher.clone(), - append_mode: self.append_mode, - merge_mode: self.merge_mode, - row_group_size: self.row_group_size, - }) + self.fork_inner(id, metadata) } fn compact(&self, for_flush: bool) -> Result<()> { @@ -714,6 +696,32 @@ impl Memtable for BulkMemtable { } impl BulkMemtable { + fn fork_inner(&self, id: MemtableId, metadata: &RegionMetadataRef) -> Arc { + Arc::new(Self { + id, + config: self.config.clone(), + parts: Arc::new(RwLock::new(BulkParts::default())), + metadata: metadata.clone(), + alloc_tracker: AllocTracker::new(self.alloc_tracker.write_buffer_manager()), + max_timestamp: AtomicI64::new(i64::MIN), + min_timestamp: AtomicI64::new(i64::MAX), + max_sequence: AtomicU64::new(0), + num_rows: AtomicUsize::new(0), + compactor: Arc::new(Mutex::new(MemtableCompactor::new( + metadata.region_id, + id, + self.config.clone(), + self.row_group_size, + self.float_field_encoding, + ))), + compact_dispatcher: self.compact_dispatcher.clone(), + append_mode: self.append_mode, + merge_mode: self.merge_mode, + row_group_size: self.row_group_size, + float_field_encoding: self.float_field_encoding, + }) + } + /// Creates a new BulkMemtable with the default row group size. pub fn new( id: MemtableId, @@ -747,6 +755,32 @@ impl BulkMemtable { append_mode: bool, merge_mode: MergeMode, row_group_size: usize, + ) -> Self { + Self::new_with_row_group_size_and_encoding( + id, + config, + metadata, + write_buffer_manager, + compact_dispatcher, + append_mode, + merge_mode, + row_group_size, + FloatFieldEncoding::default(), + ) + } + + /// Creates a new BulkMemtable with the given row group size and float encoding. + #[allow(clippy::too_many_arguments)] + fn new_with_row_group_size_and_encoding( + id: MemtableId, + config: BulkMemtableConfig, + metadata: RegionMetadataRef, + write_buffer_manager: Option, + compact_dispatcher: Option>, + append_mode: bool, + merge_mode: MergeMode, + row_group_size: usize, + float_field_encoding: FloatFieldEncoding, ) -> Self { let config = config.sanitize(); let region_id = metadata.region_id; @@ -765,11 +799,13 @@ impl BulkMemtable { id, config, row_group_size, + float_field_encoding, ))), compact_dispatcher, append_mode, merge_mode, row_group_size, + float_field_encoding, } } @@ -1179,6 +1215,7 @@ struct MemtableCompactor { config: BulkMemtableConfig, /// Max number of rows in a parquet row group for encoded parts. row_group_size: usize, + float_field_encoding: FloatFieldEncoding, } impl MemtableCompactor { @@ -1188,12 +1225,14 @@ impl MemtableCompactor { memtable_id: MemtableId, config: BulkMemtableConfig, row_group_size: usize, + float_field_encoding: FloatFieldEncoding, ) -> Self { Self { region_id, memtable_id, config, row_group_size, + float_field_encoding, } } @@ -1233,6 +1272,7 @@ impl MemtableCompactor { let encode_row_threshold = self.config.encode_row_threshold; let encode_bytes_threshold = self.config.encode_bytes_threshold; let row_group_size = self.row_group_size; + let float_field_encoding = self.float_field_encoding; // Merge all groups in parallel let merged_parts = collected @@ -1247,6 +1287,7 @@ impl MemtableCompactor { encode_row_threshold, encode_bytes_threshold, row_group_size, + float_field_encoding, ) }) .collect::>>>()?; @@ -1282,6 +1323,7 @@ impl MemtableCompactor { encode_row_threshold: usize, encode_bytes_threshold: usize, row_group_size: usize, + float_field_encoding: FloatFieldEncoding, ) -> Result> { if parts_to_merge.is_empty() { return Ok(None); @@ -1382,7 +1424,11 @@ impl MemtableCompactor { if estimated_total_rows > encode_row_threshold || estimated_total_bytes > encode_bytes_threshold { - let encoder = BulkPartEncoder::new(metadata.clone(), row_group_size)?; + let encoder = BulkPartEncoder::new_with_float_field_encoding( + metadata.clone(), + row_group_size, + float_field_encoding, + )?; let mut metrics = BulkPartEncodeMetrics::default(); let encoded_part = encoder.encode_record_batch_iter( boxed_iter, @@ -1508,6 +1554,7 @@ pub struct BulkMemtableBuilder { merge_mode: MergeMode, /// Max number of rows in a parquet row group for encoded parts. row_group_size: usize, + float_field_encoding: FloatFieldEncoding, } impl Default for BulkMemtableBuilder { @@ -1519,6 +1566,7 @@ impl Default for BulkMemtableBuilder { append_mode: false, merge_mode: MergeMode::default(), row_group_size: DEFAULT_ROW_GROUP_SIZE, + float_field_encoding: FloatFieldEncoding::default(), } } } @@ -1550,6 +1598,11 @@ impl BulkMemtableBuilder { self } + pub(super) fn with_float_field_encoding(mut self, encoding: FloatFieldEncoding) -> Self { + self.float_field_encoding = encoding; + self + } + /// Sets the compact dispatcher. pub fn with_compact_dispatcher(mut self, compact_dispatcher: Arc) -> Self { self.compact_dispatcher = Some(compact_dispatcher); @@ -1564,7 +1617,7 @@ impl BulkMemtableBuilder { impl MemtableBuilder for BulkMemtableBuilder { fn build(&self, id: MemtableId, metadata: &RegionMetadataRef) -> MemtableRef { - Arc::new(BulkMemtable::new_with_row_group_size( + Arc::new(BulkMemtable::new_with_row_group_size_and_encoding( id, self.config.clone(), metadata.clone(), @@ -1573,6 +1626,7 @@ impl MemtableBuilder for BulkMemtableBuilder { self.append_mode, self.merge_mode, self.row_group_size, + self.float_field_encoding, )) } @@ -1597,6 +1651,7 @@ mod tests { use datatypes::schema::ColumnSchema; use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; use mito_codec::row_converter::build_primary_key_codec; + use parquet::basic::{Encoding, Type as PhysicalType}; use serde_json::json; use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder, RegionMetadataRef}; @@ -1807,7 +1862,7 @@ mod tests { let second = mock_bulk_part_with_json2_values( &metadata, vec![3000, 4000], - vec![json!({"b": 3}), json!({"b": 4})], + vec![json!({"a": 3, "b": 3}), json!({"a": 4, "b": 4})], 200, )?; let parts = vec![ @@ -1829,6 +1884,7 @@ mod tests { usize::MAX, usize::MAX, DEFAULT_ROW_GROUP_SIZE, + FloatFieldEncoding::default(), )? .unwrap(); let MergedPart::Multi(part) = merged else { @@ -1879,9 +1935,18 @@ mod tests { unreachable!() }; assert_eq!( - vec![JSON2_REMAINDER_FIELD_NAME, "a", "b"], + vec![JSON2_REMAINDER_FIELD_NAME], fields.iter().map(|x| x.name().as_str()).collect::>() ); + // Neither path is explicit in every source; keep both opaque without losing values. + assert_eq!(4, batch.num_rows()); + let json = datatypes::vectors::json::array::JsonArray::from(batch.column(0)); + let restored = json.project_to_v2(schema.field(0), &ArrowDataType::Binary)?; + let restored = datatypes::vectors::json::array::JsonArray::from(&restored); + assert_eq!(json!({"a": 1}), restored.try_get_value(0)?); + assert_eq!(json!({"a": 2}), restored.try_get_value(1)?); + assert_eq!(json!({"b": 3}), restored.try_get_value(2)?); + assert_eq!(json!({"b": 4}), restored.try_get_value(3)?); Ok(()) } @@ -2134,7 +2199,7 @@ mod tests { #[test] fn test_bulk_memtable_fork() { let metadata = metadata_for_test(); - let original_memtable = BulkMemtable::new( + let original_memtable = BulkMemtable::new_with_row_group_size_and_encoding( 333, BulkMemtableConfig::default(), metadata.clone(), @@ -2142,6 +2207,8 @@ mod tests { None, false, MergeMode::LastRow, + DEFAULT_ROW_GROUP_SIZE, + FloatFieldEncoding::ByteStreamSplit, ); let bulk_part = @@ -2150,9 +2217,25 @@ mod tests { original_memtable.write_bulk(bulk_part).unwrap(); - let forked_memtable = original_memtable.fork(444, &metadata); + let forked_memtable = original_memtable.fork_inner(444, &metadata); assert_eq!(forked_memtable.id(), 444); + assert_eq!( + FloatFieldEncoding::ByteStreamSplit, + original_memtable.float_field_encoding + ); + assert_eq!( + original_memtable.float_field_encoding, + forked_memtable.float_field_encoding + ); + assert_eq!( + original_memtable.float_field_encoding, + forked_memtable + .compactor + .lock() + .unwrap() + .float_field_encoding + ); assert!(forked_memtable.is_empty()); assert_eq!(0, forked_memtable.stats().num_rows); @@ -2771,6 +2854,42 @@ mod tests { assert!(!BulkParts::is_merge_candidate_by_size(true, 9, 8)); } + #[test] + fn test_bulk_part_encoder_preserves_byte_stream_split() -> WhateverResult<()> { + let metadata = metadata_for_test(); + let bulk_part = create_bulk_part_with_converter( + "byte_stream_split", + 0, + vec![1000, 2000], + vec![Some(10.0), Some(20.0)], + 100, + )?; + let encoder = BulkPartEncoder::new_with_float_field_encoding( + metadata, + DEFAULT_ROW_GROUP_SIZE, + FloatFieldEncoding::ByteStreamSplit, + )?; + let encoded_part = encoder.encode_part(&bulk_part)?.unwrap(); + let field_column = encoded_part + .metadata() + .parquet_metadata + .row_groups() + .iter() + .flat_map(|row_group| row_group.columns()) + .find(|column| { + column.column_path().string() == "v1" + && column.column_type() == PhysicalType::DOUBLE + }) + .expect("Float64 field should be present in encoded bulk part"); + + assert!( + field_column + .encodings() + .any(|encoding| encoding == Encoding::BYTE_STREAM_SPLIT) + ); + Ok(()) + } + #[test] fn test_encoded_part_batch_size_uses_largest_uncompressed_row_group() { const NUM_ROWS: usize = 14; diff --git a/src/mito2/src/memtable/bulk/part.rs b/src/mito2/src/memtable/bulk/part.rs index 2a0c42b7035..f0af666a644 100644 --- a/src/mito2/src/memtable/bulk/part.rs +++ b/src/mito2/src/memtable/bulk/part.rs @@ -53,6 +53,7 @@ use smallvec::SmallVec; use snafu::{OptionExt, ResultExt}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{RegionMetadata, RegionMetadataRef}; +use store_api::mito_engine_options::FloatFieldEncoding; use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; use store_api::storage::{ColumnId, FileId, SequenceNumber, SequenceRange}; @@ -70,7 +71,7 @@ use crate::sst::SeriesEstimator; use crate::sst::index::IndexOutput; use crate::sst::parquet::flat_format::primary_key_column_index; use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyArrayBuilder}; -use crate::sst::parquet::{PARQUET_METADATA_KEY, SstInfo}; +use crate::sst::parquet::{PARQUET_METADATA_KEY, SstInfo, apply_float_field_encoding}; const INIT_DICT_VALUE_CAPACITY: usize = 8; @@ -1306,22 +1307,29 @@ pub struct BulkPartEncoder { impl BulkPartEncoder { pub fn new(metadata: RegionMetadataRef, row_group_size: usize) -> Result { + Self::new_with_float_field_encoding(metadata, row_group_size, FloatFieldEncoding::default()) + } + + pub(super) fn new_with_float_field_encoding( + metadata: RegionMetadataRef, + row_group_size: usize, + float_field_encoding: FloatFieldEncoding, + ) -> Result { // TODO(yingwen): Skip arrow schema if needed. let json = metadata.to_json().context(InvalidMetadataSnafu)?; let key_value_meta = parquet::file::metadata::KeyValue::new(PARQUET_METADATA_KEY.to_string(), json); // TODO(yingwen): Do we need compression? - let writer_props = Some( - WriterProperties::builder() - .set_key_value_metadata(Some(vec![key_value_meta])) - .set_write_batch_size(row_group_size) - .set_max_row_group_row_count(Some(row_group_size)) - .set_compression(Compression::ZSTD(ZstdLevel::default())) - .set_column_index_truncate_length(None) - .set_statistics_truncate_length(None) - .build(), - ); + let mut props = WriterProperties::builder() + .set_key_value_metadata(Some(vec![key_value_meta])) + .set_write_batch_size(row_group_size) + .set_max_row_group_row_count(Some(row_group_size)) + .set_compression(Compression::ZSTD(ZstdLevel::default())) + .set_column_index_truncate_length(None) + .set_statistics_truncate_length(None); + props = apply_float_field_encoding(props, &metadata, float_field_encoding); + let writer_props = Some(props.build()); Ok(Self { metadata, diff --git a/src/mito2/src/read/flat_projection.rs b/src/mito2/src/read/flat_projection.rs index 19dc962c56e..2abf6aa48af 100644 --- a/src/mito2/src/read/flat_projection.rs +++ b/src/mito2/src/read/flat_projection.rs @@ -667,6 +667,77 @@ mod tests { ); } + #[test] + fn test_single_value_dictionary_preserves_string_type() { + use datatypes::arrow::array::{DictionaryArray, TimestampMillisecondArray, UInt32Array}; + use datatypes::arrow::datatypes::UInt32Type; + + for data_type in [ + ConcreteDataType::large_string_datatype(), + ConcreteDataType::utf8_view_datatype(), + ] { + let mut builder = RegionMetadataBuilder::new(RegionId::new(1024, 0)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new("tag", data_type.clone(), true), + semantic_type: SemanticType::Tag, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 1, + }) + .primary_key(vec![0]); + let metadata = Arc::new(builder.build().unwrap()); + let mapper = FlatProjectionMapper::new(&metadata, [0, 1]).unwrap(); + + for value in [Value::from("greptime"), Value::Null] { + let mut values = data_type.create_mutable_vector(1); + values.try_push_value_ref(&value.as_value_ref()).unwrap(); + let dictionary = Arc::new( + DictionaryArray::::try_new( + UInt32Array::from(vec![0, 0, 0]), + values.to_vector().to_arrow_array(), + ) + .unwrap(), + ); + let arrays: Vec> = mapper + .batch_schema() + .iter() + .map(|(id, _)| match id { + 0 => dictionary.clone() as Arc, + 1 => Arc::new(TimestampMillisecondArray::from(vec![1, 2, 3])), + _ => unreachable!(), + }) + .collect(); + let fields = mapper + .batch_schema() + .iter() + .zip(&arrays) + .map(|((id, _), array)| { + Field::new(id.to_string(), array.data_type().clone(), true) + }) + .collect::>(); + let batch = DfRecordBatch::try_new( + Arc::new(datatypes::arrow::datatypes::Schema::new(fields)), + arrays, + ) + .unwrap(); + let output = mapper.convert(&batch, &CacheStrategy::Disabled).unwrap(); + let vector = Helper::try_into_vector(output.column(0)).unwrap(); + assert_eq!(data_type, vector.data_type()); + for row in 0..3 { + assert_eq!(value, vector.get(row)); + } + } + } + } + #[test] fn test_json_type_hint_does_not_concretize_legacy_json() { let metadata = metadata_with_legacy_json(); diff --git a/src/mito2/src/read/projection.rs b/src/mito2/src/read/projection.rs index c6d4baee3ac..fbf6a29b73d 100644 --- a/src/mito2/src/read/projection.rs +++ b/src/mito2/src/read/projection.rs @@ -19,7 +19,7 @@ use std::cmp::Ordering; use common_recordbatch::error::DataTypesSnafu; use datatypes::prelude::{ConcreteDataType, DataType}; use datatypes::value::Value; -use datatypes::vectors::VectorRef; +use datatypes::vectors::{Helper, VectorRef}; use snafu::{OptionExt, ResultExt}; use store_api::metadata::RegionMetadataRef; use store_api::storage::ColumnId; @@ -82,12 +82,21 @@ pub(crate) fn new_repeated_vector( value: &Value, num_rows: usize, ) -> common_recordbatch::error::Result { - let mut mutable_vector = data_type.create_mutable_vector(1); - mutable_vector - .try_push_value_ref(&value.as_value_ref()) - .context(DataTypesSnafu)?; - let base_vector = mutable_vector.to_vector(); - Ok(base_vector.replicate(&[num_rows])) + if let Ok(vector) = value + .try_to_scalar_value(data_type) + .and_then(|scalar| Helper::try_from_scalar_value(scalar, num_rows, Some(data_type))) + { + return Ok(vector); + } + + // Preserve extension types that cannot safely round-trip through ScalarValue. + let mut mutable_vector = data_type.create_mutable_vector(num_rows); + for _ in 0..num_rows { + mutable_vector + .try_push_value_ref(&value.as_value_ref()) + .context(DataTypesSnafu)?; + } + Ok(mutable_vector.to_vector()) } #[cfg(test)] @@ -110,6 +119,51 @@ mod tests { use crate::read::flat_projection::FlatProjectionMapper; use crate::read::read_columns::ReadColumns; + #[test] + fn test_repeated_struct_null_fields_and_json() { + use datatypes::types::{StructField, StructType}; + use datatypes::value::StructValue; + + let inner_type = StructType::from([StructField::new( + "x", + ConcreteDataType::int32_datatype(), + true, + )]); + let inner = Value::Struct(StructValue::new(vec![Value::Null], inner_type)); + let json = datatypes::json::JsonSettings::default() + .encode(serde_json::json!({"answer": 42})) + .unwrap(); + let nested_type = StructType::from([StructField::new("nested", inner.data_type(), true)]); + let json_type = StructType::from([StructField::new("json", json.data_type(), true)]); + let values = [ + inner.clone(), + Value::Struct(StructValue::new(vec![inner], nested_type)), + Value::Struct(StructValue::new(vec![], StructType::default())), + Value::Struct(StructValue::new(vec![json], json_type)), + ]; + for value in values { + let data_type = value.data_type(); + // JSON children are read back as their underlying struct values. + let expected = serde_json::Value::try_from(value.clone()).unwrap(); + for num_rows in [0, 1, 3] { + let vector = new_repeated_vector(&data_type, &value, num_rows).unwrap(); + assert_eq!(data_type, vector.data_type()); + assert_eq!( + data_type.as_arrow_type(), + *vector.to_arrow_array().data_type() + ); + assert_eq!(num_rows, vector.len()); + assert_eq!(0, vector.null_count()); + for row in 0..num_rows { + assert_eq!( + expected, + serde_json::Value::try_from(vector.get(row)).unwrap() + ); + } + } + } + } + fn print_record_batch(record_batch: RecordBatch) -> String { pretty::pretty_format_batches(&[record_batch.into_df_record_batch()]) .unwrap() diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index ebabbbd295e..0d758f676ce 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -953,6 +953,7 @@ where OptionOutputTx::none(), // We should respect the sequence in WAL during replay. Some(mutation.sequence), + false, ); } @@ -974,7 +975,8 @@ where region_write_ctx.push_bulk( OptionOutputTx::none(), part, - Some(bulk_sequence_from_wal) + Some(bulk_sequence_from_wal), + false ), RegionCorruptedSnafu { region_id, diff --git a/src/mito2/src/region/options.rs b/src/mito2/src/region/options.rs index 9a8d17c954c..a2a165c12af 100644 --- a/src/mito2/src/region/options.rs +++ b/src/mito2/src/region/options.rs @@ -32,7 +32,9 @@ use store_api::codec::PrimaryKeyEncoding; use store_api::metric_engine_consts::{ MEMTABLE_PARTITION_TREE_PRIMARY_KEY_ENCODING, PRIMARY_KEY_ENCODING, }; -use store_api::mito_engine_options::{COMPACTION_OVERRIDE, MAX_ROW_GROUP_ROW_COUNT_LIMIT}; +use store_api::mito_engine_options::{ + COMPACTION_OVERRIDE, FloatFieldEncoding, MAX_ROW_GROUP_ROW_COUNT_LIMIT, +}; use store_api::storage::{ColumnId, RegionId}; use strum::EnumString; @@ -128,6 +130,9 @@ pub struct RegionOptions { /// sequence metadata. #[serde(default, skip_serializing_if = "is_false")] pub preserve_row_sequence: bool, + /// Encoding for direct floating-point field columns in Parquet SSTs. + #[serde(default)] + pub float_field_encoding: FloatFieldEncoding, } fn is_false(value: &bool) -> bool { @@ -287,6 +292,8 @@ impl RegionOptions { sst_format = Some(FormatType::Flat); } + let float_field_encoding = options.float_field_encoding.unwrap_or_default(); + let compaction_override_flag = options_map .get(COMPACTION_OVERRIDE) .map(|v| matches!(v.to_lowercase().as_str(), "true" | "1")) @@ -322,6 +329,7 @@ impl RegionOptions { primary_key_encoding, write_buffer_size: options.write_buffer_size, preserve_row_sequence: options.preserve_row_sequence, + float_field_encoding, }; opts.validate()?; @@ -469,6 +477,9 @@ struct RegionOptionsWithoutEnum { max_row_group_row_count: Option, #[serde_as(as = "DisplayFromStr")] preserve_row_sequence: bool, + #[serde(rename = "experimental_sst_float_field_encoding")] + #[serde_as(as = "NoneAsEmptyString")] + float_field_encoding: Option, } impl Default for RegionOptionsWithoutEnum { @@ -485,6 +496,7 @@ impl Default for RegionOptionsWithoutEnum { sst_format: options.sst_format, max_row_group_row_count: options.max_row_group_row_count, preserve_row_sequence: options.preserve_row_sequence, + float_field_encoding: Some(options.float_field_encoding), } } } @@ -628,7 +640,9 @@ mod tests { use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_wal::options::KafkaWalOptions; - use store_api::mito_engine_options::{SKIP_WAL_KEY, WRITE_BUFFER_SIZE_KEY}; + use store_api::mito_engine_options::{ + EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, SKIP_WAL_KEY, WRITE_BUFFER_SIZE_KEY, + }; use super::*; @@ -646,6 +660,26 @@ mod tests { assert_eq!(RegionOptions::default(), options); } + #[test] + fn test_float_field_encoding_defaults_and_parses() { + let options = RegionOptions::try_from_options(RegionId::new(0, 0), &make_map(&[])).unwrap(); + assert_eq!(FloatFieldEncoding::Default, options.float_field_encoding); + + let map = make_map(&[(EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, "default")]); + let options = RegionOptions::try_from_options(RegionId::new(0, 0), &map).unwrap(); + assert_eq!(FloatFieldEncoding::Default, options.float_field_encoding); + + let map = make_map(&[(EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, "byte_stream_split")]); + let options = RegionOptions::try_from_options(RegionId::new(0, 0), &map).unwrap(); + assert_eq!( + FloatFieldEncoding::ByteStreamSplit, + options.float_field_encoding + ); + + let map = make_map(&[(EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, "invalid")]); + assert!(RegionOptions::try_from_options(RegionId::new(0, 0), &map).is_err()); + } + #[test] fn test_with_ttl() { let map = make_map(&[("ttl", "7d")]); @@ -1103,6 +1137,7 @@ mod tests { primary_key_encoding: None, write_buffer_size: None, preserve_row_sequence: false, + float_field_encoding: FloatFieldEncoding::default(), }; assert_eq!(expect, options); } @@ -1161,6 +1196,7 @@ mod tests { primary_key_encoding: None, write_buffer_size: Some(ReadableSize::mb(128)), preserve_row_sequence: true, + float_field_encoding: FloatFieldEncoding::default(), }; let region_options_json_str = serde_json::to_string(&options).unwrap(); assert!(region_options_json_str.contains("preserve_row_sequence")); @@ -1172,6 +1208,7 @@ mod tests { let got: RegionOptions = serde_json::from_str(old_region_options_json_str).unwrap(); assert_eq!(None, got.write_buffer_size); assert!(!got.preserve_row_sequence); + assert_eq!(FloatFieldEncoding::Default, got.float_field_encoding); let CompactionOptions::Twcs(twcs) = got.compaction; assert_eq!(16, twcs.active_window_l1_merge_trigger); assert_eq!(8, twcs.inactive_window_l1_merge_trigger); @@ -1238,6 +1275,7 @@ mod tests { primary_key_encoding: None, write_buffer_size: None, preserve_row_sequence: false, + float_field_encoding: FloatFieldEncoding::default(), }; assert_eq!(options, got); } diff --git a/src/mito2/src/region_write_ctx.rs b/src/mito2/src/region_write_ctx.rs index 400f7db1fc4..4bc0164cc01 100644 --- a/src/mito2/src/region_write_ctx.rs +++ b/src/mito2/src/region_write_ctx.rs @@ -89,12 +89,14 @@ pub(crate) struct RegionWriteCtx { /// We keep [WalEntry] instead of mutations to avoid taking mutations /// out of the context to construct the wal entry when we write to the wal. wal_entry: WalEntry, + /// Mutations that skip WAL, paired with their write notifiers. + memtable_mutations: Vec<(Mutation, WriteNotify)>, /// Wal options of the region being written to. provider: Provider, /// Notifiers to send write results to waiters. /// - /// The i-th notify is for i-th mutation. - notifiers: Vec, + /// The i-th notify is for the i-th mutation in `wal_entry`. + wal_notifiers: Vec, /// Notifiers for bulk requests. bulk_notifiers: Vec, /// Pending bulk write requests @@ -133,8 +135,9 @@ impl RegionWriteCtx { next_sequence: committed_sequence + 1, next_entry_id: last_entry_id + 1, wal_entry: WalEntry::default(), + memtable_mutations: Vec::new(), provider, - notifiers: Vec::new(), + wal_notifiers: Vec::new(), bulk_notifiers: vec![], failed: false, put_num: 0, @@ -153,21 +156,28 @@ impl RegionWriteCtx { write_hint: Option, tx: OptionOutputTx, sequence: Option, + skip_wal: bool, ) { if let Some(sequence) = sequence { self.next_sequence = sequence; } let num_rows = rows.as_ref().map(|rows| rows.rows.len()).unwrap_or(0); - self.wal_entry.mutations.push(Mutation { + let mutation = Mutation { op_type, sequence: self.next_sequence, rows, write_hint, - }); + }; + // Assign sequences before routing so concurrent memtable writes retain + // their logical order regardless of the WAL policy. let notify = WriteNotify::new(tx, num_rows); - // Notifiers are 1:1 map to mutations. - self.notifiers.push(notify); + if skip_wal { + self.memtable_mutations.push((mutation, notify)); + } else { + self.wal_entry.mutations.push(mutation); + self.wal_notifiers.push(notify); + } // Increase sequence number. self.next_sequence += num_rows as u64; @@ -206,13 +216,19 @@ impl RegionWriteCtx { /// Returns whether writes in this context should skip WAL. pub(crate) fn skip_wal(&self) -> bool { - self.provider == Provider::Noop || self.version.options.skip_wal + self.provider == Provider::Noop + || self.version.options.skip_wal + || (self.wal_entry.mutations.is_empty() && self.wal_entry.bulk_entries.is_empty()) } /// Sets error and marks all write operations are failed. pub(crate) fn set_error(&mut self, err: Arc) { // Set error for all notifiers. - for notify in &mut self.notifiers { + for notify in self + .wal_notifiers + .iter_mut() + .chain(self.memtable_mutations.iter_mut().map(|(_, notify)| notify)) + { notify.err = Some(err.clone()); } for notify in &mut self.bulk_notifiers { @@ -241,7 +257,7 @@ impl RegionWriteCtx { /// Consumes mutations and writes them into mutable memtable. pub(crate) async fn write_memtable(&mut self) { - debug_assert_eq!(self.notifiers.len(), self.wal_entry.mutations.len()); + debug_assert_eq!(self.wal_notifiers.len(), self.wal_entry.mutations.len()); if self.failed { return; @@ -254,34 +270,38 @@ impl RegionWriteCtx { None }; - let mutations = mem::take(&mut self.wal_entry.mutations) + let mut mutations = mem::take(&mut self.wal_entry.mutations) .into_iter() - .enumerate() - .filter_map(|(i, mutation)| { + .zip(&mut self.wal_notifiers) + .chain( + self.memtable_mutations + .iter_mut() + // Keep notifiers in the context until all writes complete. + .map(|(mutation, notify)| (mem::take(mutation), notify)), + ) + .filter_map(|(mutation, notify)| { let kvs = KeyValues::new(&self.version.metadata, mutation)?; - Some((i, kvs)) + Some((notify, kvs)) }) .collect::>(); if mutations.len() == 1 { if let Err(err) = mutable_memtable.write(&mutations[0].1) { - self.notifiers[mutations[0].0].err = Some(Arc::new(err)); + mutations[0].0.err = Some(Arc::new(err)); } } else { let mut tasks = FuturesUnordered::new(); - for (i, kvs) in mutations { + for (notify, kvs) in mutations { let mutable = mutable_memtable.clone(); // use tokio runtime to schedule tasks. - tasks.push(common_runtime::spawn_blocking_global(move || { - (i, mutable.write(&kvs)) - })); + let task = common_runtime::spawn_blocking_global(move || mutable.write(&kvs)); + tasks.push(async move { (notify, task.await) }); } - while let Some(result) = tasks.next().await { - // first unwrap the result from `spawn` above - let (i, result) = result.unwrap(); - if let Err(err) = result { - self.notifiers[i].err = Some(Arc::new(err)); + while let Some((notify, result)) = tasks.next().await { + // First unwrap the result from `spawn` above. + if let Err(err) = result.unwrap() { + notify.err = Some(Arc::new(err)); } } } @@ -298,24 +318,26 @@ impl RegionWriteCtx { sender: OptionOutputTx, mut bulk: BulkPart, sequence: Option, + skip_wal: bool, ) -> bool { if let Some(sequence) = sequence { self.next_sequence = sequence; } bulk.sequence = self.next_sequence; - let entry = match BulkWalEntry::try_from(&bulk) { - Ok(entry) => entry, - Err(e) => { - sender.send(Err(e)); - return false; - } - }; + if !skip_wal { + let entry = match BulkWalEntry::try_from(&bulk) { + Ok(entry) => entry, + Err(e) => { + sender.send(Err(e)); + return false; + } + }; + self.wal_entry.bulk_entries.push(entry); + } self.bulk_notifiers .push(WriteNotify::new(sender, bulk.num_rows())); - // Add bulk wal entry - self.wal_entry.bulk_entries.push(entry); self.next_sequence += bulk.num_rows() as u64; self.bulk_parts.push(bulk); true @@ -523,6 +545,7 @@ mod tests { use common_recordbatch::DfRecordBatch; use datatypes::arrow::array::{ArrayRef, TimestampMillisecondArray}; use datatypes::arrow::datatypes::{DataType, Field, Schema}; + use prost::Message; use store_api::logstore::provider::Provider; use tokio::sync::oneshot; @@ -531,8 +554,219 @@ mod tests { use crate::memtable::bulk::part::BulkPart; use crate::test_util::version_util::VersionControlBuilder; + #[test] + fn test_request_skip_wal_preserves_sequences_and_other_writes() { + // Ablate only the request flag: the workload and sequence allocation stay identical. + check_request_skip_wal_preserves_sequences_and_other_writes(false, false); + check_request_skip_wal_preserves_sequences_and_other_writes(false, true); + check_request_skip_wal_preserves_sequences_and_other_writes(true, false); + check_request_skip_wal_preserves_sequences_and_other_writes(true, true); + } + + fn check_request_skip_wal_preserves_sequences_and_other_writes( + skip_wal: bool, + bulk_skip_wal: bool, + ) { + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let mut ctx = RegionWriteCtx::new( + region_id, + &version_control, + Provider::raft_engine_provider(region_id.as_u64()), + None, + ); + for (op_type, skip, num_rows) in [ + (OpType::Put, skip_wal, 2), + (OpType::Put, false, 3), + (OpType::Delete, false, 1), + ] { + ctx.push_mutation( + op_type as i32, + Some(Rows { + schema: vec![], + rows: vec![api::v1::Row::default(); num_rows], + }), + None, + OptionOutputTx::none(), + None, + skip, + ); + } + // Unequal row counts make a notifier/mutation pairing mismatch visible. + for (mutation, notify) in ctx + .wal_entry + .mutations + .iter() + .zip(&ctx.wal_notifiers) + .chain( + ctx.memtable_mutations + .iter() + .map(|(mutation, notify)| (mutation, notify)), + ) + { + assert_eq!(mutation.rows.as_ref().unwrap().rows.len(), notify.num_rows); + } + assert!(ctx.push_bulk(OptionOutputTx::none(), new_bulk_part(), None, bulk_skip_wal)); + assert!(!ctx.skip_wal()); + assert_eq!(ctx.next_sequence, 9); + assert_eq!( + ctx.wal_entry.bulk_entries.len(), + usize::from(!bulk_skip_wal) + ); + assert_eq!(ctx.bulk_parts[0].sequence, 7); + let sequences: Vec<_> = ctx.wal_entry.mutations.iter().map(|m| m.sequence).collect(); + assert_eq!(sequences, if skip_wal { vec![3, 6] } else { vec![1, 3, 6] }); + // Check the actual WAL bytes after routing the mutations. + let encoded = crate::wal::encoder::WalEntryEncoder::new().encode_to_vec(&ctx.wal_entry); + let decoded = WalEntry::decode(encoded.as_slice()).unwrap(); + assert_eq!( + decoded + .mutations + .iter() + .map(|m| m.sequence) + .collect::>(), + sequences + ); + assert_eq!(decoded.bulk_entries, ctx.wal_entry.bulk_entries); + assert_eq!(ctx.wal_entry.mutations.len(), if skip_wal { 2 } else { 3 }); + assert_eq!(ctx.memtable_mutations.len(), usize::from(skip_wal)); + if skip_wal { + assert_eq!(ctx.memtable_mutations[0].0.sequence, 1); + } + assert_eq!( + ctx.wal_entry.mutations.last().unwrap().op_type, + OpType::Delete as i32 + ); + } + + #[test] + fn test_internal_delete_respects_skip_wal_flag() { + check_internal_delete_respects_skip_wal_flag(false); + check_internal_delete_respects_skip_wal_flag(true); + } + + fn check_internal_delete_respects_skip_wal_flag(skip_wal: bool) { + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let mut ctx = RegionWriteCtx::new( + region_id, + &version_control, + Provider::raft_engine_provider(region_id.as_u64()), + None, + ); + ctx.push_mutation( + OpType::Delete as i32, + Some(Rows { + schema: vec![], + rows: vec![api::v1::Row::default(); 2], + }), + None, + OptionOutputTx::none(), + None, + skip_wal, + ); + assert_eq!(ctx.skip_wal(), skip_wal); + assert_eq!(ctx.next_sequence, 3); + assert_eq!(ctx.delete_num, 2); + assert_eq!(ctx.wal_entry.mutations.len(), usize::from(!skip_wal)); + assert_eq!(ctx.memtable_mutations.len(), usize::from(skip_wal)); + let encoded = crate::wal::encoder::WalEntryEncoder::new().encode_to_vec(&ctx.wal_entry); + let decoded = WalEntry::decode(encoded.as_slice()).unwrap(); + if skip_wal { + assert!(decoded.mutations.is_empty()); + } else { + assert_eq!(decoded, ctx.wal_entry); + } + } + + #[test] + fn test_all_request_skip_wal_keeps_entry_id_and_propagates_errors() { + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let mut ctx = RegionWriteCtx::new( + region_id, + &version_control, + Provider::raft_engine_provider(region_id.as_u64()), + None, + ); + let (tx, rx) = oneshot::channel(); + ctx.push_mutation( + OpType::Put as i32, + Some(Rows { + schema: vec![], + rows: vec![api::v1::Row::default(); 2], + }), + None, + OptionOutputTx::from(tx), + None, + true, + ); + assert!(ctx.skip_wal()); + assert!(ctx.wal_entry.mutations.is_empty()); + assert_eq!(ctx.memtable_mutations.len(), 1); + assert_eq!(ctx.next_entry_id(), 1); + assert_eq!(ctx.next_sequence, 3); + ctx.set_error(Arc::new( + UnexpectedSnafu { + reason: "wal failed".to_string(), + } + .build(), + )); + drop(ctx); + assert!(rx.blocking_recv().unwrap().is_err()); + } + + #[test] + fn test_bulk_skip_wal_preserves_sequences() { + check_bulk_skip_wal_preserves_sequences(false); + check_bulk_skip_wal_preserves_sequences(true); + } + + fn check_bulk_skip_wal_preserves_sequences(skip_wal: bool) { + let builder = VersionControlBuilder::new(); + let region_id = builder.region_id(); + let version_control = Arc::new(builder.build()); + let mut ctx = RegionWriteCtx::new( + region_id, + &version_control, + Provider::raft_engine_provider(region_id.as_u64()), + None, + ); + // Alternate policies in one context, retaining all parts for the memtable. + for skip in [skip_wal, false, skip_wal] { + assert!(ctx.push_bulk(OptionOutputTx::none(), new_bulk_part(), None, skip)); + } + assert_eq!(ctx.next_sequence, 7); + assert_eq!(ctx.bulk_notifiers.len(), 3); + assert_eq!( + ctx.bulk_parts + .iter() + .map(|p| p.sequence) + .collect::>(), + vec![1, 3, 5] + ); + let encoded = crate::wal::encoder::WalEntryEncoder::new().encode_to_vec(&ctx.wal_entry); + let decoded = WalEntry::decode(encoded.as_slice()).unwrap(); + assert_eq!( + decoded + .bulk_entries + .iter() + .map(|e| e.sequence) + .collect::>(), + if skip_wal { vec![3] } else { vec![1, 3, 5] } + ); + } + #[test] fn test_set_error_marks_bulk_notifiers_failed() { + check_set_error_marks_bulk_notifiers_failed(false); + check_set_error_marks_bulk_notifiers_failed(true); + } + + fn check_set_error_marks_bulk_notifiers_failed(skip_wal: bool) { let builder = VersionControlBuilder::new(); let region_id = builder.region_id(); let version_control = Arc::new(builder.build()); @@ -540,7 +774,9 @@ mod tests { RegionWriteCtx::new(region_id, &version_control, Provider::noop_provider(), None); let (tx, rx) = oneshot::channel(); - assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(), None)); + assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(), None, skip_wal)); + assert_eq!(ctx.wal_entry.bulk_entries.len(), usize::from(!skip_wal)); + assert_eq!(ctx.bulk_parts.len(), 1); ctx.set_error(Arc::new( UnexpectedSnafu { reason: "wal failed".to_string(), diff --git a/src/mito2/src/request.rs b/src/mito2/src/request.rs index 94d442e261c..16b7eb4a6be 100644 --- a/src/mito2/src/request.rs +++ b/src/mito2/src/request.rs @@ -76,6 +76,8 @@ pub struct WriteRequest { pub name_to_index: HashMap, /// Whether each column has null. pub has_null: Vec, + /// Whether this insert should skip WAL. Never applies to deletes. + pub skip_wal: bool, /// Write hint. pub hint: Option, /// Region metadata on the time of this request is created. @@ -137,11 +139,18 @@ impl WriteRequest { name_to_index, has_null, hint: None, + skip_wal: false, region_metadata, partition_expr_version: None, }) } + /// Sets the request-level WAL policy. + pub fn with_skip_wal(mut self, skip_wal: bool) -> Self { + self.skip_wal = skip_wal; + self + } + /// Sets the write hint. pub fn with_hint(mut self, hint: Option) -> Self { self.hint = hint; @@ -551,6 +560,7 @@ pub(crate) struct SenderWriteRequest { } pub(crate) struct SenderBulkRequest { + pub(crate) skip_wal: bool, pub(crate) sender: OptionOutputTx, pub(crate) region_id: RegionId, pub(crate) request: BulkPart, @@ -672,6 +682,7 @@ impl WorkerRequest { let mut write_request = WriteRequest::new(region_id, OpType::Put, v.rows, region_metadata.clone())? .with_hint(v.hint) + .with_skip_wal(v.skip_wal) .with_partition_expr_version(v.partition_expr_version); if write_request.primary_key_encoding() == PrimaryKeyEncoding::Dense && let Some(region_metadata) = ®ion_metadata @@ -1873,6 +1884,25 @@ mod tests { ); } + #[test] + fn test_delete_request_defaults_to_writing_wal() { + let (request, _receiver) = WorkerRequest::try_from_region_request( + RegionId::new(1, 1), + RegionRequest::Delete(store_api::region_request::RegionDeleteRequest { + rows: Rows::default(), + hint: None, + partition_expr_version: None, + }), + None, + ) + .unwrap(); + let WorkerRequest::Write(request) = request else { + panic!("expected a write request"); + }; + assert_eq!(request.request.op_type, OpType::Delete); + assert!(!request.request.skip_wal); + } + #[test] fn test_write_request_metadata() { let rows = Rows { diff --git a/src/mito2/src/series_index.rs b/src/mito2/src/series_index.rs index f08178456a5..8173050f2ab 100644 --- a/src/mito2/src/series_index.rs +++ b/src/mito2/src/series_index.rs @@ -17,10 +17,18 @@ // These components are consumed by the upcoming query and maintenance integration. #[allow(dead_code)] mod catalog; +// Consumed by the follow-up background-maintenance integration. +#[allow(dead_code)] +mod bucket; +// Consumed by the follow-up background-maintenance integration. +#[allow(dead_code)] +mod builder; #[allow(dead_code)] mod purger; mod searcher; mod task; +#[cfg(test)] +mod tests; #[allow(dead_code)] mod version; mod writer; diff --git a/src/mito2/src/series_index/bucket.rs b/src/mito2/src/series_index/bucket.rs new file mode 100644 index 00000000000..7fecb253441 --- /dev/null +++ b/src/mito2/src/series_index/bucket.rs @@ -0,0 +1,877 @@ +// 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. + +//! Event-time bucket planning and source coverage. + +use std::collections::BTreeMap; +use std::time::Duration; + +use common_time::{TimeToLive, Timestamp}; +use smallvec::{SmallVec, smallvec}; +use store_api::storage::FileId; + +use crate::series_index::catalog::{SeriesIndexEntry, WindowSequence}; +use crate::sst::file::FileHandle; + +const SERIES_INDEX_TRIGGER_FILES: usize = 4; + +/// Index files sharing a non-overlapping, half-open time bucket. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct IndexBucket { + pub(crate) start: Timestamp, + pub(crate) end: Timestamp, + pub(crate) index_ids: SmallVec<[FileId; 2]>, + /// Zero means that merged indexes use incompatible window widths. + pub(crate) compaction_window_secs: i64, + /// Indexed SST summaries keyed by aligned start. + /// See [`SeriesIndexEntry::window_sequences`] for the layout and sequence assumption. + pub(crate) window_sequences: BTreeMap, +} + +impl IndexBucket { + pub(crate) fn from_entry(entry: &SeriesIndexEntry) -> Self { + Self { + start: entry.bucket_start, + end: entry.bucket_end, + index_ids: smallvec![entry.index_uuid], + compaction_window_secs: entry.compaction_window_secs, + window_sequences: entry.window_sequences.clone(), + } + } + + fn merge(&mut self, mut other: Self) { + self.start = self.start.min(other.start); + self.end = self.end.max(other.end); + if self.index_ids.is_empty() { + self.compaction_window_secs = other.compaction_window_secs; + self.window_sequences = std::mem::take(&mut other.window_sequences); + } else if !other.index_ids.is_empty() { + if self.compaction_window_secs == other.compaction_window_secs { + merge_window_sequences(&mut self.window_sequences, other.window_sequences); + } else { + // Do not compare maps with different window boundaries, including + // a previous incompatible merge, against current SST coverage. + self.compaction_window_secs = 0; + self.window_sequences.clear(); + } + } + self.index_ids.append(&mut other.index_ids); + } + + /// Inserts a bucket, consuming overlapping entries and expanding their interval. + /// Each map key equals the stored bucket start; adjacent intervals remain separate. + pub(crate) fn insert_into(mut self, buckets: &mut BTreeMap) { + if let Some((&previous_start, previous)) = buckets.range(..self.start).next_back() + && previous.end > self.start + { + self.start = previous_start; + } + // Expanding the end may expose more overlaps, so look up the next entry again. + while let Some((&next_start, _)) = buckets.range(self.start..self.end).next() { + if let Some(other) = buckets.remove(&next_start) { + self.merge(other); + } + } + buckets.insert(self.start, self); + } +} + +/// SSTs grouped into a half-open time interval for an aggregate series-index build. +/// +/// Reconciliation may expand the interval through existing index coverage. A changed +/// bucket is rebuilt from all its SSTs. Unknown source sequences prevent builds, +/// since their coverage cannot be determined safely. +#[derive(Debug, Clone)] +pub(crate) struct SeriesBucket { + pub(crate) start: Timestamp, + pub(crate) end: Timestamp, + pub(crate) files: SmallVec<[FileHandle; 2]>, + pub(crate) has_unknown_sequence: bool, + /// Maximum known source sequence, or zero when all sequences are unknown. + pub(crate) max_file_sequence: u64, + pub(crate) compaction_window_secs: i64, + /// SST summaries keyed by compaction-window-aligned start. + /// Each SST contributes one entry; equal starts merge by maximum end and sequence. + /// Ranges may overlap. See [`WindowSequence`] for the sequence assumption. + /// Missing SST sequences use zero as a placeholder and set `has_unknown_sequence`, + /// preventing index builds and reuse. + pub(crate) window_sequences: BTreeMap, +} + +/// Next bucket coverage and the work required to publish it, computed without I/O. +pub(crate) struct SeriesIndexPlan { + pub(crate) index_buckets: BTreeMap, + pub(crate) builds: Vec<(SeriesBucket, SeriesIndexEntry)>, + pub(crate) expired_index_ids: Vec, + /// Retire only after replacement builds and catalog publication succeed. + pub(crate) superseded_index_ids: Vec, + pub(crate) computed_buckets: usize, + pub(crate) skipped_buckets: usize, +} + +/// Plans whole-bucket replacements when source SST summaries change. The returned bucket +/// map is publishable only after every planned build and the catalog writes succeed. +pub(crate) fn plan_series_indexes( + buckets: Vec, + mut index_buckets: BTreeMap, + ttl: Option, + now_ms: i64, +) -> SeriesIndexPlan { + // Reconciliation changes geometry, not established coverage. In particular, + // a deferred bridge must not change the indexed snapshot. + let buckets = reconcile_series_buckets(buckets, &index_buckets); + let computed_buckets = buckets.len(); + let expired = |end| { + ttl.is_some_and(|ttl| { + ttl.is_expired(&end, &Timestamp::new_millisecond(now_ms)) + .unwrap_or(false) + }) + }; + let mut expired_index_ids = Vec::new(); + index_buckets.retain(|_, bucket| { + if expired(bucket.end) { + expired_index_ids.extend_from_slice(&bucket.index_ids); + false + } else { + true + } + }); + let mut builds = Vec::new(); + let mut superseded_index_ids = Vec::new(); + for bucket in buckets { + if expired(bucket.end) { + continue; + } + if bucket.has_unknown_sequence { + continue; + } + if index_buckets.get(&bucket.start).is_some_and(|indexed| { + indexed.end == bucket.end + && indexed.compaction_window_secs == bucket.compaction_window_secs + && indexed.window_sequences == bucket.window_sequences + }) { + continue; + } + if let Some(entry) = bucket.to_series_entry() { + index_buckets.retain(|_, indexed| { + if indexed.start < bucket.end && bucket.start < indexed.end { + superseded_index_ids.extend_from_slice(&indexed.index_ids); + false + } else { + true + } + }); + IndexBucket::from_entry(&entry).insert_into(&mut index_buckets); + builds.push((bucket, entry)); + } + } + index_buckets.retain(|_, bucket| !bucket.index_ids.is_empty()); + SeriesIndexPlan { + index_buckets, + skipped_buckets: computed_buckets - builds.len(), + builds, + expired_index_ids, + superseded_index_ids, + computed_buckets, + } +} + +pub(crate) fn rounded_bucket_width( + requested: Duration, + compaction_window: Duration, +) -> Option { + let window_secs = i64::try_from(compaction_window.as_secs()).ok()?.max(1); + let requested_secs = i64::try_from(requested.as_secs()) + .unwrap_or(i64::MAX) + .max(1); + let multiples = requested_secs / window_secs + i64::from(requested_secs % window_secs != 0); + multiples.checked_mul(window_secs) +} + +/// Groups SSTs across levels into sorted, disjoint time buckets without planning builds. +/// +/// Both widths must be positive, and `width_secs` must be a multiple of the compaction +/// window width. Inclusive SST ranges are rounded outward to aligned, +/// half-open intervals in seconds. Overlapping intervals merge; adjacent ones stay +/// separate. Each merged bucket tracks its maximum sequence and any unknown sequence. +pub(crate) fn group_files_into_series_buckets( + files: &[FileHandle], + width_secs: i64, + compaction_window_secs: i64, +) -> Vec { + let mut spans = files + .iter() + .map(|file| { + let start = file.time_range().0.split().0; + let end = file.time_range().1.split().0; + let sequence = file + .meta_ref() + .sequence + .map_or(0, |sequence| sequence.get()); + let first_window = start.div_euclid(compaction_window_secs); + let last_window = end.div_euclid(compaction_window_secs); + let window_sequences = BTreeMap::from([( + first_window.saturating_mul(compaction_window_secs), + WindowSequence { + start: first_window.saturating_mul(compaction_window_secs), + end: last_window + .saturating_add(1) + .saturating_mul(compaction_window_secs), + max_sequence: sequence, + }, + )]); + SeriesBucket { + start: Timestamp::new_second( + start.div_euclid(width_secs).saturating_mul(width_secs), + ), + end: Timestamp::new_second( + end.div_euclid(width_secs) + .saturating_add(1) + .saturating_mul(width_secs), + ), + files: smallvec![file.clone()], + has_unknown_sequence: file.meta_ref().sequence.is_none(), + max_file_sequence: sequence, + compaction_window_secs, + window_sequences, + } + }) + .collect::>(); + spans.sort_unstable_by(|a, b| a.start.cmp(&b.start).then_with(|| b.end.cmp(&a.end))); + group_series_buckets(spans) +} + +/// Expands SST buckets through existing index coverage before grouping build inputs. +/// +/// `buckets` must be sorted by start and disjoint. `index_buckets` must contain +/// disjoint intervals keyed by their starts. Expansion preserves start ordering, +/// but may introduce overlaps, which are merged in the returned disjoint buckets. +/// Established index coverage is borrowed so deferred builds cannot modify it. +fn reconcile_series_buckets( + mut buckets: Vec, + index_buckets: &BTreeMap, +) -> Vec { + for bucket in &mut buckets { + if let Some((&start, previous)) = index_buckets.range(..=bucket.start).next_back() + && previous.end > bucket.start + { + bucket.start = start; + bucket.end = bucket.end.max(previous.end); + } + // Disjoint index intervals make the last overlapping interval's end the + // furthest boundary. Expanding to it cannot expose another index interval. + if let Some((_, last)) = index_buckets.range(bucket.start..bucket.end).next_back() { + bucket.end = bucket.end.max(last.end); + } + } + group_series_buckets(buckets) +} + +/// Merges overlapping spans sorted by nondecreasing start; equal starts are allowed. +/// Returns sorted, disjoint buckets. Adjacent half-open intervals remain separate. +/// All spans must use the same compaction-window width for their SST summaries. +fn group_series_buckets(spans: Vec) -> Vec { + let mut buckets: Vec = Vec::new(); + for mut span in spans { + if let Some(last) = buckets.last_mut() + && span.start < last.end + { + last.end = last.end.max(span.end); + last.files.append(&mut span.files); + last.has_unknown_sequence |= span.has_unknown_sequence; + last.max_file_sequence = last.max_file_sequence.max(span.max_file_sequence); + merge_window_sequences(&mut last.window_sequences, span.window_sequences); + } else { + buckets.push(span); + } + } + buckets +} + +/// Merges summaries sharing a start, relying on [`WindowSequence`]'s sequence assumption. +fn merge_window_sequences( + target: &mut BTreeMap, + source: BTreeMap, +) { + for (start, summary) in source { + target + .entry(start) + .and_modify(|current| { + current.end = current.end.max(summary.end); + current.max_sequence = current.max_sequence.max(summary.max_sequence); + }) + .or_insert(summary); + } +} + +impl SeriesBucket { + /// Creates entry metadata with a fresh UUID and sorted source file IDs. + /// Returns `None` for unknown sequences or too few files. + fn to_series_entry(&self) -> Option { + if self.has_unknown_sequence || self.files.len() < SERIES_INDEX_TRIGGER_FILES { + return None; + } + let mut source_file_ids = self + .files + .iter() + .map(|file| file.file_id().file_id()) + .collect::>(); + source_file_ids.sort_unstable_by(|left, right| left.as_bytes().cmp(right.as_bytes())); + let min_file_sequence = self + .files + .iter() + .filter_map(|file| file.meta_ref().sequence.map(|sequence| sequence.get())) + .min()?; + Some(SeriesIndexEntry { + index_uuid: FileId::random(), + bucket_start: self.start, + bucket_end: self.end, + source_file_ids, + min_file_sequence, + max_file_sequence: self.max_file_sequence, + compaction_window_secs: self.compaction_window_secs, + window_sequences: self.window_sequences.clone(), + }) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + use std::num::NonZeroU64; + + use common_time::timestamp::TimeUnit; + + use super::*; + use crate::sst::file::FileMeta; + use crate::test_util::new_noop_file_purger; + + fn coverage(intervals: &[(i64, i64, u64)]) -> BTreeMap { + intervals + .iter() + .map(|&(start, end, max_sequence)| { + ( + start, + WindowSequence { + start, + end, + max_sequence, + }, + ) + }) + .collect() + } + + fn file(sequence: Option, level: u8, start: Timestamp, end: Timestamp) -> FileHandle { + FileHandle::new( + FileMeta { + file_id: FileId::random(), + sequence: sequence.and_then(NonZeroU64::new), + level, + time_range: (start, end), + ..Default::default() + }, + new_noop_file_purger(), + ) + } + + #[test] + fn test_second_resolution_buckets_merge_spans_across_levels() { + let width = rounded_bucket_width(Duration::from_secs(11), Duration::from_secs(10)).unwrap(); + let files = [ + file( + Some(1), + 0, + Timestamp::new_millisecond(-1), + Timestamp::new_millisecond(19999), + ), + file( + Some(2), + 1, + Timestamp::new_microsecond(19000000), + Timestamp::new_microsecond(39000000), + ), + file( + Some(3), + 2, + Timestamp::new_nanosecond(39000000000), + Timestamp::new_nanosecond(40000000000), + ), + file( + Some(4), + 1, + Timestamp::new_second(60), + Timestamp::new_second(61), + ), + ]; + let buckets = group_files_into_series_buckets(&files, width, 10); + let spans = buckets + .iter() + .map(|b| (b.start, b.end, b.files.len())) + .collect::>(); + assert_eq!( + vec![ + (Timestamp::new_second(-20), Timestamp::new_second(60), 3), + (Timestamp::new_second(60), Timestamp::new_second(80), 1), + ], + spans + ); + assert_eq!(3, buckets[0].max_file_sequence); + assert_eq!(4, buckets[1].max_file_sequence); + assert!(buckets[0].to_series_entry().is_none()); + assert!(buckets[1].to_series_entry().is_none()); + let mut files = files.to_vec(); + files.push(file( + None, + 0, + Timestamp::new_second(0), + Timestamp::new_second(1), + )); + assert!( + group_files_into_series_buckets(&files, width, 10)[0] + .to_series_entry() + .is_none() + ); + } + + #[test] + fn test_seconds_do_not_require_millisecond_conversion() { + let start = Timestamp::new_second(i64::MAX / 1000 + 100); + assert!(start.convert_to(TimeUnit::Millisecond).is_none()); + let buckets = group_files_into_series_buckets(&[file(Some(1), 0, start, start)], 1, 1); + assert_eq!(start, buckets[0].start); + assert_eq!(Timestamp::new_second(start.value() + 1), buckets[0].end); + + let width = rounded_bucket_width(Duration::ZERO, Duration::from_millis(100)).unwrap(); + let buckets = group_files_into_series_buckets( + &[file( + Some(1), + 0, + Timestamp::new_nanosecond(-1), + Timestamp::new_microsecond(1), + )], + width, + 1, + ); + assert_eq!( + (Timestamp::new_second(-1), Timestamp::new_second(1)), + (buckets[0].start, buckets[0].end) + ); + } + + #[test] + fn test_reconcile_bridge_expands_through_index_and_sst_buckets() { + let ts = Timestamp::new_second; + let mut indexes = BTreeMap::new(); + let ids = [FileId::random(), FileId::random(), FileId::random()]; + for (start, end, max_sequence, id) in [ + (0, 20, 10, ids[0]), + (30, 60, 30, ids[1]), + (70, 100, 20, ids[2]), + ] { + IndexBucket { + start: ts(start), + end: ts(end), + index_ids: smallvec![id], + compaction_window_secs: 10, + window_sequences: coverage(&[(start, end, max_sequence)]), + } + .insert_into(&mut indexes); + } + let files = [ + file(Some(15), 1, ts(10), ts(30)), + file(Some(31), 0, ts(50), ts(70)), + file(Some(32), 0, ts(90), ts(95)), + file(Some(32), 1, ts(90), ts(95)), + file(Some(33), 0, ts(90), ts(95)), + file(Some(34), 0, ts(100), ts(105)), + ]; + let planned = group_files_into_series_buckets(&files, 10, 10); + assert_eq!(4, planned.len()); + let plan = plan_series_indexes(planned, indexes, None, 0); + assert_eq!((2, 1), (plan.computed_buckets, plan.skipped_buckets)); + let [(bucket, entry)] = plan.builds.as_slice() else { + panic!("expected one replacement build"); + }; + assert_eq!((ts(0), ts(100)), (bucket.start, bucket.end)); + // Include sequence 15 and both files at 32, but exclude the adjacent bucket. + let expected = files[..5] + .iter() + .map(|file| file.file_id().file_id()) + .collect::>(); + assert_eq!( + expected, + bucket + .files + .iter() + .map(|file| file.file_id().file_id()) + .collect() + ); + assert_eq!(expected, entry.source_file_ids.iter().copied().collect()); + assert_eq!((15, 33), (entry.min_file_sequence, entry.max_file_sequence)); + assert!(plan.expired_index_ids.is_empty()); + assert_eq!(1, plan.index_buckets.len()); + let merged = &plan.index_buckets[&ts(0)]; + assert_eq!((ts(0), ts(100)), (merged.start, merged.end)); + assert_eq!(33, bucket.max_file_sequence); + assert_eq!(bucket.window_sequences, merged.window_sequences); + assert_eq!(ids.as_slice(), plan.superseded_index_ids.as_slice()); + assert_eq!([entry.index_uuid].as_slice(), merged.index_ids.as_slice()); + } + + #[test] + fn test_plan_reuses_replaced_sources_and_rebuilds_changed_bucket() { + let ts = Timestamp::new_second; + let make_file = |sequence| file(Some(sequence), 0, ts(1), ts(2)); + let original = (1..=4).map(make_file).collect::>(); + let initial = plan_series_indexes( + group_files_into_series_buckets(&original, 10, 10), + BTreeMap::new(), + None, + 0, + ); + let first_id = initial.builds[0].1.index_uuid; + + // New file IDs with already indexed sequences do not invalidate the index. + let mut files = (1..=4).map(make_file).collect::>(); + let replaced = plan_series_indexes( + group_files_into_series_buckets(&files, 10, 10), + initial.index_buckets.clone(), + None, + 0, + ); + assert!(replaced.builds.is_empty()); + assert!(replaced.expired_index_ids.is_empty()); + assert_eq!(initial.index_buckets, replaced.index_buckets); + + files.push(make_file(5)); + let ready = plan_series_indexes( + group_files_into_series_buckets(&files, 10, 10), + replaced.index_buckets, + None, + 0, + ); + let [(bucket, entry)] = ready.builds.as_slice() else { + panic!("a changed window must rebuild the whole bucket"); + }; + let expected = files[..] + .iter() + .map(|file| file.file_id().file_id()) + .collect::>(); + assert_eq!( + expected, + bucket + .files + .iter() + .map(|file| file.file_id().file_id()) + .collect() + ); + assert_eq!(expected, entry.source_file_ids.iter().copied().collect()); + assert_eq!((1, 5), (entry.min_file_sequence, entry.max_file_sequence)); + assert_eq!( + [entry.index_uuid].as_slice(), + ready.index_buckets[&ts(0)].index_ids.as_slice() + ); + assert_eq!(vec![first_id], ready.superseded_index_ids); + let repeated = plan_series_indexes( + group_files_into_series_buckets(&files, 10, 10), + ready.index_buckets.clone(), + None, + 0, + ); + assert!(repeated.builds.is_empty()); + assert_eq!(ready.index_buckets, repeated.index_buckets); + + // TTL retires the replacement; the previous index is already superseded. + let expired = plan_series_indexes( + Vec::new(), + ready.index_buckets, + Some(TimeToLive::Duration(Duration::from_secs(10))), + 21_000, + ); + assert!(expired.builds.is_empty()); + assert!(expired.index_buckets.is_empty()); + assert_eq!( + [entry.index_uuid].as_slice(), + expired.expired_index_ids.as_slice() + ); + } + + #[test] + fn test_sst_summaries_round_ranges_outward() { + let ts = Timestamp::new_millisecond; + let files = [ + file(Some(10), 0, ts(-1), ts(19_999)), + file(Some(30), 1, ts(10_000), ts(20_000)), + ]; + let buckets = group_files_into_series_buckets(&files, 100, 10); + assert_eq!(1, buckets.len()); + assert_eq!( + coverage(&[(-10, 20, 10), (10, 30, 30)]), + buckets[0].window_sequences + ); + } + + #[rstest::rstest] + #[case(32)] + #[case(33)] + #[case(100_000)] + fn test_wide_sst_reuses_unchanged_inputs_and_rebuilds_after_splitting(#[case] windows: i64) { + let ts = Timestamp::new_second; + let end = windows * 10; + let files = (1..=4) + .map(|seq| file(Some(seq), 0, ts(0), ts(end - 1))) + .collect::>(); + let initial = plan_series_indexes( + group_files_into_series_buckets(&files, end, 10), + BTreeMap::new(), + None, + 0, + ); + assert_eq!(1, initial.builds.len()); + assert_eq!( + coverage(&[(0, end, 4)]), + initial.builds[0].1.window_sequences + ); + let repeated = plan_series_indexes( + group_files_into_series_buckets(&files, end, 10), + initial.index_buckets.clone(), + None, + 0, + ); + assert!(repeated.builds.is_empty()); + assert_eq!(initial.index_buckets, repeated.index_buckets); + + // Splitting changes the summaries and conservatively triggers one rebuild. + let split = (1..=4) + .rev() + .flat_map(|seq| { + [ + file(Some(seq), 1, ts(end / 2), ts(end - 1)), + file(Some(seq), 1, ts(0), ts(end / 2 - 1)), + ] + }) + .collect::>(); + let replaced = plan_series_indexes( + group_files_into_series_buckets(&split, end, 10), + repeated.index_buckets, + None, + 0, + ); + assert_eq!(1, replaced.builds.len()); + assert_eq!(1, replaced.superseded_index_ids.len()); + let repeated = plan_series_indexes( + group_files_into_series_buckets(&split, end, 10), + replaced.index_buckets.clone(), + None, + 0, + ); + assert!(repeated.builds.is_empty()); + assert_eq!(replaced.index_buckets, repeated.index_buckets); + } + + #[rstest::rstest] + #[case::existing_start(0)] + #[case::new_start(20)] + fn test_new_data_changes_sst_summary(#[case] start: i64) { + let ts = Timestamp::new_second; + let mut files = vec![ + file(Some(1), 0, ts(1), ts(29)), + file(Some(10), 0, ts(1), ts(29)), + file(Some(30), 0, ts(1), ts(9)), + file(Some(20), 0, ts(41), ts(49)), + ]; + let initial = plan_series_indexes( + group_files_into_series_buckets(&files, 100, 10), + BTreeMap::new(), + None, + 0, + ); + assert_eq!(1, initial.builds.len()); + // The maximum end survives even when the highest sequence is in a shorter SST. + assert_eq!( + coverage(&[(0, 30, 30), (40, 50, 20)]), + initial.builds[0].1.window_sequences + ); + // New data has a sequence greater than those in the indexed snapshot. + files.push(file(Some(31), 0, ts(start + 1), ts(start + 2))); + let plan = plan_series_indexes( + group_files_into_series_buckets(&files, 100, 10), + initial.index_buckets, + None, + 0, + ); + let [(bucket, entry)] = plan.builds.as_slice() else { + panic!("new data must rebuild the bucket"); + }; + assert_eq!(files.len(), bucket.files.len()); + assert_eq!(31, entry.max_file_sequence); + let expected = if start == 0 { + coverage(&[(0, 30, 31), (40, 50, 20)]) + } else { + coverage(&[(0, 30, 30), (20, 30, 31), (40, 50, 20)]) + }; + assert_eq!(expected, entry.window_sequences); + let repeated = plan_series_indexes( + group_files_into_series_buckets(&files, 100, 10), + plan.index_buckets, + None, + 0, + ); + assert!(repeated.builds.is_empty()); + } + + #[test] + fn test_deferred_bridge_preserves_established_coverage() { + let ts = Timestamp::new_second; + let mut indexes = BTreeMap::new(); + for (start, seq) in [(0, 10), (20, 30)] { + IndexBucket { + start: ts(start), + end: ts(start + 10), + index_ids: smallvec![FileId::random()], + compaction_window_secs: 10, + window_sequences: coverage(&[(start, start + 10, seq)]), + } + .insert_into(&mut indexes); + } + let mut files = vec![ + file(Some(13), 0, ts(1), ts(21)), + file(Some(30), 0, ts(21), ts(22)), + ]; + let deferred = plan_series_indexes( + group_files_into_series_buckets(&files, 10, 10), + indexes.clone(), + None, + 0, + ); + assert!(deferred.builds.is_empty()); + assert!(deferred.superseded_index_ids.is_empty()); + assert_eq!(indexes, deferred.index_buckets); + files.extend((11..=12).map(|seq| file(Some(seq), 0, ts(1), ts(2)))); + let plan = plan_series_indexes( + group_files_into_series_buckets(&files, 10, 10), + deferred.index_buckets, + None, + 0, + ); + assert_eq!(1, plan.builds.len()); + assert_eq!(4, plan.builds[0].0.files.len()); + assert_eq!(2, plan.superseded_index_ids.len()); + assert_eq!( + coverage(&[(0, 30, 13), (20, 30, 30)]), + plan.builds[0].1.window_sequences + ); + + files.push(file(None, 0, ts(1), ts(2))); + let unknown = plan_series_indexes( + group_files_into_series_buckets(&files, 10, 10), + indexes.clone(), + None, + 0, + ); + assert!(unknown.builds.is_empty()); + assert_eq!(indexes, unknown.index_buckets); + } + + #[rstest::rstest] + #[case::removed_window(100, 10, false)] + #[case::added_window(100, 10, true)] + #[case::window_width(100, 20, false)] + #[case::bucket_width(200, 10, false)] + fn test_coverage_shape_changes_rebuild( + #[case] bucket_width: i64, + #[case] window_width: i64, + #[case] add_window: bool, + ) { + let ts = Timestamp::new_second; + let mut files = (27..=30) + .map(|seq| file(Some(seq), 0, ts(11), ts(12))) + .collect::>(); + let mut original = files.clone(); + original.push(file(Some(10), 0, ts(1), ts(2))); + let initial = plan_series_indexes( + group_files_into_series_buckets(&original, 100, 10), + BTreeMap::new(), + None, + 0, + ); + if add_window { + files = original; + files.push(file(Some(15), 0, ts(21), ts(22))); + } else if bucket_width != 100 || window_width != 10 { + files = original; + } + let plan = plan_series_indexes( + group_files_into_series_buckets(&files, bucket_width, window_width), + initial.index_buckets, + None, + 0, + ); + assert_eq!(1, plan.builds.len()); + assert_eq!(1, plan.superseded_index_ids.len()); + let repeated = plan_series_indexes( + group_files_into_series_buckets(&files, bucket_width, window_width), + plan.index_buckets, + None, + 0, + ); + assert!(repeated.builds.is_empty()); + } + + #[test] + fn test_index_map_merge_is_order_independent() { + let ts = Timestamp::new_second; + let make_index = |start, end, width, windows: &[(i64, u64)]| IndexBucket { + start: ts(start), + end: ts(end), + index_ids: smallvec![FileId::random()], + compaction_window_secs: width, + window_sequences: windows + .iter() + .map(|&(start, max_sequence)| { + ( + start, + WindowSequence { + start, + end: start + width, + max_sequence, + }, + ) + }) + .collect(), + }; + let indexes = [ + make_index(0, 20, 10, &[(0, 10), (10, 20)]), + make_index(10, 30, 10, &[(10, 30), (20, 15)]), + make_index(20, 40, 10, &[(20, 25), (30, 5)]), + ]; + for order in [[0, 1, 2], [2, 1, 0], [1, 0, 2], [0, 2, 1]] { + let mut map = BTreeMap::new(); + for i in order { + indexes[i].clone().insert_into(&mut map); + } + assert_eq!(1, map.len()); + assert_eq!(10, map[&ts(0)].compaction_window_secs); + assert_eq!( + coverage(&[(0, 10, 10), (10, 20, 30), (20, 30, 25), (30, 40, 5)]), + map[&ts(0)].window_sequences + ); + make_index(10, 30, 20, &[(0, 30)]).insert_into(&mut map); + indexes[0].clone().insert_into(&mut map); + assert_eq!(0, map[&ts(0)].compaction_window_secs); + assert!(map[&ts(0)].window_sequences.is_empty()); + } + } +} diff --git a/src/mito2/src/series_index/builder.rs b/src/mito2/src/series_index/builder.rs new file mode 100644 index 00000000000..8d361be364e --- /dev/null +++ b/src/mito2/src/series_index/builder.rs @@ -0,0 +1,742 @@ +// 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. + +//! Independently builds range and aggregate series indexes from SST readers. + +use std::sync::Arc; + +use async_stream::try_stream; +use common_telemetry::warn; +use futures::TryStreamExt; +use object_store::ObjectStore; +use snafu::{OptionExt, ensure}; +use store_api::storage::FileId; + +use crate::error::{Result, UnexpectedSnafu}; +use crate::read::BoxedRecordBatchStream; +use crate::read::compat::FlatCompatBatch; +use crate::read::flat_merge::FlatMergeReader; +use crate::read::flat_projection::FlatProjectionMapper; +use crate::read::prune::FlatPruneReader; +use crate::read::read_columns::ReadColumns; +use crate::region::MitoRegionRef; +use crate::region::version::VersionRef; +use crate::series_index::bucket::SeriesBucket; +use crate::series_index::catalog::{ + SeriesIndexEntry, range_index_path, series_index_path, series_metadata, +}; +use crate::series_index::purger::{IndexFilePurger, IndexFileType, file_operation}; +use crate::series_index::version::SeriesIndexFileHandle; +use crate::series_index::{SeriesIndexWriter, SeriesIndexWriterOptions}; +use crate::sst::file::FileHandle; +use crate::sst::parquet::reader::{FlatRowGroupReader, ReaderMetrics}; +use crate::sst::parquet::row_group::ParquetFetchMetrics; +use crate::sst::range_index::{SstRangeIndexWriter, SstRangeIndexWriterOptions}; + +async fn reader_input( + region: &MitoRegionRef, + file: FileHandle, +) -> Result< + Option<( + Arc, + crate::sst::parquet::row_selection::RowGroupSelection, + )>, +> { + Ok(region + .access_layer + .read_sst(file) + .projection(Some(ReadColumns::new([]))) + .build_reader_input(&mut ReaderMetrics::default()) + .await? + .map(|(context, selection)| (Arc::new(context), selection))) +} + +/// Builds one range index, or returns `None` when the SST has no readable input. +pub(crate) async fn build_range_index( + store: &ObjectStore, + region: &MitoRegionRef, + version: &VersionRef, + file: FileHandle, +) -> Result> { + let file_id = file.file_id().file_id(); + let Some((context, mut selection)) = reader_input(region, file).await? else { + return Ok(None); + }; + let mapper = FlatProjectionMapper::new(&version.metadata, [])?; + let compat = FlatCompatBatch::try_new(&mapper, context.read_format(), false)?; + let path = range_index_path(region.region_id, file_id); + let mut writer = SstRangeIndexWriter::try_new( + version.metadata.clone(), + store.clone(), + &path, + SstRangeIndexWriterOptions::default(), + ) + .await?; + let result: Result<()> = async { + let fetch_metrics = ParquetFetchMetrics::default(); + while let Some((row_group_id, row_selection)) = selection.pop_first() { + let parquet_reader = context + .reader_builder() + .build(context.build_context( + row_group_id, + Some(row_selection), + Some(&fetch_metrics), + )) + .await?; + let mut reader = FlatPruneReader::new_with_row_group_reader( + context.clone(), + FlatRowGroupReader::new(context.clone(), parquet_reader), + context.pre_filter_mode().skip_fields(), + ); + while let Some(batch) = reader.next_batch().await? { + let batch = match &compat { + Some(compat) => compat.compat(batch)?, + None => batch, + }; + writer.write(row_group_id as u32, &batch).await?; + } + } + Ok(()) + } + .await; + if let Err(error) = result { + if let Err(cleanup_error) = writer.abort().await { + warn!(cleanup_error; "Failed to abort range-index build"); + } + return Err(error); + } + writer.finish().await?; + file_operation(IndexFileType::Range, "build", "success"); + Ok(Some(file_id)) +} + +/// Builds only the series index. Callers build needed range indexes separately. +pub(crate) async fn build_series_index( + store: &ObjectStore, + region: &MitoRegionRef, + version: &VersionRef, + bucket: &SeriesBucket, + entry: &SeriesIndexEntry, + purger: &IndexFilePurger, +) -> Result { + let mut sources = Vec::::new(); + let mapper = FlatProjectionMapper::new(&version.metadata, [])?; + let schema = mapper.input_arrow_schema(false); + for file in &bucket.files { + let Some((context, mut selection)) = reader_input(region, file.clone()).await? else { + continue; + }; + let compat = FlatCompatBatch::try_new(&mapper, context.read_format(), false)?; + sources.push(Box::pin(try_stream! { + let fetch_metrics = ParquetFetchMetrics::default(); + while let Some((row_group_id, row_selection)) = selection.pop_first() { + let parquet_reader = context.reader_builder().build(context.build_context( + row_group_id, + Some(row_selection), + Some(&fetch_metrics), + )).await?; + let mut reader = FlatPruneReader::new_with_row_group_reader( + context.clone(), + FlatRowGroupReader::new(context.clone(), parquet_reader), + context.pre_filter_mode().skip_fields(), + ); + while let Some(batch) = reader.next_batch().await? { + yield match &compat { + Some(compat) => compat.compat(batch)?, + None => batch, + }; + } + } + })); + } + ensure!( + !sources.is_empty(), + UnexpectedSnafu { + reason: "series-index bucket has no readable SST", + } + ); + let mut visible: BoxedRecordBatchStream = if sources.len() == 1 { + sources.pop().context(UnexpectedSnafu { + reason: "series-index source disappeared", + })? + } else { + Box::pin( + FlatMergeReader::new(schema, sources, 8192, None) + .await? + .into_stream(), + ) + }; + // TODO(yingwen): Deduplicate update-mode rows before series indexes are used by queries. + let path = series_index_path(region.region_id, entry.index_uuid); + let mut writer = SeriesIndexWriter::try_new( + version.metadata.clone(), + store.clone(), + &path, + SeriesIndexWriterOptions::default(), + Some(series_metadata(entry)?), + ) + .await?; + let result: Result<()> = async { + while let Some(batch) = visible.try_next().await? { + writer.write(&batch).await?; + } + Ok(()) + } + .await; + if let Err(error) = result { + if let Err(cleanup_error) = writer.abort().await { + warn!(cleanup_error; "Failed to abort series-index build"); + } + return Err(error); + } + writer.finish().await?; + file_operation(IndexFileType::Series, "build", "success"); + Ok(SeriesIndexFileHandle::new( + region.region_id, + entry.clone(), + purger.clone(), + )) +} + +#[cfg(test)] +mod tests { + use std::collections::{BTreeMap, HashMap}; + use std::sync::Mutex; + use std::sync::atomic::{AtomicBool, Ordering}; + + use datatypes::data_type::ConcreteDataType; + use object_store::layers::mock::{self, MockLayerBuilder, oio}; + use object_store::services::Memory; + use store_api::region_engine::RegionEngine; + + use super::*; + use crate::series_index::bucket::{group_files_into_series_buckets, plan_series_indexes}; + use crate::series_index::catalog::{ + SeriesIndexCatalog, load_version_control, series_catalog_path, store_catalog, + }; + use crate::series_index::purger::series_index_channel; + use crate::series_index::tests::prepare_region; + use crate::test_util::TestEnv; + + fn build_input(version: &VersionRef) -> (SeriesBucket, SeriesIndexEntry) { + let files = version + .ssts + .levels() + .iter() + .flat_map(|level| level.files()) + .cloned() + .collect::>(); + let mut plan = plan_series_indexes( + group_files_into_series_buckets(&files, 100, 10), + BTreeMap::new(), + None, + 0, + ); + assert_eq!(1, plan.builds.len()); + plan.builds.pop().unwrap() + } + + fn metadata_with_seconds(version: &VersionRef) -> store_api::metadata::RegionMetadataRef { + let mut metadata = (*version.metadata).clone(); + let time_index = metadata.time_index_column_pos(); + metadata.column_metadatas[time_index] + .column_schema + .data_type = ConcreteDataType::timestamp_second_datatype(); + Arc::new(metadata) + } + + #[rstest::rstest] + #[case::range_then_series(true, false)] + #[case::series_only(false, false)] + #[case::series_failure(true, true)] + #[tokio::test] + async fn test_build_indexes_without_publication( + #[case] build_ranges: bool, + #[case] fail_series: bool, + ) { + let mut env = TestEnv::with_prefix("series-builder").await; + let (engine, region) = prepare_region(&mut env).await; + let version = region.version(); + let (bucket, entry) = build_input(&version); + let files = &bucket.files; + let store = ObjectStore::new(Memory::default()).unwrap(); + let (purger, mut receiver) = series_index_channel(store.clone()); + let mut range_bytes = HashMap::new(); + if build_ranges { + for file in files { + let id = build_range_index(&store, ®ion, &version, file.clone()) + .await + .unwrap() + .unwrap(); + range_bytes.insert( + id, + store + .read(&range_index_path(region.region_id, id)) + .await + .unwrap() + .to_bytes(), + ); + } + } + // Both builders use the captured version, even if live metadata changes. + region + .version_control + .alter_metadata(metadata_with_seconds(&version)); + let layer = writer_layer(&WriterStates::default(), move |path| { + assert!(path.contains("/series/"), "series stage opened {path}"); + if fail_series { + WriterFailure::Finish + } else { + WriterFailure::None + } + }); + let result = build_series_index( + &store.clone().layer(layer), + ®ion, + &version, + &bucket, + &entry, + &purger, + ) + .await; + if fail_series { + assert!(result.is_err()); + } else { + let handle = result.unwrap(); + assert_eq!(handle.entry(), &entry); + assert!( + store + .exists(&series_index_path(region.region_id, entry.index_uuid)) + .await + .unwrap() + ); + store_catalog( + &store, + &series_catalog_path(region.region_id), + &SeriesIndexCatalog { + indexes: vec![handle.entry().clone()], + }, + ) + .await + .unwrap(); + let recovered = load_version_control(&store, region.region_id, &purger) + .await + .current(); + let repeated = plan_series_indexes( + group_files_into_series_buckets(files, 100, 10), + recovered.index_buckets.clone(), + None, + 0, + ); + assert!(repeated.builds.is_empty()); + assert_eq!(recovered.index_buckets, repeated.index_buckets); + } + // Series success or failure leaves completed range indexes unchanged. + for file in files { + let id = file.file_id().file_id(); + let path = range_index_path(region.region_id, id); + if let Some(bytes) = range_bytes.get(&id) { + assert_eq!(*bytes, store.read(&path).await.unwrap().to_bytes()); + } else { + assert!(!store.exists(&path).await.unwrap()); + } + } + assert!(region.series_index_version().range_indexes.is_empty()); + assert!(region.series_index_version().series_indexes.is_empty()); + assert!(receiver.try_recv().is_err()); + engine.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_build_indexes_after_time_index_widening() { + use api::v1::helper::row; + use api::v1::value::ValueData; + use api::v1::{ColumnDataType, Rows, SemanticType, WriteHint}; + use datatypes::arrow::array::{TimestampMicrosecondArray, UInt64Array}; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use store_api::region_request::{ + AlterKind, ModifyColumnType, RegionAlterRequest, RegionPutRequest, RegionRequest, + }; + use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; + + use crate::test_util::sst_util::new_sparse_primary_key; + use crate::test_util::{CreateRequestBuilder, flush_region, rows_schema}; + + let mut env = TestEnv::with_prefix("series-builder-widen").await; + let (engine, region) = prepare_region(&mut env).await; + let metadata = region.version().metadata.clone(); + engine + .handle_request( + region.region_id, + RegionRequest::Alter(RegionAlterRequest { + kind: AlterKind::ModifyColumnTypes { + columns: vec![ModifyColumnType { + column_name: metadata.time_index_column().column_schema.name.clone(), + target_type: ConcreteDataType::timestamp_microsecond_datatype(), + }], + }, + }), + ) + .await + .unwrap(); + + let mut request = CreateRequestBuilder::new().build(); + request.column_metadatas = metadata.column_metadatas.clone(); + request.primary_key = metadata.primary_key.clone(); + let full_schema = rows_schema(&request); + let mut pk = full_schema[0].clone(); + pk.column_name = PRIMARY_KEY_COLUMN_NAME.to_string(); + pk.datatype = ColumnDataType::Binary.into(); + pk.semantic_type = SemanticType::Tag.into(); + let mut ts = full_schema[5].clone(); + ts.datatype = ColumnDataType::TimestampMicrosecond.into(); + engine + .handle_request( + region.region_id, + RegionRequest::Put(RegionPutRequest { + skip_wal: false, + rows: Rows { + schema: vec![pk, ts, full_schema[4].clone()], + rows: [500_500, 2_500_500, 3_500_500] + .into_iter() + .map(|ts| { + row(vec![ + ValueData::BinaryValue(new_sparse_primary_key( + &["a", "x"], + &metadata, + 10, + 0, + )), + ValueData::TimestampMicrosecondValue(ts), + ValueData::U64Value(1), + ]) + }) + .collect(), + }, + hint: Some(WriteHint { + primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse.into(), + }), + partition_expr_version: None, + }), + ) + .await + .unwrap(); + flush_region(&engine, region.region_id, None).await; + let version = region.version(); + let (bucket, entry) = build_input(&version); + let store = ObjectStore::new(Memory::default()).unwrap(); + let (purger, _receiver) = series_index_channel(store.clone()); + for file in &bucket.files { + build_range_index(&store, ®ion, &version, file.clone()) + .await + .unwrap() + .unwrap(); + } + let _handle = build_series_index(&store, ®ion, &version, &bucket, &entry, &purger) + .await + .unwrap(); + let bytes = store + .read(&series_index_path(region.region_id, entry.index_uuid)) + .await + .unwrap() + .to_bytes(); + let batches = ParquetRecordBatchReaderBuilder::try_new(bytes) + .unwrap() + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(1, batches.len()); + let batch = &batches[0]; + assert_eq!(1, batch.num_rows()); + for (column, expected) in [(0, 500_500), (1, 4_000_000)] { + assert_eq!( + expected, + batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + ); + } + assert_eq!( + 7, + batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + ); + engine.stop().await.unwrap(); + } + + #[tokio::test] + async fn test_range_failure_preserves_completed_indexes_and_stops_series_stage() { + let mut env = TestEnv::with_prefix("range-builder-stage-failure").await; + let (engine, region) = prepare_region(&mut env).await; + let version = region.version(); + let (bucket, entry) = build_input(&version); + let files = &bucket.files; + let failed_path = range_index_path(region.region_id, files[1].file_id().file_id()); + let states = WriterStates::default(); + let layer = writer_layer(&states, move |path| { + if path == failed_path { + WriterFailure::Finish + } else { + WriterFailure::None + } + }); + let store = ObjectStore::new(Memory::default()).unwrap().layer(layer); + let (purger, _receiver) = series_index_channel(store.clone()); + let mut completed = Vec::new(); + let result: Result> = async { + for file in files { + let Some(id) = build_range_index(&store, ®ion, &version, file.clone()).await? + else { + // Defer series construction if a needed range is not ready. + return Ok(None); + }; + completed.push(id); + } + build_series_index(&store, ®ion, &version, &bucket, &entry, &purger) + .await + .map(Some) + } + .await; + assert!(format!("{:?}", result.unwrap_err()).contains("injected index finish failure")); + assert_eq!(vec![files[0].file_id().file_id()], completed); + { + let states = states.lock().unwrap(); + assert_eq!(2, states.len()); + assert!(states.keys().all(|path| !path.contains("/series/"))); + assert_eq!( + 1, + states[&range_index_path(region.region_id, completed[0])].closed + ); + assert_eq!( + 1, + states[&range_index_path(region.region_id, files[1].file_id().file_id())].aborted + ); + } + for (i, file) in files.iter().enumerate() { + assert_eq!( + i == 0, + store + .exists(&range_index_path( + region.region_id, + file.file_id().file_id() + )) + .await + .unwrap() + ); + } + engine.stop().await.unwrap(); + } + + #[derive(Default)] + struct WriterState { + closed: usize, + aborted: usize, + } + + type WriterStates = Arc>>; + + enum WriterFailure { + None, + Finish, + Abort, + } + + fn writer_layer( + states: &WriterStates, + failure: impl Fn(&str) -> WriterFailure + Send + Sync + 'static, + ) -> mock::MockLayer { + let states = states.clone(); + MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _, inner| -> oio::Writer { + states + .lock() + .unwrap() + .insert(path.to_string(), WriterState::default()); + Box::new(RecordingWriter { + inner, + path: path.to_string(), + states: states.clone(), + failure: failure(path), + }) + })) + .build() + .unwrap() + } + + struct RecordingWriter { + inner: oio::Writer, + path: String, + states: WriterStates, + failure: WriterFailure, + } + + impl mock::Write for RecordingWriter { + async fn write(&mut self, buffer: mock::Buffer) -> mock::Result<()> { + self.inner.write(buffer).await + } + + async fn close(&mut self) -> mock::Result { + if matches!(self.failure, WriterFailure::Finish) { + return Err(mock::Error::new( + mock::ErrorKind::Unexpected, + "injected index finish failure", + )); + } + let metadata = self.inner.close().await?; + self.states + .lock() + .unwrap() + .get_mut(&self.path) + .unwrap() + .closed += 1; + Ok(metadata) + } + + async fn abort(&mut self) -> mock::Result<()> { + self.states + .lock() + .unwrap() + .get_mut(&self.path) + .unwrap() + .aborted += 1; + if matches!(self.failure, WriterFailure::Abort) { + return Err(mock::Error::new( + mock::ErrorKind::Unexpected, + "injected abort failure", + )); + } + self.inner.abort().await + } + } + + struct FailingSstReader { + inner: oio::Reader, + fail: Arc, + } + + impl mock::Read for FailingSstReader { + async fn read( + &self, + range: mock::BytesRange, + ) -> mock::Result<(mock::RpRead, mock::Buffer)> { + if self.fail.load(Ordering::Relaxed) { + return Err(mock::Error::new( + mock::ErrorKind::Unexpected, + "injected SST read failure", + )); + } + self.inner.read(range).await + } + + async fn open( + &self, + range: mock::BytesRange, + ) -> mock::Result<(mock::RpRead, Box)> { + if self.fail.load(Ordering::Relaxed) { + return Err(mock::Error::new( + mock::ErrorKind::Unexpected, + "injected SST read failure", + )); + } + self.inner.open(range).await + } + } + + #[rstest::rstest] + #[case::range(false, false, false)] + #[case::range_abort_failure(false, false, true)] + #[case::series_setup(true, true, false)] + #[case::series_read(true, false, false)] + #[case::series_abort_failure(true, false, true)] + #[tokio::test] + async fn test_read_failure_aborts_unfinished_outputs( + #[case] series: bool, + #[case] fail_before_open: bool, + #[case] fail_abort: bool, + ) { + let fail = Arc::new(AtomicBool::new(false)); + let read_fail = fail.clone(); + let read_layer = MockLayerBuilder::default() + .reader_factory(Arc::new(move |path, _, inner| -> oio::Reader { + if path.ends_with(".parquet") { + Box::new(FailingSstReader { + inner, + fail: read_fail.clone(), + }) + } else { + inner + } + })) + .build() + .unwrap(); + let mut env = TestEnv::with_prefix("series-builder-read-failure") + .await + .with_mock_layer(read_layer); + let (engine, region) = prepare_region(&mut env).await; + let version = region.version(); + let (mut bucket, mut entry) = build_input(&version); + // A single source is first polled after opening the series writer. + bucket.files.truncate(1); + entry.source_file_ids = vec![bucket.files[0].file_id().file_id()]; + let states = WriterStates::default(); + let write_fail = fail.clone(); + let write_layer = writer_layer(&states, move |_| { + write_fail.store(true, Ordering::Relaxed); + if fail_abort { + WriterFailure::Abort + } else { + WriterFailure::None + } + }); + let store = ObjectStore::new(Memory::default()) + .unwrap() + .layer(write_layer); + fail.store(fail_before_open, Ordering::Relaxed); + let result = if !series { + build_range_index(&store, ®ion, &version, bucket.files[0].clone()) + .await + .map(|_| ()) + } else { + let (purger, _receiver) = series_index_channel(store.clone()); + build_series_index(&store, ®ion, &version, &bucket, &entry, &purger) + .await + .map(|_| ()) + }; + fail.store(false, Ordering::Relaxed); + let error = result.unwrap_err(); + assert!( + format!("{error:?}").contains("injected SST read failure"), + "{error:?}" + ); + { + let states = states.lock().unwrap(); + assert_eq!(usize::from(!fail_before_open), states.len()); + for (path, state) in states.iter() { + assert_eq!(0, state.closed, "{path}"); + assert_eq!(1, state.aborted, "{path}"); + } + } + assert!(store.list("/").await.unwrap().is_empty()); + assert!(region.series_index_version().series_indexes.is_empty()); + engine.stop().await.unwrap(); + } +} diff --git a/src/mito2/src/series_index/catalog.rs b/src/mito2/src/series_index/catalog.rs index a59c1926fc0..56c61da54f9 100644 --- a/src/mito2/src/series_index/catalog.rs +++ b/src/mito2/src/series_index/catalog.rs @@ -14,6 +14,8 @@ //! Index catalog persistence, coverage metadata, and file paths. +use std::collections::BTreeMap; + use common_telemetry::warn; use common_time::Timestamp; use object_store::{ErrorKind, ObjectStore}; @@ -28,11 +30,29 @@ use crate::series_index::purger::IndexFilePurger; use crate::series_index::version::{ SeriesIndexFileHandle, SeriesIndexVersion, SeriesIndexVersionControl, }; +pub(crate) use crate::sst::range_index::range_index_path; + const SERIES_DIR: &str = "series"; const RANGE_CATALOG: &str = "range-index.json"; const SERIES_CATALOG: &str = "series-index.json"; const SERIES_METADATA_KEY: &str = "greptime.series_index"; +/// Summary of SSTs sharing a compaction-window-aligned start. +/// +/// New data changes must have sequences greater than those already indexed. A new +/// start adds a map entry; new data at an existing start raises its maximum sequence. +/// The maximum end and sequence may come from different files, so this summary +/// does not establish uniform sequence coverage throughout the interval. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct WindowSequence { + /// Inclusive start in epoch seconds, equal to the key in `window_sequences`. + pub(crate) start: i64, + /// Exclusive interval end in epoch seconds. + pub(crate) end: i64, + /// Maximum sequence among the SSTs sharing this aligned start. + pub(crate) max_sequence: u64, +} + /// Self-describing coverage stored in a series-index Parquet footer. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub(crate) struct SeriesIndexEntry { @@ -41,9 +61,19 @@ pub(crate) struct SeriesIndexEntry { pub(crate) bucket_start: Timestamp, /// Exclusive bucket end. pub(crate) bucket_end: Timestamp, + /// Source SST IDs retained for debugging only. Compaction can replace these + /// files without changing indexed data, so IDs must not determine index reuse. pub(crate) source_file_ids: Vec, pub(crate) min_file_sequence: u64, pub(crate) max_file_sequence: u64, + /// Width used to align the half-open compaction windows, in seconds. + pub(crate) compaction_window_secs: i64, + /// Source SST summaries keyed by aligned start; intervals may overlap. + /// Each file contributes one summary regardless of its span. Equal starts merge + /// by taking the maximum end and sequence. See [`WindowSequence`] for the + /// sequence assumption used to detect new data. Compaction changing summary + /// boundaries may conservatively trigger a rebuild. + pub(crate) window_sequences: BTreeMap, } #[derive(Debug, Default, Serialize, Deserialize)] @@ -135,9 +165,9 @@ pub(crate) async fn load_version_control( .await .unwrap_or_default(); // TODO: Handle catalog entries whose index files are missing from storage. - let version = SeriesIndexVersion { - range_indexes: range.indexes.into_iter().collect(), - series_indexes: series + let version = SeriesIndexVersion::new( + range.indexes.into_iter().collect(), + series .indexes .into_iter() .map(|entry| { @@ -147,7 +177,7 @@ pub(crate) async fn load_version_control( ) }) .collect(), - }; + ); let control = SeriesIndexVersionControl::default(); control.publish(std::sync::Arc::new(version)); control @@ -155,6 +185,7 @@ pub(crate) async fn load_version_control( #[cfg(test)] mod tests { + use std::collections::BTreeMap; use std::sync::Arc; use common_time::Timestamp; @@ -164,8 +195,9 @@ mod tests { use store_api::storage::{FileId, RegionId}; use crate::series_index::catalog::{ - SeriesIndexCatalog, SeriesIndexEntry, load_catalog, load_version_control, - series_catalog_path, series_metadata, store_catalog, + RangeIndexCatalog, SeriesIndexCatalog, SeriesIndexEntry, WindowSequence, load_catalog, + load_version_control, range_catalog_path, series_catalog_path, series_metadata, + store_catalog, }; use crate::series_index::purger::series_index_channel; @@ -194,31 +226,54 @@ mod tests { } #[tokio::test] - async fn test_load_catalog_returns_none_on_error() { + async fn test_load_catalog_defaults_on_missing_invalid_or_unreadable_catalog() { let store = ObjectStore::new(Memory::default()).unwrap(); - let path = series_catalog_path(RegionId::new(1, 1)); - // Missing catalog. + let region_id = RegionId::new(1, 1); + let (purger, _receiver) = series_index_channel(store.clone()); assert!( - load_catalog::(&store, &path) - .await - .is_none() - ); - store.write(&path, "invalid").await.unwrap(); - assert!( - load_catalog::(&store, &path) + load_catalog::(&store, &series_catalog_path(region_id)) .await .is_none() ); + let control = load_version_control(&store, region_id, &purger).await; + assert!(control.current().range_indexes.is_empty()); + assert!(control.current().series_indexes.is_empty()); + let file_id = FileId::random(); + store + .write( + &range_catalog_path(region_id), + serde_json::to_vec(&RangeIndexCatalog { + indexes: vec![file_id], + }) + .unwrap(), + ) + .await + .unwrap(); + store + .write(&series_catalog_path(region_id), "invalid") + .await + .unwrap(); + let control = load_version_control(&store, region_id, &purger).await; + assert!(control.current().range_indexes.contains(&file_id)); + assert!(control.current().series_indexes.is_empty()); let layer = MockLayerBuilder::default() .reader_factory(Arc::new(|_, _, _| Box::new(FailingCatalogReader))) .build() .unwrap(); - let store = store.layer(layer); assert!( - load_catalog::(&store, &path) + load_catalog::(&store, &series_catalog_path(region_id)) .await .is_none() ); + let store = store.layer(layer); + assert!( + load_catalog::(&store, &range_catalog_path(region_id)) + .await + .is_none() + ); + let control = load_version_control(&store, region_id, &purger).await; + assert!(control.current().range_indexes.is_empty()); + assert!(control.current().series_indexes.is_empty()); } #[tokio::test] @@ -232,6 +287,25 @@ mod tests { source_file_ids: vec![FileId::random()], min_file_sequence: 1, max_file_sequence: 2, + compaction_window_secs: 10, + window_sequences: BTreeMap::from([ + ( + 0, + WindowSequence { + start: 0, + end: 20, + max_sequence: 1, + }, + ), + ( + 20, + WindowSequence { + start: 20, + end: 100, + max_sequence: 2, + }, + ), + ]), }; store_catalog( &store, @@ -248,6 +322,13 @@ mod tests { .current(); assert!(current.range_indexes.is_empty()); assert_eq!(&entry, current.series_indexes[&entry.index_uuid].entry()); + assert_eq!(current.index_buckets.len(), 1); + let bucket = ¤t.index_buckets[&entry.bucket_start]; + assert_eq!(bucket.start, entry.bucket_start); + assert_eq!(bucket.end, entry.bucket_end); + assert_eq!(bucket.index_ids.as_slice(), &[entry.index_uuid]); + assert_eq!(bucket.compaction_window_secs, entry.compaction_window_secs); + assert_eq!(bucket.window_sequences, entry.window_sequences); let metadata = series_metadata(&entry).unwrap(); let decoded: SeriesIndexEntry = diff --git a/src/mito2/src/series_index/tests.rs b/src/mito2/src/series_index/tests.rs new file mode 100644 index 00000000000..fa65dadcddb --- /dev/null +++ b/src/mito2/src/series_index/tests.rs @@ -0,0 +1,114 @@ +// 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. + +//! Shared fixtures for series-index construction tests. + +use std::sync::Arc; + +use api::v1::helper::row; +use api::v1::value::ValueData; +use api::v1::{ColumnDataType, Rows, SemanticType, WriteHint}; +use store_api::codec::PrimaryKeyEncoding; +use store_api::metric_engine_consts::PRIMARY_KEY_ENCODING; +use store_api::region_engine::RegionEngine; +use store_api::region_request::{RegionPutRequest, RegionRequest}; +use store_api::storage::RegionId; +use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; + +use crate::config::MitoConfig; +use crate::engine::MitoEngine; +use crate::region::MitoRegionRef; +use crate::test_util::sst_util::{new_sparse_primary_key, sst_region_metadata_with_encoding}; +use crate::test_util::{CreateRequestBuilder, TestEnv, flush_region, rows_schema}; + +/// Builds real sparse SSTs; background maintenance is disabled so tests control publication. +pub(super) async fn prepare_region(env: &mut TestEnv) -> (MitoEngine, MitoRegionRef) { + prepare_region_with_timestamps(env, &[1000, 2000, 3000, 4000]).await +} + +async fn prepare_region_with_timestamps( + env: &mut TestEnv, + timestamps: &[i64], +) -> (MitoEngine, MitoRegionRef) { + let engine = env.create_engine(MitoConfig::default()).await; + let metadata = Arc::new(sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + )); + let region_id = RegionId::new(1, 1); + let mut request = CreateRequestBuilder::new().build(); + request.column_metadatas = metadata.column_metadatas.clone(); + request.primary_key = metadata.primary_key.clone(); + request + .options + .insert(PRIMARY_KEY_ENCODING.to_string(), "sparse".to_string()); + request + .options + .insert("memtable.type".to_string(), "bulk".to_string()); + request + .options + .insert("sst_format".to_string(), "flat".to_string()); + request + .options + .insert("compaction.type".to_string(), "twcs".to_string()); + request.options.insert( + "compaction.twcs.time_window".to_string(), + "100s".to_string(), + ); + // Keep the source SSTs stable while tests reconcile multiple buckets. + request.options.insert( + "compaction.twcs.trigger_file_num".to_string(), + "100".to_string(), + ); + let full_schema = rows_schema(&request); + let mut pk_column = full_schema[0].clone(); + pk_column.column_name = PRIMARY_KEY_COLUMN_NAME.to_string(); + pk_column.datatype = ColumnDataType::Binary.into(); + pk_column.semantic_type = SemanticType::Tag.into(); + let schema = vec![pk_column, full_schema[5].clone(), full_schema[4].clone()]; + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + for &ts in timestamps { + engine + .handle_request( + region_id, + RegionRequest::Put(RegionPutRequest { + skip_wal: false, + rows: Rows { + schema: schema.clone(), + rows: vec![row(vec![ + ValueData::BinaryValue(new_sparse_primary_key( + &["a", "x"], + &metadata, + 10, + 0, + )), + ValueData::TimestampMillisecondValue(ts), + ValueData::U64Value(1), + ])], + }, + hint: Some(WriteHint { + primary_key_encoding: api::v1::PrimaryKeyEncoding::Sparse.into(), + }), + partition_expr_version: None, + }), + ) + .await + .unwrap(); + flush_region(&engine, region_id, None).await; + } + let region = engine.get_region(region_id).unwrap(); + (engine, region) +} diff --git a/src/mito2/src/series_index/version.rs b/src/mito2/src/series_index/version.rs index 7db23da2f93..4db883e7626 100644 --- a/src/mito2/src/series_index/version.rs +++ b/src/mito2/src/series_index/version.rs @@ -14,13 +14,15 @@ //! Immutable index snapshots and aggregate series-file handles. -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::fmt::{self, Debug, Formatter}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, RwLock}; +use common_time::Timestamp; use store_api::storage::{FileId, RegionId}; +use crate::series_index::bucket::IndexBucket; use crate::series_index::catalog::SeriesIndexEntry; use crate::series_index::purger::{IndexFilePurger, PurgeRequest}; use crate::sst::file::RegionFileId; @@ -85,11 +87,30 @@ impl Drop for SeriesIndexFileHandleInner { /// Immutable series-index snapshot for one region. #[derive(Debug, Default)] pub(crate) struct SeriesIndexVersion { + /// Range indexes for visible SSTs; reconciliation removes IDs absent from its SST snapshot. + /// Physical deletion is independently handled by the SST file purger. pub(crate) range_indexes: HashSet, pub(crate) series_indexes: HashMap, + pub(crate) index_buckets: BTreeMap, } impl SeriesIndexVersion { + /// Restores bucket lookup from immutable index coverage stored in the catalog. + pub(crate) fn new( + range_indexes: HashSet, + series_indexes: HashMap, + ) -> Self { + let mut index_buckets = BTreeMap::new(); + for handle in series_indexes.values() { + IndexBucket::from_entry(handle.entry()).insert_into(&mut index_buckets); + } + Self { + range_indexes, + series_indexes, + index_buckets, + } + } + fn mark_all_deleted(&self) { self.series_indexes .values() diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 3d90a360925..87eeb7bae5d 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -17,9 +17,14 @@ use std::collections::BTreeMap; use std::sync::Arc; +use api::v1::SemanticType; use common_base::readable_size::ReadableSize; use datatypes::json::JsonSettings; use parquet::file::metadata::ParquetMetaData; +use parquet::file::properties::WriterPropertiesBuilder; +use parquet::schema::types::ColumnPath; +use store_api::metadata::RegionMetadataRef; +use store_api::mito_engine_options::FloatFieldEncoding; use store_api::storage::{ColumnId, FileId}; use crate::sst::DEFAULT_WRITE_BUFFER_SIZE; @@ -72,6 +77,27 @@ pub(crate) struct Json2TargetLayout { /// batching without changing the row group layout of newly written SSTs. pub const DEFAULT_ROW_GROUP_SIZE: usize = 100 * 1024; +/// Applies the configured encoding to direct floating-point field columns. +pub(crate) fn apply_float_field_encoding( + mut builder: WriterPropertiesBuilder, + metadata: &RegionMetadataRef, + encoding: FloatFieldEncoding, +) -> WriterPropertiesBuilder { + if encoding == FloatFieldEncoding::ByteStreamSplit { + for column in &metadata.column_metadatas { + if column.semantic_type == SemanticType::Field + && column.column_schema.data_type.is_float() + { + let path = ColumnPath::new(vec![column.column_schema.name.clone()]); + builder = builder + .set_column_encoding(path.clone(), parquet::basic::Encoding::BYTE_STREAM_SPLIT) + .set_column_dictionary_enabled(path, false); + } + } + } + builder +} + /// Parquet write options. #[derive(Debug, Clone)] pub struct WriteOptions { @@ -83,6 +109,8 @@ pub struct WriteOptions { /// Note: This is not a hard limit as we can only observe the file size when /// ArrowWrite writes to underlying writers. pub max_file_size: Option, + /// Encoding policy for direct floating-point field columns. + pub float_field_encoding: FloatFieldEncoding, } impl Default for WriteOptions { @@ -91,6 +119,7 @@ impl Default for WriteOptions { write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE, row_group_size: DEFAULT_ROW_GROUP_SIZE, max_file_size: None, + float_field_encoding: FloatFieldEncoding::default(), } } } @@ -125,6 +154,7 @@ mod tests { use std::sync::Arc; use api::v1::{OpType, SemanticType}; + use bytes::Bytes; use common_function::function::FunctionRef; use common_function::function_factory::ScalarFunctionFactory; use common_function::scalars::matches::MatchesFunction; @@ -135,20 +165,24 @@ mod tests { use datafusion_expr::{BinaryExpr, Expr, Literal, Operator, col, lit}; use datatypes::arrow; use datatypes::arrow::array::{ - ArrayRef, AsArray, BinaryDictionaryBuilder, RecordBatch, StringArray, - StringDictionaryBuilder, TimestampMillisecondArray, UInt8Array, UInt64Array, + Array, ArrayRef, AsArray, BinaryDictionaryBuilder, Float32Array, Float64Array, Int32Array, + RecordBatch, StringArray, StringDictionaryBuilder, TimestampMillisecondArray, UInt8Array, + UInt64Array, }; - use datatypes::arrow::datatypes::{DataType, Field, Schema, UInt32Type}; + use datatypes::arrow::datatypes::{DataType, Field, Schema, TimeUnit, UInt32Type}; use datatypes::arrow::util::pretty::pretty_format_batches; use datatypes::prelude::ConcreteDataType; use datatypes::schema::{FulltextAnalyzer, FulltextBackend, FulltextOptions}; use object_store::ObjectStore; - use parquet::arrow::AsyncArrowWriter; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use parquet::arrow::{ArrowWriter, AsyncArrowWriter}; use parquet::basic::{Compression, Encoding, ZstdLevel}; use parquet::file::metadata::{KeyValue, PageIndexPolicy}; use parquet::file::properties::WriterProperties; + use parquet::schema::types::ColumnPath; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder}; + use store_api::mito_engine_options::FloatFieldEncoding; use store_api::region_request::PathType; use store_api::storage::{ColumnSchema, RegionId}; use table::predicate::Predicate; @@ -187,6 +221,176 @@ mod tests { const FILE_DIR: &str = "/"; const REGION_ID: RegionId = RegionId::new(0, 0); + #[test] + fn test_float_field_encoding_properties_and_roundtrip() { + let mut metadata_builder = RegionMetadataBuilder::new(REGION_ID); + metadata_builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new("f32", ConcreteDataType::float32_datatype(), true), + semantic_type: SemanticType::Field, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new("f64", ConcreteDataType::float64_datatype(), true), + semantic_type: SemanticType::Field, + column_id: 1, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new("tag", ConcreteDataType::float32_datatype(), true), + semantic_type: SemanticType::Tag, + column_id: 2, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new("i32", ConcreteDataType::int32_datatype(), true), + semantic_type: SemanticType::Field, + column_id: 3, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 4, + }); + metadata_builder.primary_key(vec![2]); + let metadata = Arc::new(metadata_builder.build().unwrap()); + + let f32_values = [ + Some(1.5_f32), + Some(2.0), + Some(0.0), + Some(-0.0), + Some(f32::INFINITY), + Some(f32::from_bits(0x7fc0_1234)), + None, + ]; + let f64_values = [ + Some(1.5_f64), + Some(2.0), + Some(0.0), + Some(-0.0), + Some(f64::NEG_INFINITY), + Some(f64::from_bits(0x7ff8_0000_0000_1234)), + None, + ]; + let schema = Arc::new(Schema::new(vec![ + Field::new("f32", DataType::Float32, true), + Field::new("f64", DataType::Float64, true), + Field::new("tag", DataType::Float32, true), + Field::new("i32", DataType::Int32, true), + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Float32Array::from(f32_values.to_vec())) as ArrayRef, + Arc::new(Float64Array::from(f64_values.to_vec())) as ArrayRef, + Arc::new(Float32Array::from(vec![ + Some(1.0), + None, + Some(-0.0), + Some(2.0), + Some(3.0), + Some(4.0), + None, + ])), + Arc::new(Int32Array::from(vec![ + Some(1), + None, + Some(3), + Some(4), + Some(5), + Some(6), + None, + ])), + Arc::new(TimestampMillisecondArray::from_iter_values(0..7)), + ], + ) + .unwrap(); + + let path = |name: &str| ColumnPath::new(vec![name.to_string()]); + let bss = apply_float_field_encoding( + WriterProperties::builder(), + &metadata, + FloatFieldEncoding::ByteStreamSplit, + ); + assert_eq!( + Some(Encoding::BYTE_STREAM_SPLIT), + bss.clone().build().encoding(&path("f32")) + ); + assert_eq!( + Some(Encoding::BYTE_STREAM_SPLIT), + bss.clone().build().encoding(&path("f64")) + ); + assert!(!bss.clone().build().dictionary_enabled(&path("f32"))); + assert!(!bss.clone().build().dictionary_enabled(&path("f64"))); + assert_eq!(None, bss.clone().build().encoding(&path("i32"))); + assert!(bss.clone().build().dictionary_enabled(&path("tag"))); + + let default = apply_float_field_encoding( + WriterProperties::builder().set_encoding(Encoding::PLAIN), + &metadata, + FloatFieldEncoding::Default, + ) + .build(); + assert_eq!(Some(Encoding::PLAIN), default.encoding(&path("f32"))); + assert!(default.dictionary_enabled(&path("f32"))); + assert!(default.dictionary_enabled(&path("f64"))); + + let mut bytes = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut bytes, schema, Some(bss.build())).unwrap(); + writer.write(&batch).unwrap(); + let footer = writer.finish().unwrap(); + drop(writer); + for name in ["f32", "f64"] { + let column = footer.row_groups()[0] + .columns() + .iter() + .find(|column| column.column_path().string() == name) + .unwrap(); + assert!( + column + .encodings() + .any(|encoding| encoding == Encoding::BYTE_STREAM_SPLIT) + ); + } + let mut reader = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(bytes)) + .unwrap() + .build() + .unwrap(); + let actual = reader.next().unwrap().unwrap(); + let actual_f32 = actual + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let actual_f64 = actual + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + for (index, value) in f32_values.into_iter().enumerate() { + assert_eq!(value.is_none(), actual_f32.is_null(index)); + if let Some(value) = value { + assert_eq!(value.to_bits(), actual_f32.value(index).to_bits()); + } + } + for (index, value) in f64_values.into_iter().enumerate() { + assert_eq!(value.is_none(), actual_f64.is_null(index)); + if let Some(value) = value { + assert_eq!(value.to_bits(), actual_f64.value(index).to_bits()); + } + } + assert!(actual.column(2).is_null(1)); + assert!(actual.column(3).is_null(1)); + } + #[derive(Clone)] struct FixedPathProvider { region_file_id: RegionFileId, diff --git a/src/mito2/src/sst/parquet/flat_format.rs b/src/mito2/src/sst/parquet/flat_format.rs index 5432c112f49..b6dfdf0ec47 100644 --- a/src/mito2/src/sst/parquet/flat_format.rs +++ b/src/mito2/src/sst/parquet/flat_format.rs @@ -560,7 +560,8 @@ struct ParquetFlat { arrow_schema: SchemaRef, /// Projection computed for the flat format. format_projection: FormatProjection, - /// Column id to index in SST. + /// Column id to top-level SST index. Shared statistics helpers resolve + /// physical leaves from each file's actual Parquet schema. column_id_to_sst_index: HashMap, } @@ -618,7 +619,6 @@ impl ParquetFlat { // No such column in the SST. return StatValues::NoColumn; }; - let stats = column_null_counts(row_groups, *index); StatValues::from_stats_opt(stats) } @@ -633,8 +633,9 @@ impl ParquetFlat { // No such column in the SST. return StatValues::NoColumn; }; - // Safety: `column_id_to_sst_index` is built from `metadata`. - let index = self.column_id_to_sst_index.get(&column_id).unwrap(); + let Some(index) = self.column_id_to_sst_index.get(&column_id) else { + return StatValues::NoStats; + }; let stats = column_values(row_groups, column, *index, is_min); StatValues::from_stats_opt(stats) @@ -895,16 +896,25 @@ mod tests { use api::v1::SemanticType; use datatypes::arrow::array::{ - ArrayRef, BinaryArray, TimestampMillisecondArray, UInt8Array, UInt32Array, UInt64Array, + ArrayRef, BinaryArray, Int64Array, TimestampMillisecondArray, UInt8Array, UInt32Array, + UInt64Array, }; - use datatypes::arrow::datatypes::DataType as ArrowDataType; + use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, TimeUnit}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; + use datatypes::types::json_type::{JsonNativeType, JsonObjectType}; + use parquet::arrow::ArrowSchemaConverter; + use parquet::basic::{Repetition, Type as PhysicalType}; + use parquet::file::metadata::{ColumnChunkMetaData, RowGroupMetaData}; + use parquet::file::statistics::Statistics; + use parquet::schema::types::{SchemaDescriptor, Type}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder}; use store_api::storage::RegionId; - use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME; + use store_api::storage::consts::{ + OP_TYPE_COLUMN_NAME, PRIMARY_KEY_COLUMN_NAME, SEQUENCE_COLUMN_NAME, + }; use super::*; use crate::read::read_columns::ReadColumns; @@ -964,6 +974,331 @@ mod tests { builder.build().unwrap() } + /// Builds the metadata of a table with a JSON2 struct field column: + /// `[tag_0, field_0, payload, nullable_after_payload, ts]` with primary key `tag_0`. + fn metadata_with_struct_field() -> RegionMetadata { + let mut builder = RegionMetadataBuilder::new(RegionId::new(0, 0)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "tag_0".to_string(), + ConcreteDataType::string_datatype(), + true, + ), + semantic_type: SemanticType::Tag, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "field_0".to_string(), + ConcreteDataType::int64_datatype(), + true, + ), + semantic_type: SemanticType::Field, + column_id: 1, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "payload".to_string(), + ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())), + false, + ), + semantic_type: SemanticType::Field, + column_id: 2, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "nullable_after_payload".to_string(), + ConcreteDataType::int64_datatype(), + true, + ), + semantic_type: SemanticType::Field, + column_id: 4, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts".to_string(), + ConcreteDataType::timestamp_nanosecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 3, + }); + builder.primary_key(vec![0]); + builder.primary_key_encoding(PrimaryKeyEncoding::Dense); + builder.build().unwrap() + } + + /// Builds a file schema and a row group in which the `payload` struct + /// column expands to three leaf columns. The `ns_edge` leaf carries small + /// Int64 statistics that must not be mistaken for the statistics of `ts`. + fn struct_column_file_and_row_group(raw_pk_columns: bool) -> (SchemaRef, RowGroupMetaData) { + let mut arrow_fields = vec![ + Field::new("tag_0", ArrowDataType::Utf8, true), + Field::new("field_0", ArrowDataType::Int64, true), + Field::new( + "payload", + ArrowDataType::Struct( + vec![ + Field::new("metadata", ArrowDataType::Binary, false), + Field::new("value", ArrowDataType::Binary, false), + Field::new("ns_edge", ArrowDataType::Int64, true), + ] + .into(), + ), + false, + ), + Field::new("nullable_after_payload", ArrowDataType::Int64, true), + Field::new( + "ts", + ArrowDataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ), + Field::new(PRIMARY_KEY_COLUMN_NAME, ArrowDataType::Binary, false), + Field::new(SEQUENCE_COLUMN_NAME, ArrowDataType::UInt64, false), + Field::new(OP_TYPE_COLUMN_NAME, ArrowDataType::UInt8, false), + ]; + if !raw_pk_columns { + arrow_fields.remove(0); + } + let file_schema = Arc::new(Schema::new(arrow_fields)); + + let leaf = |name: &str, physical: PhysicalType| { + Arc::new( + Type::primitive_type_builder(name, physical) + .with_repetition(Repetition::OPTIONAL) + .build() + .unwrap(), + ) + }; + let payload = Arc::new( + Type::group_type_builder("payload") + .with_repetition(Repetition::OPTIONAL) + .with_fields(vec![ + leaf("metadata", PhysicalType::BYTE_ARRAY), + leaf("value", PhysicalType::BYTE_ARRAY), + leaf("ns_edge", PhysicalType::INT64), + ]) + .build() + .unwrap(), + ); + let mut parquet_fields = vec![ + leaf("tag_0", PhysicalType::BYTE_ARRAY), + leaf("field_0", PhysicalType::INT64), + payload, + leaf("nullable_after_payload", PhysicalType::INT64), + leaf("ts", PhysicalType::INT64), + leaf(PRIMARY_KEY_COLUMN_NAME, PhysicalType::BYTE_ARRAY), + leaf(SEQUENCE_COLUMN_NAME, PhysicalType::INT64), + leaf(OP_TYPE_COLUMN_NAME, PhysicalType::INT32), + ]; + if !raw_pk_columns { + parquet_fields.remove(0); + } + let schema_descr = Arc::new(SchemaDescriptor::new(Arc::new( + Type::group_type_builder("schema") + .with_fields(parquet_fields) + .build() + .unwrap(), + ))); + + // Omitted raw tags shift all subsequent leaves in primary-key SSTs. + let ns_edge_leaf = if raw_pk_columns { 4 } else { 3 }; + let nullable_leaf = ns_edge_leaf + 1; + let ts_leaf = nullable_leaf + 1; + let chunks: Vec<_> = (0..schema_descr.num_columns()) + .map(|i| { + let mut builder = ColumnChunkMetaData::builder(schema_descr.column(i)); + if i == ns_edge_leaf { + // Small values from the JSON payload, not timestamps. + builder = builder.set_statistics(Statistics::int64( + Some(0), + Some(86_400_000_000_000), + None, + Some(65), + true, + )); + } else if i == nullable_leaf { + builder = builder.set_statistics(Statistics::int64( + Some(100), + Some(200), + None, + Some(7), + true, + )); + } else if i == ts_leaf { + builder = builder.set_statistics(Statistics::int64( + Some(1_788_998_400_000_000_000), + Some(1_789_084_800_000_000_000), + None, + Some(0), + true, + )); + } + builder.build().unwrap() + }) + .collect(); + let row_group = RowGroupMetaData::builder(schema_descr) + .set_num_rows(69) + .set_total_byte_size(0) + .set_column_metadata(chunks) + .build() + .unwrap(); + + (file_schema, row_group) + } + + /// Regression test: row group statistics must be looked up by parquet leaf + /// column index. A struct field column (e.g. JSON2) expands to multiple + /// leaf columns, so statistics of columns after it must not be read from + /// the struct's leaves. Otherwise min-max pruning can drop a whole row + /// group by mistake (e.g. pruning `ts` with the small `ns_edge` stats), + /// which caused data loss during SWCS compaction. + #[test] + fn test_stats_with_struct_field_column() { + for (encoding, raw_pk_columns) in [ + (PrimaryKeyEncoding::Dense, true), + (PrimaryKeyEncoding::Dense, false), + (PrimaryKeyEncoding::Sparse, false), + ] { + let mut metadata = metadata_with_struct_field(); + metadata.primary_key_encoding = encoding; + let metadata = Arc::new(metadata); + let (file_schema, row_group) = struct_column_file_and_row_group(raw_pk_columns); + let read_format = FlatReadFormat::new( + metadata, + ReadColumns::new([0, 1, 2, 3, 4]), + Some(file_schema), + "test", + false, + ) + .unwrap(); + let row_groups = [&row_group]; + + // Statistics of `ts` come from the `ts` leaf column, not the leaves of + // the payload struct. + let StatValues::Values(min) = read_format.min_values(&row_groups, 3) else { + panic!("expected ts min values") + }; + let min = min.as_any().downcast_ref::().unwrap(); + assert_eq!(1_788_998_400_000_000_000, min.value(0)); + let StatValues::Values(max) = read_format.max_values(&row_groups, 3) else { + panic!("expected ts max values") + }; + let max = max.as_any().downcast_ref::().unwrap(); + assert_eq!(1_789_084_800_000_000_000, max.value(0)); + + let stats = crate::sst::parquet::stats::RowGroupPruningStats::new( + &row_groups, + &read_format, + None, + false, + ); + for (start, end, keep) in [ + (1_788_998_400_000_000_000, 1_789_084_800_000_000_000, true), + (1_789_084_800_000_000_001, 1_789_171_200_000_000_000, false), + ] { + let predicate = table::predicate::Predicate::new(vec![ + datafusion_expr::col("ts").gt_eq(datafusion_expr::lit( + datafusion_common::ScalarValue::TimestampNanosecond(Some(start), None), + )), + datafusion_expr::col("ts").lt(datafusion_expr::lit( + datafusion_common::ScalarValue::TimestampNanosecond(Some(end), None), + )), + ]); + assert_eq!( + vec![keep], + predicate + .prune_with_stats(&stats, read_format.metadata().schema.arrow_schema(),) + ); + } + + // Null counts of `ts` also read the correct leaf column. + let StatValues::Values(nulls) = read_format.null_counts(&row_groups, 3) else { + panic!("expected ts null counts") + }; + let nulls = nulls.as_any().downcast_ref::().unwrap(); + assert!(nulls.is_valid(0)); + assert_eq!(0, nulls.value(0)); + + // A null slot may contain an underlying zero. Check validity and + // a nonzero count to distinguish unknown or wrong-leaf statistics. + let StatValues::Values(nulls) = read_format.null_counts(&row_groups, 4) else { + panic!("expected nullable field null counts") + }; + let nulls = nulls.as_any().downcast_ref::().unwrap(); + assert!(nulls.is_valid(0)); + assert_eq!(7, nulls.value(0)); + + // A column that expands to multiple leaf columns has no single column + // statistics. + assert!(matches!( + read_format.min_values(&row_groups, 2), + StatValues::NoStats + )); + assert!(matches!( + read_format.max_values(&row_groups, 2), + StatValues::NoStats + )); + assert!(matches!( + read_format.null_counts(&row_groups, 2), + StatValues::NoStats + )); + } + } + + /// Even one leaf cannot supply the null count of a nested parent: + /// {"a": null} and [null] are non-null parents with null children. + #[test] + fn test_single_leaf_nested_columns_have_no_root_stats() { + let child = Arc::new(Field::new("a", ArrowDataType::Int64, true)); + for nested_type in [ + ArrowDataType::Struct(vec![child.clone()].into()), + ArrowDataType::List(child.clone()), + ArrowDataType::LargeList(child.clone()), + ArrowDataType::FixedSizeList(child, 1), + ] { + let metadata = Arc::new(metadata_with_struct_field()); + let (schema, _) = struct_column_file_and_row_group(true); + let mut fields = schema.fields().to_vec(); + fields[2] = Arc::new(Field::new("payload", nested_type.clone(), true)); + let schema = Arc::new(Schema::new(fields)); + let descriptor = Arc::new(ArrowSchemaConverter::new().convert(&schema).unwrap()); + let chunks = descriptor + .columns() + .iter() + .map(|column| { + ColumnChunkMetaData::builder(column.clone()) + .build() + .unwrap() + }) + .collect(); + let row_group = RowGroupMetaData::builder(descriptor) + .set_num_rows(1) + .set_total_byte_size(0) + .set_column_metadata(chunks) + .build() + .unwrap(); + let format = ParquetFlat::new(metadata, ReadColumns::new([0, 1, 2, 3]), schema); + + // The nested root has no usable statistics, independently of the + // values stored in any row group. Its leaf still occupies a slot. + let groups = &[row_group]; + assert!( + matches!(format.min_values(groups, 2), StatValues::NoStats), + "{nested_type:?}" + ); + assert!( + matches!(format.max_values(groups, 2), StatValues::NoStats), + "{nested_type:?}" + ); + assert!( + matches!(format.null_counts(groups, 2), StatValues::NoStats), + "{nested_type:?}" + ); + } + } + #[test] fn test_field_column_start() { // (num_tags, num_fields, encoding, expected) diff --git a/src/mito2/src/sst/parquet/format.rs b/src/mito2/src/sst/parquet/format.rs index 1a4325b7a7e..6ca7a486624 100644 --- a/src/mito2/src/sst/parquet/format.rs +++ b/src/mito2/src/sst/parquet/format.rs @@ -46,6 +46,7 @@ use mito_codec::row_converter::{ }; use parquet::file::metadata::{ParquetMetaData, RowGroupMetaData}; use parquet::file::statistics::Statistics; +use parquet::schema::types::SchemaDescriptor; use snafu::{OptionExt, ResultExt, ensure}; use store_api::metadata::{ColumnMetadata, RegionMetadataRef}; use store_api::storage::{ColumnId, NestedPath, SequenceNumber}; @@ -144,13 +145,16 @@ pub(crate) fn column_values( ) } -/// Returns min/max values of a parquet column with the given Arrow data type. +/// Returns min/max values for a top-level column with the given Arrow data type. +/// Resolves its leaf from the actual Parquet schema, not an inferred Arrow layout. pub(crate) fn column_values_by_type( row_groups: &[impl Borrow], data_type: &ArrowDataType, column_index: usize, is_min: bool, ) -> Option { + let column_index = + scalar_leaf_index(row_groups.first()?.borrow().schema_descr(), column_index)?; let null_scalar: ScalarValue = data_type.try_into().ok()?; let scalar_values = row_groups .iter() @@ -207,12 +211,14 @@ pub(crate) fn column_values_by_type( ScalarValue::iter_to_array(scalar_values).ok() } -/// Returns null counts of specific columns. +/// Returns null counts of a top-level column. /// The column should not be encoded as a part of a primary key. pub(crate) fn column_null_counts( row_groups: &[impl Borrow], column_index: usize, ) -> Option { + let column_index = + scalar_leaf_index(row_groups.first()?.borrow().schema_descr(), column_index)?; let values = row_groups.iter().map(|meta| { let col = meta.borrow().column(column_index); let stat = col.statistics()?; @@ -221,6 +227,21 @@ pub(crate) fn column_null_counts( Some(Arc::new(UInt64Array::from_iter(values))) } +/// Maps a scalar root to its physical leaf. Nested roots have no root-level +/// statistics, even with one leaf: child null counts do not describe parents. +/// All row groups of a file share the same schema. +fn scalar_leaf_index(schema: &SchemaDescriptor, root_index: usize) -> Option { + if !schema + .root_schema() + .get_fields() + .get(root_index)? + .is_primitive() + { + return None; + } + (0..schema.num_columns()).find(|&leaf| schema.get_column_root_idx(leaf) == root_index) +} + /// Helper for reading the SST format. pub struct PrimaryKeyReadFormat { /// The metadata stored in the SST. @@ -539,11 +560,12 @@ impl PrimaryKeyReadFormat { .into_iter(), ); + let primary_key_leaf = scalar_leaf_index( + row_groups.first()?.borrow().schema_descr(), + self.primary_key_position(), + )?; let values = row_groups.iter().map(|meta| { - let stats = meta - .borrow() - .column(self.primary_key_position()) - .statistics()?; + let stats = meta.borrow().column(primary_key_leaf).statistics()?; match stats { Statistics::Boolean(_) => None, Statistics::Int32(_) => None, diff --git a/src/mito2/src/sst/parquet/prefilter.rs b/src/mito2/src/sst/parquet/prefilter.rs index d39f540c186..4362440e415 100644 --- a/src/mito2/src/sst/parquet/prefilter.rs +++ b/src/mito2/src/sst/parquet/prefilter.rs @@ -24,15 +24,19 @@ use std::sync::Arc; use api::v1::SemanticType; use common_recordbatch::filter::SimpleFilterEvaluator; +use datafusion_common::ScalarValue; use datafusion_expr::Expr; use datatypes::arrow::array::{Array, BinaryArray, BooleanArray, BooleanBufferBuilder}; use datatypes::arrow::buffer::BooleanBuffer; use datatypes::arrow::datatypes::SchemaRef; use datatypes::arrow::record_batch::RecordBatch; +use datatypes::prelude::ConcreteDataType; +use datatypes::value::Value; use futures::StreamExt; use mito_codec::row_converter::{PrimaryKeyCodec, PrimaryKeyFilter, build_primary_key_codec}; use parquet::arrow::ProjectionMask; -use parquet::arrow::arrow_reader::RowSelection; +use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; +use parquet::file::metadata::ParquetMetaData; use parquet::schema::types::SchemaDescriptor; use smallvec::{SmallVec, smallvec}; use snafu::{OptionExt, ResultExt}; @@ -47,7 +51,7 @@ use crate::error::{ }; use crate::sst::parquet::file_range::PreFilterMode; use crate::sst::parquet::flat_format::FlatReadFormat; -use crate::sst::parquet::format::PrimaryKeyArray; +use crate::sst::parquet::format::{PrimaryKeyArray, StatValues}; use crate::sst::parquet::reader::{ MaybeFilter, PhysicalFilterContext, RowGroupBuildContext, RowGroupReaderBuilder, SimpleFilterContext, @@ -398,6 +402,7 @@ pub(crate) fn build_bulk_filter_plan( /// the prefilter pass. A caller can postpone simple timestamp filters to the normal /// precise-filter path when the scan time range covers the SST. When predicate /// prefiltering is disabled, all simple filters remain on the normal path instead. +#[allow(clippy::too_many_arguments)] pub(crate) fn build_reader_filter_plan( predicate: Option<&Predicate>, expected_metadata: Option<&RegionMetadata>, @@ -406,6 +411,7 @@ pub(crate) fn build_reader_filter_plan( postpone_time_index_filter: bool, read_format: &FlatReadFormat, codec: &Arc, + parquet_metadata: &ParquetMetaData, ) -> ReaderFilterPlan { let Some(predicate) = predicate else { return ReaderFilterPlan { @@ -533,6 +539,7 @@ pub(crate) fn build_reader_filter_plan( prefilter_simple_filters.clone(), prefilter_physical_filters, schema_version, + parquet_metadata, ); if prefilter_builder.is_some() { @@ -567,6 +574,8 @@ pub(crate) struct PrefilterContext { pk_filter_expr_strs: Option>, /// Arrow schema used to build narrowed prefilter projections. arrow_schema: SchemaRef, + /// Simple filters already proven SQL-true by this row group's statistics. + proven_simple_filters: Vec, } /// Pre-built state for constructing [PrefilterContext] per row group. @@ -583,6 +592,8 @@ pub(crate) struct PrefilterContextBuilder { metadata: RegionMetadataRef, schema_version: u64, arrow_schema: SchemaRef, + /// Per-row-group simple filters already proven SQL-true by column statistics. + proven_simple_filters: Vec>, } impl PrefilterContextBuilder { @@ -592,6 +603,7 @@ impl PrefilterContextBuilder { /// - The read format doesn't use flat layout /// - No prefilter columns are selected /// - Prefilter would read the full projection without any PK filter + #[allow(clippy::too_many_arguments)] pub(crate) fn new( read_format: &FlatReadFormat, codec: &Arc, @@ -600,6 +612,7 @@ impl PrefilterContextBuilder { filters: Vec, physical_filters: Vec, schema_version: u64, + parquet_metadata: &ParquetMetaData, ) -> Option { let metadata = read_format.metadata(); let use_raw_tag_columns = read_format.batch_has_raw_pk_columns(); @@ -646,6 +659,9 @@ impl PrefilterContextBuilder { return None; } + let proven_simple_filters = + simple_filter_stats_proofs(read_format, parquet_metadata.row_groups(), &filters); + Some(Self { pk_filters, pk_filter_expr_strs, @@ -655,11 +671,12 @@ impl PrefilterContextBuilder { metadata: metadata.clone(), schema_version, arrow_schema: read_format.arrow_schema().clone(), + proven_simple_filters, }) } /// Builds a [PrefilterContext] for a specific row group. - pub(crate) fn build(&self) -> PrefilterContext { + pub(crate) fn build(&self, row_group_idx: usize) -> PrefilterContext { let pk_filter = self .build_primary_key_filter() .map(|filter| Box::new(filter) as Box); @@ -670,6 +687,11 @@ impl PrefilterContextBuilder { schema_version: self.schema_version, pk_filter_expr_strs: self.pk_filter_expr_strs.clone(), arrow_schema: self.arrow_schema.clone(), + proven_simple_filters: self + .proven_simple_filters + .get(row_group_idx) + .cloned() + .unwrap_or_else(|| vec![false; self.filters.len()]), } } @@ -687,6 +709,114 @@ impl PrefilterContextBuilder { const PREFILTER_COLUMN_RATIO_THRESHOLD: f64 = 0.5; const PREFILTER_MIN_REMAINING_COLUMNS: usize = 2; +/// Returns row-group-major proof bits for simple filters. Statistics are +/// extracted once per eligible filter across all row groups. +fn simple_filter_stats_proofs( + read_format: &FlatReadFormat, + row_groups: &[parquet::file::metadata::RowGroupMetaData], + filters: &[SimpleFilterContext], +) -> Vec> { + let mut proofs = vec![vec![false; filters.len()]; row_groups.len()]; + for (filter_idx, filter_ctx) in filters.iter().enumerate() { + let Some((filter, literal)) = eligible_simple_filter(read_format, filter_ctx) else { + continue; + }; + let (StatValues::Values(mins), StatValues::Values(maxs), StatValues::Values(null_counts)) = ( + read_format.min_values(row_groups, filter_ctx.column_id()), + read_format.max_values(row_groups, filter_ctx.column_id()), + read_format.null_counts(row_groups, filter_ctx.column_id()), + ) else { + continue; + }; + for (row_group_idx, proof) in proofs.iter_mut().enumerate() { + proof[filter_idx] = simple_filter_is_true_by_values( + filter, + &literal, + stat_value_at(&mins, row_group_idx), + stat_value_at(&maxs, row_group_idx), + stat_value_at(&null_counts, row_group_idx), + ); + } + } + proofs +} + +fn eligible_simple_filter<'a>( + read_format: &FlatReadFormat, + filter_ctx: &'a SimpleFilterContext, +) -> Option<(&'a SimpleFilterEvaluator, Value)> { + if filter_ctx.semantic_type() != SemanticType::Field { + return None; + } + let filter = filter_ctx.filter().as_filter()?; + let literal = filter.literal_value()?; + let column = read_format + .metadata() + .column_by_id(filter_ctx.column_id())?; + column_type_matches_literal(&column.column_schema.data_type, &literal) + .then_some((filter, literal)) +} + +fn simple_filter_is_true_by_values( + filter: &SimpleFilterEvaluator, + literal: &Value, + min: Option, + max: Option, + null_count: Option, +) -> bool { + let (Some(min), Some(max), Some(null_count)) = (min, max, null_count) else { + return false; + }; + if null_count != Value::UInt64(0) + || !same_supported_value_type(&min, literal) + || !same_supported_value_type(&max, literal) + || min > max + { + return false; + } + + if filter.is_gt() { + min > *literal + } else if filter.is_gt_eq() { + min >= *literal + } else if filter.is_lt() { + max < *literal + } else if filter.is_lt_eq() { + max <= *literal + } else if filter.is_eq() { + min == *literal && max == *literal + } else if filter.is_not_eq() { + max < *literal || min > *literal + } else { + false + } +} + +fn column_type_matches_literal(data_type: &ConcreteDataType, literal: &Value) -> bool { + matches!( + (data_type, literal), + (ConcreteDataType::Int32(_), Value::Int32(_)) + | (ConcreteDataType::UInt32(_), Value::UInt32(_)) + | (ConcreteDataType::Int64(_), Value::Int64(_)) + | (ConcreteDataType::UInt64(_), Value::UInt64(_)) + ) +} + +fn same_supported_value_type(left: &Value, right: &Value) -> bool { + matches!( + (left, right), + (Value::Int32(_), Value::Int32(_)) + | (Value::UInt32(_), Value::UInt32(_)) + | (Value::Int64(_), Value::Int64(_)) + | (Value::UInt64(_), Value::UInt64(_)) + ) +} + +fn stat_value_at(values: &datatypes::arrow::array::ArrayRef, index: usize) -> Option { + let scalar = ScalarValue::try_from_array(values, index).ok()?; + Value::try_from(scalar).ok() +} + /// Result of prefiltering a row group. pub(crate) struct PrefilterResult { /// Refined row selection after prefiltering. @@ -941,6 +1071,9 @@ async fn execute_prefilter_by_reading_columns( build_ctx: &RowGroupBuildContext<'_>, ) -> Result { let entries = all_prefilter_entries(prefilter_ctx); + if entries.is_empty() { + return Ok(identity_prefilter_result(reader_builder, build_ctx)); + } let (mask, rows_before_filter) = build_prefilter_masks(prefilter_ctx, reader_builder, build_ctx, &entries).await?; @@ -965,6 +1098,13 @@ fn all_prefilter_entries(prefilter_ctx: &PrefilterContext) -> Vec, +) -> PrefilterResult { + let row_count = reader_builder + .parquet_metadata() + .row_group(build_ctx.row_group_idx) + .num_rows() as usize; + PrefilterResult { + refined_selection: identity_row_selection(&build_ctx.row_selection, row_count), + filtered_rows: 0, + } +} + +fn identity_row_selection( + original_selection: &Option, + row_count: usize, +) -> RowSelection { + original_selection + .clone() + .unwrap_or_else(|| RowSelection::from(vec![RowSelector::select(row_count)])) +} + fn rows_before_filter( reader_builder: &RowGroupReaderBuilder, build_ctx: &RowGroupBuildContext<'_>, @@ -1200,15 +1371,24 @@ mod tests { use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; + use bytes::Bytes; use common_recordbatch::filter::SimpleFilterEvaluator; use datafusion_common::ScalarValue; use datafusion_expr::{col, lit}; use datatypes::arrow::array::{ - ArrayRef, DictionaryArray, TimestampMillisecondArray, UInt8Array, UInt32Array, UInt64Array, + ArrayRef, DictionaryArray, Int32Array, TimestampMillisecondArray, UInt8Array, UInt32Array, + UInt64Array, }; use datatypes::arrow::datatypes::{DataType, Field, Schema, UInt32Type}; + use datatypes::arrow::record_batch::RecordBatch; + use datatypes::prelude::ConcreteDataType; + use datatypes::value::Value; use mito_codec::row_converter::{PrimaryKeyFilter, build_primary_key_codec}; + use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::{ParquetRecordBatchReaderBuilder, RowSelector}; use store_api::codec::PrimaryKeyEncoding; + use store_api::metadata::RegionMetadataBuilder; + use store_api::region_request::{AlterKind, ModifyColumnType}; use super::*; use crate::read::read_columns::ReadColumns; @@ -1268,6 +1448,83 @@ mod tests { .collect() } + fn stats_metadata_with_options( + values: &[Vec], + nulls: Option<&[bool]>, + writer_options: Option, + ) -> Arc { + let first = new_record_batch_with_custom_sequence(&["a", "x"], 0, values[0].len(), 1); + let mut bytes = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut bytes, first.schema(), writer_options).unwrap(); + for (idx, values) in values.iter().enumerate() { + let has_null = nulls.and_then(|nulls| nulls.get(idx)).copied() == Some(true); + let batch = new_record_batch_with_custom_sequence( + &["a", "x"], + 0, + if has_null { + values.len().max(2) + } else { + values.len() + }, + 1, + ); + let mut columns = batch.columns().to_vec(); + columns[2] = Arc::new(if has_null { + UInt64Array::from(vec![Some(values[0]), None]) + } else { + UInt64Array::from_iter_values(values.iter().copied()) + }); + writer + .write(&RecordBatch::try_new(batch.schema(), columns).unwrap()) + .unwrap(); + writer.flush().unwrap(); + } + writer.close().unwrap(); + ParquetRecordBatchReaderBuilder::try_new(Bytes::from(bytes)) + .unwrap() + .metadata() + .clone() + } + + fn stats_metadata(values: &[Vec], nulls: Option<&[bool]>) -> Arc { + stats_metadata_with_options(values, nulls, None) + } + + fn int32_stats_metadata(values: &[i32]) -> Arc { + let batch = new_record_batch_with_custom_sequence(&["a", "x"], 0, values.len(), 1); + let mut fields = batch + .schema() + .fields() + .iter() + .map(|field| field.as_ref().clone()) + .collect::>(); + fields[2].set_data_type(DataType::Int32); + let mut columns = batch.columns().to_vec(); + columns[2] = Arc::new(Int32Array::from(values.to_vec())); + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(); + let mut bytes = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut bytes, batch.schema(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + ParquetRecordBatchReaderBuilder::try_new(Bytes::from(bytes)) + .unwrap() + .metadata() + .clone() + } + + fn metadata_with_field_type(data_type: ConcreteDataType) -> RegionMetadataRef { + let mut builder = RegionMetadataBuilder::from_existing(sst_region_metadata()); + builder + .alter(AlterKind::ModifyColumnTypes { + columns: vec![ModifyColumnType { + column_name: "field_0".to_string(), + target_type: data_type, + }], + }) + .unwrap(); + Arc::new(builder.build().unwrap()) + } + fn new_physical_filter_contexts( metadata: &RegionMetadataRef, read_format: &FlatReadFormat, @@ -1502,10 +1759,292 @@ mod tests { Vec::new(), Vec::new(), metadata.schema_version, + &stats_metadata(&[vec![1]], None), ); assert!(builder.is_none()); } + #[test] + fn test_simple_filter_stats_uses_real_metadata() { + let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); + let read_format = FlatReadFormat::new( + metadata.clone(), + ReadColumns::new( + metadata + .column_metadatas + .iter() + .map(|column| column.column_id), + ), + None, + "test", + true, + ) + .unwrap(); + let parquet_metadata = stats_metadata(&[vec![i64::MAX as u64 + 2]], None); + let filters = new_simple_filter_contexts( + &metadata, + &[ + col("field_0").gt(lit(i64::MAX as u64 + 1)), + col("field_0").gt(lit(i64::MAX as u64 + 2)), + lit(i64::MAX as u64 + 1).lt(col("field_0")), + ], + ); + + let proofs = + simple_filter_stats_proofs(&read_format, parquet_metadata.row_groups(), &filters); + assert_eq!(proofs, vec![vec![true, false, true]]); + } + + #[test] + fn test_simple_filter_stats_retain_narrow_or_incomplete_metadata() { + let parquet_metadata = stats_metadata(&[vec![2]], Some(&[true])); + let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); + let read_format = FlatReadFormat::new( + metadata.clone(), + ReadColumns::new( + metadata + .column_metadatas + .iter() + .map(|column| column.column_id), + ), + None, + "test", + true, + ) + .unwrap(); + let filter = new_simple_filter_contexts(&metadata, &[col("field_0").gt(lit(1_u64))]); + assert_eq!( + simple_filter_stats_proofs(&read_format, parquet_metadata.row_groups(), &filter), + vec![vec![false]], + ); + let missing_stats = stats_metadata_with_options( + &[vec![2]], + None, + Some( + parquet::file::properties::WriterProperties::builder() + .set_statistics_enabled(parquet::file::properties::EnabledStatistics::None) + .build(), + ), + ); + assert_eq!( + simple_filter_stats_proofs(&read_format, missing_stats.row_groups(), &filter), + vec![vec![false]], + ); + + // Parquet INT32 stats decode as Value::Int32; without the actual SST + // type gate this Int8 field would be incorrectly proven true. + let narrow_metadata = metadata_with_field_type(ConcreteDataType::int8_datatype()); + let narrow_read_format = FlatReadFormat::new( + narrow_metadata.clone(), + ReadColumns::new( + narrow_metadata + .column_metadatas + .iter() + .map(|column| column.column_id), + ), + None, + "test", + true, + ) + .unwrap(); + let narrow_filter = + new_simple_filter_contexts(&narrow_metadata, &[col("field_0").gt(lit(1_i32))]); + let narrow_stats = int32_stats_metadata(&[2]); + assert_eq!( + simple_filter_stats_proofs( + &narrow_read_format, + narrow_stats.row_groups(), + &narrow_filter, + ), + vec![vec![false]], + ); + } + + #[test] + fn test_prefilter_builder_uses_row_group_proofs_by_index() { + let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); + let read_format = FlatReadFormat::new( + metadata.clone(), + ReadColumns::new( + metadata + .column_metadatas + .iter() + .map(|column| column.column_id), + ), + None, + "test", + true, + ) + .unwrap(); + let codec = build_primary_key_codec(metadata.as_ref()); + let builder = PrefilterContextBuilder::new( + &read_format, + &codec, + None, + None, + new_simple_filter_contexts(&metadata, &[col("field_0").gt(lit(1_u64))]), + Vec::new(), + metadata.schema_version, + &stats_metadata(&[vec![2], vec![1]], None), + ) + .unwrap(); + + assert!(builder.build(0).proven_simple_filters[0]); + assert!(!builder.build(1).proven_simple_filters[0]); + } + + #[test] + fn test_simple_filter_stats_prove_integer_predicates() { + macro_rules! assert_all_operators { + ($literal:expr, $base:expr, $below:expr, $above:expr) => {{ + for (expr, min, max) in [ + (col("x").gt(lit($literal)), $above.clone(), $above.clone()), + (col("x").gt_eq(lit($literal)), $base.clone(), $above.clone()), + (col("x").lt(lit($literal)), $below.clone(), $below.clone()), + (col("x").lt_eq(lit($literal)), $below.clone(), $base.clone()), + (col("x").eq(lit($literal)), $base.clone(), $base.clone()), + ( + col("x").not_eq(lit($literal)), + $below.clone(), + $below.clone(), + ), + ] { + let filter = SimpleFilterEvaluator::try_new(&expr).unwrap(); + assert!(simple_filter_is_true_by_values( + &filter, + &$base, + Some(min), + Some(max), + Some(Value::UInt64(0)), + )); + } + }}; + } + + assert_all_operators!(5_i32, Value::Int32(5), Value::Int32(-1), Value::Int32(6)); + assert_all_operators!(5_u32, Value::UInt32(5), Value::UInt32(4), Value::UInt32(6)); + assert_all_operators!(5_i64, Value::Int64(5), Value::Int64(-1), Value::Int64(6)); + assert_all_operators!( + i64::MAX as u64 + 2, + Value::UInt64(i64::MAX as u64 + 2), + Value::UInt64(i64::MAX as u64 + 1), + Value::UInt64(i64::MAX as u64 + 3) + ); + + // Boundary cases must remain unproven: the min/max interval still + // admits a row that violates the predicate. + for (expr, min, max) in [ + (col("x").lt(lit(5_i64)), 1, 5), + (col("x").gt(lit(5_i64)), 5, 9), + (col("x").lt_eq(lit(5_i64)), 1, 6), + (col("x").gt_eq(lit(5_i64)), 4, 9), + (col("x").eq(lit(5_i64)), 5, 6), + (col("x").not_eq(lit(5_i64)), 5, 9), + (col("x").not_eq(lit(5_i64)), 1, 5), + (col("x").not_eq(lit(5_i64)), 1, 9), + ] { + let filter = SimpleFilterEvaluator::try_new(&expr).unwrap(); + assert!( + !simple_filter_is_true_by_values( + &filter, + &Value::Int64(5), + Some(Value::Int64(min)), + Some(Value::Int64(max)), + Some(Value::UInt64(0)), + ), + "{expr:?} must not be proven by stats {min}..={max}", + ); + } + } + + #[test] + fn test_simple_filter_stats_retain_unknown_or_unsupported_values() { + let filter = SimpleFilterEvaluator::try_new(&col("x").gt(lit(1_i32))).unwrap(); + let literal = Value::Int32(1); + for (min, max, null_count) in [ + ( + Some(Value::Int32(2)), + Some(Value::Int32(3)), + Some(Value::UInt64(1)), + ), + (Some(Value::Int32(2)), Some(Value::Int32(3)), None), + ( + Some(Value::Null), + Some(Value::Int32(3)), + Some(Value::UInt64(0)), + ), + (None, Some(Value::Int32(3)), Some(Value::UInt64(0))), + (Some(Value::Int32(2)), None, Some(Value::UInt64(0))), + ( + Some(Value::Int64(2)), + Some(Value::Int64(3)), + Some(Value::UInt64(0)), + ), + ( + Some(Value::Int32(3)), + Some(Value::Int32(2)), + Some(Value::UInt64(0)), + ), + ] { + assert!(!simple_filter_is_true_by_values( + &filter, &literal, min, max, null_count + )); + } + + let float_filter = SimpleFilterEvaluator::try_new(&col("x").gt(lit(1.0_f64))).unwrap(); + assert!(!simple_filter_is_true_by_values( + &float_filter, + &Value::Float64(1.0.into()), + Some(Value::Float64(2.0.into())), + Some(Value::Float64(3.0.into())), + Some(Value::UInt64(0)), + )); + } + + #[test] + fn test_prefilter_entries_keep_only_unproven_simple_filter_indices() { + let metadata: RegionMetadataRef = Arc::new(sst_region_metadata()); + let filters = new_simple_filter_contexts( + &metadata, + &[col("field_0").gt(lit(1_u64)), col("field_0").lt(lit(9_u64))], + ); + let context = PrefilterContext { + pk_filter: None, + filters, + physical_filters: Vec::new(), + schema_version: metadata.schema_version, + pk_filter_expr_strs: None, + arrow_schema: metadata.schema.arrow_schema().clone(), + proven_simple_filters: vec![true, false], + }; + + let entries = all_prefilter_entries(&context); + assert!(matches!( + entries.as_slice(), + [PrefilterEntry { + kind: PrefilterEntryKind::Simple(1), + .. + }] + )); + } + + #[test] + fn test_identity_row_selection_preserves_input() { + let sparse = RowSelection::from(vec![ + RowSelector::skip(2), + RowSelector::select(3), + RowSelector::skip(1), + ]); + assert_eq!(identity_row_selection(&Some(sparse.clone()), 6), sparse); + + let empty = RowSelection::from(vec![]); + assert_eq!(identity_row_selection(&Some(empty.clone()), 6), empty); + assert_eq!( + identity_row_selection(&None, 6), + RowSelection::from(vec![RowSelector::select(6)]) + ); + } + #[test] fn test_should_use_prefilter() { assert!(should_use_prefilter(1, 5, 6)); @@ -1602,6 +2141,7 @@ mod tests { false, &full_read_format, &codec, + &stats_metadata(&[vec![1]], None), ); assert!(skip_fields_plan.prefilter_builder.is_some()); assert_eq!( @@ -1621,6 +2161,7 @@ mod tests { true, &full_read_format, &codec, + &stats_metadata(&[vec![1]], None), ); assert!(postponed_time_plan.prefilter_builder.is_some()); assert_eq!( @@ -1638,6 +2179,7 @@ mod tests { true, &full_read_format, &codec, + &stats_metadata(&[vec![1]], None), ); assert!(postponed_time_only_plan.prefilter_builder.is_none()); assert_eq!( @@ -1667,6 +2209,7 @@ mod tests { false, &projected_read_format, &metric_codec, + &stats_metadata(&[vec![1]], None), ); assert!(pk_prefilter_plan.prefilter_builder.is_some()); assert!( @@ -1691,6 +2234,7 @@ mod tests { true, &projected_read_format, &metric_codec, + &stats_metadata(&[vec![1]], None), ); assert!(disabled_plan.prefilter_builder.is_none()); assert_eq!( @@ -1724,6 +2268,7 @@ mod tests { false, &read_format, &codec, + &stats_metadata(&[vec![1]], None), ); let plan_b_a = build_reader_filter_plan( Some(&Predicate::new(vec![expr_b, expr_a])), @@ -1733,6 +2278,7 @@ mod tests { false, &read_format, &codec, + &stats_metadata(&[vec![1]], None), ); let exprs_ab = plan_ab.prefilter_builder.unwrap().pk_filter_expr_strs; diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index ff58d920b1f..b8588bd36cd 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -598,6 +598,7 @@ impl ParquetReaderBuilder { self.postpone_time_index_filter, &read_format, &codec, + &parquet_meta, ); if self.defer_optional_page_index @@ -1975,7 +1976,10 @@ impl RowGroupReaderBuilder { &self, build_ctx: RowGroupBuildContext<'_>, ) -> Result { - let prefilter_ctx = self.prefilter_builder.as_ref().map(|b| b.build()); + let prefilter_ctx = self + .prefilter_builder + .as_ref() + .map(|b| b.build(build_ctx.row_group_idx)); let Some(mut prefilter_ctx) = prefilter_ctx else { // No prefilter applicable, build stream with full projection. @@ -2553,9 +2557,11 @@ impl FlatRowGroupReader { #[cfg(test)] mod tests { use std::any::Any; + use std::collections::HashMap; use std::fmt::{Debug, Formatter}; use std::sync::{Arc, LazyLock}; + use api::v1::OpType; use common_error::ext::WhateverResult; use common_function::scalars::json::json_get::JsonGetWithType; use common_function::scalars::udf::create_udf; @@ -2567,25 +2573,271 @@ mod tests { ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, col, lit, }; - use datatypes::arrow::array::{ArrayRef, Int64Array, StringArray, StructArray}; - use datatypes::arrow::datatypes::{Fields, Schema}; + use datatypes::arrow::array::{ + ArrayRef, BinaryDictionaryBuilder, Int64Array, StringArray, StructArray, + TimestampMillisecondArray, UInt8Array, UInt64Array, + }; + use datatypes::arrow::datatypes::{Fields, Schema, UInt32Type}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::extension::json::Json2ExtensionType; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; use object_store::services::Memory; use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::RowSelector; use parquet::file::properties::WriterProperties; + use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder}; use store_api::region_request::PathType; use store_api::storage::RegionId; use table::predicate::Predicate; use super::*; + use crate::cache::CacheManager; use crate::sst::parquet::metadata::MetadataLoader; + use crate::sst::parquet::prefilter::{build_reader_filter_plan, execute_prefilter}; use crate::sst::parquet::read_columns::{ParquetReadColumn, ParquetReadColumns}; + use crate::sst::parquet::row_group::ParquetFetchMetrics; use crate::test_util::sst_util::{sst_file_handle, sst_region_metadata}; + async fn prefilter_test_builder( + object_store: ObjectStore, + predicate: Predicate, + cache_strategy: CacheStrategy, + ) -> (RowGroupReaderBuilder, Arc) { + let metadata = Arc::new( + crate::test_util::sst_util::sst_region_metadata_with_encoding( + PrimaryKeyEncoding::Sparse, + ), + ); + let batch = |start: i64, end: i64| { + let mut primary_key = BinaryDictionaryBuilder::::new(); + let mut fields = Vec::new(); + let mut timestamps = Vec::new(); + for value in start..end { + let tag = if value == 4 { "b" } else { "a" }; + primary_key + .append(crate::test_util::sst_util::new_sparse_primary_key( + &[tag, "x"], + &metadata, + 1, + 100, + )) + .unwrap(); + fields.push(value as u64); + timestamps.push(value); + } + RecordBatch::try_new( + crate::sst::to_flat_sst_arrow_schema( + &metadata, + &crate::sst::FlatSchemaOptions::from_encoding(PrimaryKeyEncoding::Sparse), + ), + vec![ + Arc::new(UInt64Array::from(fields)) as ArrayRef, + Arc::new(TimestampMillisecondArray::from(timestamps)) as ArrayRef, + Arc::new(primary_key.finish()) as ArrayRef, + Arc::new(UInt64Array::from_value(1, (end - start) as usize)) as ArrayRef, + Arc::new(UInt8Array::from_value( + OpType::Put as u8, + (end - start) as usize, + )) as ArrayRef, + ], + ) + .unwrap() + }; + let first_batch = batch(0, 3); + let second_batch = batch(3, 6); + let mut bytes = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut bytes, first_batch.schema(), None).unwrap(); + writer.write(&first_batch).unwrap(); + writer.flush().unwrap(); + writer.write(&second_batch).unwrap(); + writer.close().unwrap(); + + let file_handle = sst_file_handle(0, 6); + let file_path = file_handle.file_path("prefilter_test", PathType::Bare); + let file_size = bytes.len() as u64; + object_store.write(&file_path, bytes).await.unwrap(); + + let mut cache_metrics = MetadataCacheMetrics::default(); + let parquet_meta = Arc::new( + MetadataLoader::new(object_store.clone(), &file_path, file_size) + .load(&mut cache_metrics) + .await + .unwrap(), + ); + let read_format = FlatReadFormat::new( + metadata.clone(), + ReadColumns::new( + metadata + .column_metadatas + .iter() + .map(|column| column.column_id), + ), + None, + &file_path, + false, + ) + .unwrap(); + let codec = build_primary_key_codec(metadata.as_ref()); + let filter_plan = build_reader_filter_plan( + Some(&predicate), + None, + PreFilterMode::All, + true, + false, + &read_format, + &codec, + &parquet_meta, + ); + assert!(filter_plan.prefilter_builder.is_some()); + + let output_schema = read_format.arrow_schema().clone(); + let parquet_schema = parquet_meta.file_metadata().schema_descr(); + let projection = build_projection_plan(read_format.parquet_read_columns(), parquet_schema); + let arrow_metadata = + ArrowReaderMetadata::try_new(parquet_meta.clone(), ArrowReaderOptions::new()).unwrap(); + ( + RowGroupReaderBuilder { + file_handle: file_handle.clone(), + file_path, + parquet_meta, + parquet_metadata_size: 0, + arrow_metadata, + output_schema, + json2_rewrite_targets: HashMap::new(), + object_store, + projection, + has_nested_projection: false, + cache_strategy, + prefilter_builder: filter_plan.prefilter_builder, + batch_size: DEFAULT_READ_BATCH_SIZE, + }, + metadata, + ) + } + + #[tokio::test(flavor = "current_thread")] + async fn test_execute_prefilter_proven_filters_preserve_selection_without_fetching() { + let object_store = ObjectStore::new(Memory::default()).unwrap(); + let predicate = Predicate::new(vec![col("field_0").gt_eq(lit(0_u64))]); + let (reader_builder, _) = + prefilter_test_builder(object_store, predicate, CacheStrategy::Disabled).await; + let prefilter_builder = reader_builder.prefilter_builder.as_ref().unwrap(); + + for original_selection in [ + None, + Some(RowSelection::from(vec![ + RowSelector::skip(1), + RowSelector::select(1), + RowSelector::skip(1), + ])), + Some(RowSelection::from(vec![])), + ] { + let mut prefilter_ctx = prefilter_builder.build(0); + let fetch_metrics = ParquetFetchMetrics::default(); + let result = execute_prefilter( + &mut prefilter_ctx, + &reader_builder, + &RowGroupBuildContext { + row_group_idx: 0, + row_selection: original_selection.clone(), + fetch_metrics: Some(&fetch_metrics), + }, + ) + .await + .unwrap(); + + let expected = original_selection.unwrap_or_else(|| { + RowSelection::from(vec![RowSelector::select( + reader_builder.parquet_meta.row_group(0).num_rows() as usize, + )]) + }); + assert_eq!(result.refined_selection, expected); + assert_eq!(result.filtered_rows, 0); + let metrics = fetch_metrics.data.lock().unwrap(); + assert_eq!(metrics.pages_to_fetch_store, 0); + assert_eq!(metrics.pages_to_fetch_mem, 0); + assert_eq!(metrics.pages_to_fetch_write_cache, 0); + } + } + + #[tokio::test(flavor = "current_thread")] + async fn test_execute_prefilter_mixed_filters_use_nonzero_row_group_and_cache() { + let object_store = ObjectStore::new(Memory::default()).unwrap(); + let predicate = Predicate::new(vec![ + col("field_0").gt_eq(lit(3_u64)), + col("ts").lt(lit(ScalarValue::TimestampMillisecond(Some(6), None))), + col("field_0").in_list(vec![lit(3_u64), lit(4_u64)], false), + col("tag_0").eq(lit("a")), + ]); + let cache = CacheStrategy::EnableAll(Arc::new( + CacheManager::builder() + .prefilter_result_cache_size(1024) + .build(), + )); + let (reader_builder, _) = + prefilter_test_builder(object_store, predicate, cache.clone()).await; + let prefilter_builder = reader_builder.prefilter_builder.as_ref().unwrap(); + + for pass in 0..2 { + let mut prefilter_ctx = prefilter_builder.build(1); + let fetch_metrics = ParquetFetchMetrics::default(); + let fetch_metrics_ref = (pass == 1).then_some(&fetch_metrics); + let result = execute_prefilter( + &mut prefilter_ctx, + &reader_builder, + &RowGroupBuildContext { + row_group_idx: 1, + row_selection: Some(RowSelection::from(vec![ + RowSelector::select(2), + RowSelector::skip(1), + ])), + fetch_metrics: fetch_metrics_ref, + }, + ) + .await + .unwrap(); + assert_eq!(result.filtered_rows, 1); + assert_eq!( + result.refined_selection, + RowSelection::from(vec![RowSelector::select(1), RowSelector::skip(2),]) + ); + if pass == 1 { + assert_eq!(fetch_metrics.data.lock().unwrap().pages_to_fetch_store, 0); + } + } + + let disabled = prefilter_test_builder( + ObjectStore::new(Memory::default()).unwrap(), + Predicate::new(vec![ + col("field_0").gt_eq(lit(3_u64)), + col("ts").lt(lit(ScalarValue::TimestampMillisecond(Some(6), None))), + col("field_0").in_list(vec![lit(3_u64), lit(4_u64)], false), + col("tag_0").eq(lit("a")), + ]), + CacheStrategy::Disabled, + ) + .await; + let mut prefilter_ctx = disabled.0.prefilter_builder.as_ref().unwrap().build(1); + let result = execute_prefilter( + &mut prefilter_ctx, + &disabled.0, + &RowGroupBuildContext { + row_group_idx: 1, + row_selection: None, + fetch_metrics: None, + }, + ) + .await + .unwrap(); + assert_eq!(result.filtered_rows, 2); + assert_eq!( + result.refined_selection, + RowSelection::from(vec![RowSelector::select(1), RowSelector::skip(2)]) + ); + } + #[test] fn test_skip_prefilter_for_json_get() -> WhateverResult<()> { fn json_get_expr(base: Expr, path: &str) -> Expr { diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index b9d5016173e..9888b4ead53 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/src/mito2/src/sst/parquet/writer.rs @@ -38,7 +38,7 @@ use object_store::{FuturesAsyncWriter, ObjectStore}; use parquet::arrow::AsyncArrowWriter; use parquet::basic::{Compression, Encoding, ZstdLevel}; use parquet::file::metadata::KeyValue; -use parquet::file::properties::{WriterProperties, WriterPropertiesBuilder}; +use parquet::file::properties::WriterProperties; use parquet::schema::types::ColumnPath; use smallvec::smallvec; use snafu::{OptionExt, ResultExt}; @@ -61,7 +61,9 @@ use crate::sst::parquet::flat_format::{ FlatWriteFormat, primary_key_column_index, time_index_column_index, }; use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyWriteFormat}; -use crate::sst::parquet::{PARQUET_METADATA_KEY, SstInfo, WriteOptions}; +use crate::sst::parquet::{ + PARQUET_METADATA_KEY, SstInfo, WriteOptions, apply_float_field_encoding, +}; use crate::sst::{ DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator, maybe_wrap_schema, @@ -473,29 +475,6 @@ where Ok(results) } - /// Customizes per-column config according to schema and maybe column cardinality. - fn customize_column_config( - builder: WriterPropertiesBuilder, - region_metadata: &RegionMetadataRef, - ) -> WriterPropertiesBuilder { - let ts_col = ColumnPath::new(vec![ - region_metadata - .time_index_column() - .column_schema - .name - .clone(), - ]); - let seq_col = ColumnPath::new(vec![SEQUENCE_COLUMN_NAME.to_string()]); - let op_type_col = ColumnPath::new(vec![OP_TYPE_COLUMN_NAME.to_string()]); - - builder - .set_column_encoding(seq_col.clone(), Encoding::DELTA_BINARY_PACKED) - .set_column_dictionary_enabled(seq_col, false) - .set_column_encoding(ts_col.clone(), Encoding::DELTA_BINARY_PACKED) - .set_column_dictionary_enabled(ts_col, false) - .set_column_compression(op_type_col, Compression::UNCOMPRESSED) - } - async fn append_flat_batch( &mut self, batch: &RecordBatch, @@ -564,8 +543,22 @@ where .set_max_row_group_row_count(Some(opts.row_group_size)) .set_column_index_truncate_length(None) .set_statistics_truncate_length(None); - - let props_builder = Self::customize_column_config(props_builder, &self.metadata); + let ts_col = ColumnPath::new(vec![ + self.metadata.time_index_column().column_schema.name.clone(), + ]); + let seq_col = ColumnPath::new(vec![SEQUENCE_COLUMN_NAME.to_string()]); + let op_type_col = ColumnPath::new(vec![OP_TYPE_COLUMN_NAME.to_string()]); + let props_builder = props_builder + .set_column_encoding(seq_col.clone(), Encoding::DELTA_BINARY_PACKED) + .set_column_dictionary_enabled(seq_col, false) + .set_column_encoding(ts_col.clone(), Encoding::DELTA_BINARY_PACKED) + .set_column_dictionary_enabled(ts_col, false) + .set_column_compression(op_type_col, Compression::UNCOMPRESSED); + let props_builder = apply_float_field_encoding( + props_builder, + &self.metadata, + opts.float_field_encoding, + ); let writer_props = props_builder.build(); let sst_file_path = self.path_provider.build_sst_file_path(RegionFileId::new( diff --git a/src/mito2/src/sst/range_index.rs b/src/mito2/src/sst/range_index.rs index c4e576b6cd9..6d9eadc5532 100644 --- a/src/mito2/src/sst/range_index.rs +++ b/src/mito2/src/sst/range_index.rs @@ -28,8 +28,6 @@ pub use writer::{ }; pub use crate::sst::range_index::deleter::RangeIndexDeleter; -// Used by the upcoming query and index-building integration. -#[allow(unused_imports)] pub(crate) use crate::sst::range_index::deleter::range_index_path; const ROW_GROUP_ID_COLUMN: &str = "row_group_id"; diff --git a/src/mito2/src/test_util.rs b/src/mito2/src/test_util.rs index 8b0302a365e..dbfbc53db25 100644 --- a/src/mito2/src/test_util.rs +++ b/src/mito2/src/test_util.rs @@ -682,6 +682,7 @@ impl TestEnv { match log_store { LogStoreImpl::RaftEngine(log_store) => WorkerGroup::start( + &data_home, Arc::new(config), log_store, Arc::new(object_store_manager), @@ -693,6 +694,7 @@ impl TestEnv { .await .unwrap(), LogStoreImpl::Kafka(log_store) => WorkerGroup::start( + &data_home, Arc::new(config), log_store, Arc::new(object_store_manager), @@ -1351,6 +1353,7 @@ pub async fn put_rows(engine: &MitoEngine, region_id: RegionId, rows: Rows) { .handle_request( region_id, RegionRequest::Put(RegionPutRequest { + skip_wal: false, rows, hint: None, partition_expr_version: None, diff --git a/src/mito2/src/wal/entry_distributor.rs b/src/mito2/src/wal/entry_distributor.rs index e48c21967a5..b2b774f67de 100644 --- a/src/mito2/src/wal/entry_distributor.rs +++ b/src/mito2/src/wal/entry_distributor.rs @@ -133,24 +133,7 @@ impl WalEntryReader for WalEntryReceiver { } let stream = stream! { - let mut buffered_entry: Option = None; - while let Some(next_entry) = entry_receiver.recv().await { - match buffered_entry.take() { - Some(entry) => { - if entry.is_complete() { - yield decode_raw_entry(entry); - } else { - warn!("Ignoring incomplete entry: {}", entry); - } - buffered_entry = Some(next_entry); - }, - None => { - buffered_entry = Some(next_entry); - } - }; - } - if let Some(entry) = buffered_entry { - // Ignores tail corrupted data. + while let Some(entry) = entry_receiver.recv().await { if entry.is_complete() { yield decode_raw_entry(entry); } else { @@ -220,8 +203,10 @@ pub fn build_wal_entry_distributor_and_receivers( #[cfg(test)] mod tests { + use std::time::Duration; + use api::v1::{Mutation, OpType, WalEntry}; - use futures::{TryStreamExt, stream}; + use futures::{StreamExt, TryStreamExt, stream}; use prost::Message; use store_api::logstore::entry::{Entry, MultiplePartEntry, MultiplePartHeader, NaiveEntry}; @@ -247,6 +232,34 @@ mod tests { } } + #[tokio::test] + async fn test_delivers_complete_entry_while_channel_is_alive() { + let provider = Provider::kafka_provider("my_topic".to_string()); + let region_id = RegionId::new(1024, 1); + let wal_entry = WalEntry::default(); + let (sender, receiver) = mpsc::channel(1); + let (arg_sender, arg_receiver) = oneshot::channel(); + let mut receiver = WalEntryReceiver::new(receiver, arg_sender); + let mut stream = receiver.read(&provider, 0).unwrap(); + assert_eq!(arg_receiver.await.unwrap(), 0); + sender + .send(Entry::Naive(NaiveEntry { + provider, + region_id, + entry_id: 1, + data: wal_entry.encode_to_vec(), + })) + .await + .unwrap(); + + let entry = tokio::time::timeout(Duration::from_secs(1), stream.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(entry, (1, wal_entry)); + } + #[tokio::test] async fn test_wal_entry_distributor_without_receivers() { let provider = Provider::kafka_provider("my_topic".to_string()); diff --git a/src/mito2/src/wal/entry_reader.rs b/src/mito2/src/wal/entry_reader.rs index 272aa251736..67944351cf7 100644 --- a/src/mito2/src/wal/entry_reader.rs +++ b/src/mito2/src/wal/entry_reader.rs @@ -79,24 +79,8 @@ impl WalEntryReader for LogStoreEntryReader { let mut stream = reader.read(ns, start_id)?; let stream = stream! { - let mut buffered_entry: Option = None; while let Some(next_entry) = stream.next().await { - match buffered_entry.take() { - Some(entry) => { - if entry.is_complete() { - yield decode_raw_entry(entry); - } else { - warn!("Ignoring incomplete entry: {}", entry); - } - buffered_entry = Some(next_entry?); - }, - None => { - buffered_entry = Some(next_entry?); - } - }; - } - if let Some(entry) = buffered_entry { - // Ignores tail corrupted data. + let entry = next_entry?; if entry.is_complete() { yield decode_raw_entry(entry); } else { @@ -112,15 +96,54 @@ impl WalEntryReader for LogStoreEntryReader { #[cfg(test)] mod tests { + use std::time::Duration; + use api::v1::{Mutation, OpType, WalEntry}; - use futures::TryStreamExt; + use futures::{StreamExt, TryStreamExt, stream}; use prost::Message; use store_api::logstore::entry::{Entry, MultiplePartEntry, MultiplePartHeader}; use store_api::logstore::provider::Provider; use store_api::storage::RegionId; + use crate::error::Result; use crate::test_util::wal_util::MockRawEntryStream; + use crate::wal::EntryId; use crate::wal::entry_reader::{LogStoreEntryReader, WalEntryReader}; + use crate::wal::raw_entry_reader::{EntryStream, RawEntryReader}; + + struct LiveRawEntryReader { + entry: Entry, + } + + impl RawEntryReader for LiveRawEntryReader { + fn read(&self, _ns: &Provider, _start_id: EntryId) -> Result> { + let stream = stream::iter([Ok(self.entry.clone())]).chain(stream::pending()); + Ok(stream.boxed()) + } + } + + #[tokio::test] + async fn test_delivers_complete_entry_while_stream_is_alive() { + let provider = Provider::kafka_provider("my_topic".to_string()); + let wal_entry = WalEntry::default(); + let raw_entry_reader = LiveRawEntryReader { + entry: Entry::Naive(store_api::logstore::entry::NaiveEntry { + provider: provider.clone(), + region_id: RegionId::new(1, 1), + entry_id: 1, + data: wal_entry.encode_to_vec(), + }), + }; + let mut reader = LogStoreEntryReader::new(raw_entry_reader); + let mut entries = reader.read(&provider, 0).unwrap(); + + let entry = tokio::time::timeout(Duration::from_secs(1), entries.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(entry, (1, wal_entry)); + } #[tokio::test] async fn test_tail_corrupted_stream() { diff --git a/src/mito2/src/worker.rs b/src/mito2/src/worker.rs index 8efcf47ab68..0483217f968 100644 --- a/src/mito2/src/worker.rs +++ b/src/mito2/src/worker.rs @@ -169,7 +169,9 @@ impl WorkerGroup { /// Starts a worker group. /// /// The number of workers should be power of two. + #[allow(clippy::too_many_arguments)] pub(crate) async fn start( + data_home: &str, config: Arc, log_store: Arc, object_store_manager: ObjectStoreManagerRef, @@ -195,11 +197,7 @@ impl WorkerGroup { .with_buffer_size(Some(config.index.write_buffer_size.as_bytes() as _)); let index_build_job_pool = Arc::new(LocalScheduler::new(config.max_background_index_builds)); - let series_index_store = if config.experimental_series_index_root.trim().is_empty() { - None - } else { - Some(new_fs_cache_store(&config.experimental_series_index_root).await?) - }; + let series_index_store = series_index_store_from_config(&config, data_home).await?; let flush_job_pool = Arc::new(LocalScheduler::new(config.max_background_flushes)); let compact_job_pool = Arc::new(LocalScheduler::new(config.max_background_compactions)); let flush_semaphore = Arc::new(Semaphore::new(config.max_background_flushes)); @@ -391,6 +389,7 @@ impl WorkerGroup { /// The number of workers should be power of two. #[allow(clippy::too_many_arguments)] pub(crate) async fn start_for_test( + data_home: &str, config: Arc, log_store: Arc, object_store_manager: ObjectStoreManagerRef, @@ -410,11 +409,7 @@ impl WorkerGroup { }); let index_build_job_pool = Arc::new(LocalScheduler::new(config.max_background_index_builds)); - let series_index_store = if config.experimental_series_index_root.trim().is_empty() { - None - } else { - Some(new_fs_cache_store(&config.experimental_series_index_root).await?) - }; + let series_index_store = series_index_store_from_config(&config, data_home).await?; let flush_job_pool = Arc::new(LocalScheduler::new(config.max_background_flushes)); let compact_job_pool = Arc::new(LocalScheduler::new(config.max_background_compactions)); let flush_semaphore = Arc::new(Semaphore::new(config.max_background_flushes)); @@ -517,6 +512,19 @@ fn region_id_to_index(id: RegionId, num_workers: usize) -> usize { % num_workers } +/// Opens the fixed local series-index store only when the feature is enabled. +async fn series_index_store_from_config( + config: &MitoConfig, + data_home: &str, +) -> Result> { + if !config.experimental_enable_series_index { + return Ok(None); + } + + let root = Path::new(data_home).join("series_index"); + new_fs_cache_store(&root.to_string_lossy()).await.map(Some) +} + pub async fn write_cache_from_config( config: &MitoConfig, puffin_manager_factory: PuffinManagerFactory, @@ -1620,18 +1628,40 @@ mod tests { #[tokio::test] async fn test_worker_group_start_stop() { - let env = TestEnv::with_prefix("group-stop").await; - let group = env - .create_worker_group(MitoConfig { - num_workers: 4, - experimental_series_index_root: "series-index".to_string(), - ..Default::default() - }) - .await; + for (enabled, existing_index) in [(false, false), (true, false), (false, true)] { + // Use a fresh WAL directory for each case: stopping workers does not stop the log store. + let env = TestEnv::with_prefix("group-stop").await; + let root = env.data_home().join("series_index"); + if existing_index { + tokio::fs::create_dir_all(&root).await.unwrap(); + tokio::fs::write(root.join("retained"), b"index") + .await + .unwrap(); + } + let group = env + .create_worker_group(MitoConfig { + num_workers: 4, + experimental_enable_series_index: enabled, + ..Default::default() + }) + .await; - tokio::time::timeout(Duration::from_secs(5), group.stop()) - .await - .expect("series-index tasks should stop without waiting for their interval") - .unwrap(); + for worker in &group.workers { + assert_eq!(worker.series_index_task_state.is_some(), enabled); + assert_eq!(worker.series_index_handle.lock().await.is_some(), enabled); + } + assert_eq!(root.is_dir(), enabled || existing_index); + if existing_index { + assert_eq!( + tokio::fs::read(root.join("retained")).await.unwrap(), + b"index" + ); + } + + tokio::time::timeout(Duration::from_secs(5), group.stop()) + .await + .expect("series-index tasks should stop without waiting for their interval") + .unwrap(); + } } } diff --git a/src/mito2/src/worker/handle_bulk_insert.rs b/src/mito2/src/worker/handle_bulk_insert.rs index b4a7ae4a182..c094c1fa6dc 100644 --- a/src/mito2/src/worker/handle_bulk_insert.rs +++ b/src/mito2/src/worker/handle_bulk_insert.rs @@ -96,6 +96,7 @@ impl RegionWorkerLoop { None }; pending_bulk_request.push(SenderBulkRequest { + skip_wal: request.skip_wal, sender, request: part, region_id: request.region_id, diff --git a/src/mito2/src/worker/handle_write.rs b/src/mito2/src/worker/handle_write.rs index 96940756248..6fee8fb4b2a 100644 --- a/src/mito2/src/worker/handle_write.rs +++ b/src/mito2/src/worker/handle_write.rs @@ -451,6 +451,7 @@ impl RegionWorkerLoop { sender_req.request.hint, sender_req.sender, None, + sender_req.request.skip_wal, ); } } @@ -560,7 +561,7 @@ impl RegionWorkerLoop { } // Collect requests by region. - if !region_ctx.push_bulk(bulk_req.sender, bulk_req.request, None) { + if !region_ctx.push_bulk(bulk_req.sender, bulk_req.request, None, bulk_req.skip_wal) { return; } } @@ -613,14 +614,21 @@ async fn write_wal( region_ctxs: &mut HashMap, ) -> bool { let mut wal_writer = wal.writer(); + let mut has_wal_entries = false; for region_ctx in region_ctxs.values_mut() { if region_ctx.skip_wal() { continue; } if let Err(e) = region_ctx.add_wal_entry(&mut wal_writer).map_err(Arc::new) { region_ctx.set_error(e); + } else { + has_wal_entries = true; } } + // All-skipped batches should not touch the log store, even with an empty append. + if !has_wal_entries { + return true; + } match wal_writer.write_to_wal().await.map_err(Arc::new) { Ok(response) => { for (region_id, region_ctx) in region_ctxs.iter_mut() { @@ -880,6 +888,7 @@ mod tests { fn new_region_ctx( region_id: RegionId, + skip_wal: bool, ) -> (RegionWriteCtx, oneshot::Receiver>) { let version_control = Arc::new(VersionControlBuilder::new().build()); let mut ctx = RegionWriteCtx::new( @@ -908,10 +917,57 @@ mod tests { None, OptionOutputTx::from(tx), None, + skip_wal, ); (ctx, rx) } + #[tokio::test] + async fn test_request_skip_wal_does_not_append_empty_batch() { + // Only change the request flag. A failing log store demonstrates that + // the all-skipped path never invokes append_batch, including empty appends. + check_request_skip_wal_does_not_append_empty_batch(false, false).await; + check_request_skip_wal_does_not_append_empty_batch(false, true).await; + check_request_skip_wal_does_not_append_empty_batch(true, false).await; + check_request_skip_wal_does_not_append_empty_batch(true, true).await; + } + + async fn check_request_skip_wal_does_not_append_empty_batch(skip_wal: bool, bulk: bool) { + let region_id = RegionId::new(1, 1); + let wal = Wal::new(Arc::new(MockLogStore { + fail_append: true, + ..Default::default() + })); + let (ctx, rx) = if bulk { + let version_control = Arc::new(VersionControlBuilder::new().build()); + let mut ctx = RegionWriteCtx::new( + region_id, + &version_control, + Provider::raft_engine_provider(region_id.as_u64()), + None, + ); + let (tx, rx) = oneshot::channel(); + assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(1), None, skip_wal)); + (ctx, rx) + } else { + new_region_ctx(region_id, skip_wal) + }; + let version_control = ctx.version_control().clone(); + let mut contexts = HashMap::from([(region_id, ctx)]); + assert_eq!(write_wal(&wal, &mut contexts).await, skip_wal); + if skip_wal { + let ctx = contexts.get_mut(®ion_id).unwrap(); + assert_eq!(ctx.next_entry_id(), 1); + ctx.write_memtable().await; + ctx.write_bulk().await; + ctx.publish_sequence_and_entry_id(); + assert_eq!(version_control.committed_sequence(), 1); + assert_eq!(version_control.current().last_entry_id, 0); + } + drop(contexts); + assert_eq!(rx.await.unwrap().is_ok(), skip_wal); + } + #[tokio::test] async fn test_write_wal_skips_region_failed_to_build_entry() { let failing_region = RegionId::new(1, 1); @@ -922,10 +978,10 @@ mod tests { })); let mut region_ctxs = HashMap::new(); - let (ctx, failing_rx) = new_region_ctx(failing_region); + let (ctx, failing_rx) = new_region_ctx(failing_region, false); let failing_committed_sequence = ctx.version_control().committed_sequence(); region_ctxs.insert(failing_region, ctx); - let (ctx, ok_rx) = new_region_ctx(ok_region); + let (ctx, ok_rx) = new_region_ctx(ok_region, false); let ok_committed_sequence = ctx.version_control().committed_sequence(); region_ctxs.insert(ok_region, ctx); let entry_id = region_ctxs[&ok_region].next_entry_id(); @@ -963,6 +1019,11 @@ mod tests { #[tokio::test] async fn test_bulk_write_sequence_not_committed_before_install_worker_level() { + check_bulk_write_sequence_not_committed_before_install_worker_level(false).await; + check_bulk_write_sequence_not_committed_before_install_worker_level(true).await; + } + + async fn check_bulk_write_sequence_not_committed_before_install_worker_level(skip_wal: bool) { let region_id = RegionId::new(1, 1); let version_control = Arc::new(VersionControlBuilder::new().build()); @@ -974,7 +1035,7 @@ mod tests { None, ); let (tx, rx) = oneshot::channel(); - assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(3), None)); + assert!(ctx.push_bulk(OptionOutputTx::from(tx), new_bulk_part(3), None, skip_wal)); region_ctxs.insert(region_id, ctx); let wal = Wal::new(Arc::new(MockLogStore::default())); @@ -1026,7 +1087,7 @@ mod tests { })); let mut region_ctxs = HashMap::new(); - let (ctx, rx) = new_region_ctx(failing_region); + let (ctx, rx) = new_region_ctx(failing_region, false); region_ctxs.insert(failing_region, ctx); // Writing an empty batch to the WAL succeeds, the failed region must not panic @@ -1047,7 +1108,7 @@ mod tests { })); let mut region_ctxs = HashMap::new(); - let (ctx, rx) = new_region_ctx(region_id); + let (ctx, rx) = new_region_ctx(region_id, false); region_ctxs.insert(region_id, ctx); assert!(!write_wal(&wal, &mut region_ctxs).await); diff --git a/src/object-store/src/secure_fs.rs b/src/object-store/src/secure_fs.rs index 581cc2837c0..b940755c5b3 100644 --- a/src/object-store/src/secure_fs.rs +++ b/src/object-store/src/secure_fs.rs @@ -73,6 +73,20 @@ impl SecureFsRoot { &self.path } + /// Checks for a regular file without following the final symbolic link. + pub async fn is_file(&self, path: &str) -> Result { + let path = backend_path(path).map_err(new_std_io_error)?; + let root = self.clone(); + common_runtime::spawn_blocking_global(move || { + root.dir + .symlink_metadata(path) + .map(|metadata| metadata.is_file()) + }) + .await + .map_err(new_task_join_error)? + .map_err(new_std_io_error) + } + /// Opens a descendant directory without leaving this capability root. pub fn open_subdir(&self, path: impl AsRef) -> io::Result { let path = normalize_relative_path(path.as_ref())?; diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index c9258df2c0f..4303e831002 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -85,7 +85,9 @@ tonic.workspace = true tracing.workspace = true [dev-dependencies] +axum.workspace = true catalog = { workspace = true, features = ["testing"] } common-meta = { workspace = true, features = ["testing"] } common-test-util.workspace = true path-slash = "0.2" +url.workspace = true diff --git a/src/operator/src/bulk_insert.rs b/src/operator/src/bulk_insert.rs index 22fae236196..e35779ec979 100644 --- a/src/operator/src/bulk_insert.rs +++ b/src/operator/src/bulk_insert.rs @@ -44,6 +44,7 @@ impl Inserter { raw_flight_data: FlightData, record_batch: RecordBatch, schema_bytes: Bytes, + skip_wal: bool, ) -> error::Result { let table_info = table.table_info(); let table_id = table_info.table_id(); @@ -103,6 +104,7 @@ impl Inserter { ..Default::default() }), body: Some(region_request::Body::BulkInsert(BulkInsertRequest { + skip_wal, region_id: region_id.as_u64(), partition_expr_version: partition_expr_version .map(|value| PartitionExprVersion { value }), @@ -218,6 +220,7 @@ impl Inserter { ..Default::default() }), body: Some(region_request::Body::BulkInsert(BulkInsertRequest { + skip_wal, region_id: region_id.as_u64(), partition_expr_version: partition_expr_version .map(|value| PartitionExprVersion { value }), diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs index 0ed3c74bbed..ae680be33ba 100644 --- a/src/operator/src/error.rs +++ b/src/operator/src/error.rs @@ -1014,8 +1014,10 @@ impl ErrorExt for Error { Error::BuildDfLogicalPlan { .. } | Error::BuildTableMeta { .. } | Error::MissingInsertBody { .. } => StatusCode::Internal, - Error::ExecuteAdminFunction { .. } - | Error::EncodeJson { .. } + Error::ExecuteAdminFunction { error, .. } => admin_external_error(error) + .map(ErrorExt::status_code) + .unwrap_or(StatusCode::Unexpected), + Error::EncodeJson { .. } | Error::DeserializePartitionExpr { .. } | Error::SerializePartitionExpr { .. } => StatusCode::Unexpected, Error::AdminFunctionCancelled => StatusCode::Cancelled, @@ -1135,6 +1137,9 @@ impl ErrorExt for Error { | Error::MissingTimeIndexColumn { source, .. } => source.retry_hint(), Error::Cast { source, .. } => source.retry_hint(), + Error::ExecuteAdminFunction { error, .. } => admin_external_error(error) + .map(ErrorExt::retry_hint) + .unwrap_or(RetryHint::NonRetryable), Error::ParseSql { source, .. } => source.retry_hint(), Error::Catalog { source, .. } => source.retry_hint(), Error::SubstraitCodec { source, .. } => source.retry_hint(), @@ -1154,4 +1159,36 @@ impl ErrorExt for Error { } } +fn admin_external_error(error: &DataFusionError) -> Option<&BoxedError> { + match error { + DataFusionError::External(error) => error.downcast_ref::(), + DataFusionError::Diagnostic(_, error) => admin_external_error(error), + _ => None, + } +} + define_into_tonic_status!(Error); + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn admin_function_preserves_external_error_metadata() { + let external = BoxedError::new(meta_client::error::Error::MetaServer { + code: StatusCode::TableUnavailable, + msg: "leader changed".to_string(), + tonic_code: tonic::Code::Unavailable, + retry_hint: RetryHint::Retryable, + location: snafu::location!(), + }); + let error = Error::ExecuteAdminFunction { + msg: "test_admin".to_string(), + error: DataFusionError::External(Box::new(external)), + location: snafu::location!(), + }; + + assert_eq!(StatusCode::TableUnavailable, error.status_code()); + assert_eq!(RetryHint::Retryable, error.retry_hint()); + } +} diff --git a/src/operator/src/expr_helper.rs b/src/operator/src/expr_helper.rs index 925e062cf5a..b28af78d383 100644 --- a/src/operator/src/expr_helper.rs +++ b/src/operator/src/expr_helper.rs @@ -826,9 +826,16 @@ pub(crate) fn to_alter_table_expr( AlterTableOperation::ModifyColumnType { column_name, target_type, + .. } => { let target_type = sql_data_type_to_concrete_data_type(&target_type).context(ParseSqlSnafu)?; + if target_type.is_json2() { + return NotSupportedSnafu { + feat: "ALTER TABLE MODIFY COLUMN to JSON2 type", + } + .fail(); + } let (target_type, target_type_extension) = ColumnDataTypeWrapper::try_from(target_type) .map(|w| w.to_parts()) .context(ColumnDataTypeSnafu)?; @@ -1744,6 +1751,34 @@ SELECT max(c1), min(c2) FROM schema_2.table_2;"; assert!(modify_column_type.target_type_extension.is_none()); } + #[test] + fn test_alter_json2_is_not_supported() { + for sql in [ + "ALTER TABLE monitor MODIFY COLUMN payload JSON2;", + "ALTER TABLE monitor MODIFY COLUMN payload JSON2 (service STRING);", + ] { + let stmt = ParserContext::create_with_dialect( + sql, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap() + .pop() + .unwrap(); + + let Statement::AlterTable(alter_table) = stmt else { + unreachable!() + }; + let err = to_alter_table_expr(alter_table, &QueryContext::arc()).unwrap_err(); + + assert!(matches!(err, crate::error::Error::NotSupported { .. })); + assert_eq!( + "Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type", + err.to_string() + ); + } + } + #[test] fn test_to_repartition_request() { let sql = r#" diff --git a/src/operator/src/insert.rs b/src/operator/src/insert.rs index 52a73892dc0..a4a69ef8ba9 100644 --- a/src/operator/src/insert.rs +++ b/src/operator/src/insert.rs @@ -265,6 +265,8 @@ impl Inserter { accommodate_existing_schema: bool, is_single_value: bool, ) -> Result { + let skip_wal = ctx.skip_wal(); + // remove empty requests requests.inserts.retain(|req| { req.rows @@ -297,7 +299,7 @@ impl Inserter { instant_table_ids, self.partition_manager.as_ref(), ) - .convert(requests) + .convert(requests, skip_wal) .await?; self.do_request(inserts, &table_infos, &ctx).await @@ -311,6 +313,8 @@ impl Inserter { statement_executor: &StatementExecutor, physical_table: String, ) -> Result { + let skip_wal = ctx.skip_wal(); + // remove empty requests requests.inserts.retain(|req| { req.rows @@ -343,7 +347,7 @@ impl Inserter { .map(|info| (info.name.clone(), info.clone())) .collect::>(); let inserts = RowToRegion::new(name_to_info, instant_table_ids, &self.partition_manager) - .convert(requests) + .convert(requests, skip_wal) .await?; self.do_request(inserts, &table_infos, &ctx).await @@ -1743,6 +1747,24 @@ mod tests { assert!(table_is_native_histogram(&table)); } + #[test] + fn test_skip_wal_does_not_change_table_options() { + check_skip_wal_does_not_change_table_options(false); + check_skip_wal_does_not_change_table_options(true); + } + + fn check_skip_wal_does_not_change_table_options(skip_wal: bool) { + let ctx = Arc::new(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + )); + ctx.set_skip_wal(skip_wal); + let mut options = Default::default(); + fill_table_options_for_create(&mut options, &AutoCreateTableType::Physical, &ctx); + assert!(!options.contains_key(session::hints::INSERT_SKIP_WAL_HINT)); + assert!(!options.contains_key("skip_wal")); + } + #[test] fn test_last_non_null_create_options_preserve_default_without_append_mode() { let ctx = Arc::new(QueryContext::with( diff --git a/src/operator/src/req_convert/common/partitioner.rs b/src/operator/src/req_convert/common/partitioner.rs index b52a6f6be6c..3353e2fcda7 100644 --- a/src/operator/src/req_convert/common/partitioner.rs +++ b/src/operator/src/req_convert/common/partitioner.rs @@ -34,6 +34,7 @@ impl<'a> Partitioner<'a> { &self, table_info: &TableInfo, rows: Rows, + skip_wal: bool, ) -> Result> { let table_id = table_info.table_id(); let requests = self @@ -44,6 +45,7 @@ impl<'a> Partitioner<'a> { .into_iter() .map( |(region_number, (rows, partition_expr_version))| InsertRequest { + skip_wal, region_id: RegionId::new(table_id, region_number).into(), rows: Some(rows), partition_expr_version: partition_expr_version diff --git a/src/operator/src/req_convert/insert/row_to_region.rs b/src/operator/src/req_convert/insert/row_to_region.rs index 9a03cd71219..f512a6c9acb 100644 --- a/src/operator/src/req_convert/insert/row_to_region.rs +++ b/src/operator/src/req_convert/insert/row_to_region.rs @@ -45,6 +45,7 @@ impl<'a> RowToRegion<'a> { pub async fn convert( &self, requests: RowInsertRequests, + skip_wal: bool, ) -> Result { let mut region_request = Vec::with_capacity(requests.inserts.len()); let mut instant_request = Vec::with_capacity(requests.inserts.len()); @@ -55,7 +56,7 @@ impl<'a> RowToRegion<'a> { let table_id = table_info.table_id(); let requests = Partitioner::new(self.partition_manager) - .partition_insert_requests(table_info, rows) + .partition_insert_requests(table_info, rows, skip_wal) .await?; if self.instant_table_ids.contains(&table_id) { @@ -81,3 +82,79 @@ impl<'a> RowToRegion<'a> { .context(TableNotFoundSnafu { table_name }) } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use api::v1::helper::tag_column_schema; + use api::v1::value::ValueData; + use api::v1::{ColumnDataType, Row, RowInsertRequest, Rows, Value}; + + use super::*; + use crate::tests::{ + create_partition_rule_manager, new_test_table_info, prepare_mocked_backend, + }; + + #[tokio::test] + async fn test_partitioned_insert_skip_wal_normal_and_instant() { + check_partitioned_insert_skip_wal(false, false).await; + check_partitioned_insert_skip_wal(false, true).await; + check_partitioned_insert_skip_wal(true, false).await; + check_partitioned_insert_skip_wal(true, true).await; + } + + async fn check_partitioned_insert_skip_wal(instant: bool, skip_wal: bool) { + let backend = prepare_mocked_backend().await; + let partition_manager = create_partition_rule_manager(backend).await; + let table_info = Arc::new(new_test_table_info(1, "table_1", [1, 2, 3].into_iter())); + let instant_table_ids = if instant { + HashSet::from_iter([1]) + } else { + HashSet::default() + }; + let converter = RowToRegion::new( + HashMap::from_iter([("table_1".to_string(), table_info)]), + instant_table_ids, + &partition_manager, + ); + let requests = RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "table_1".to_string(), + rows: Some(Rows { + schema: vec![tag_column_schema("a", ColumnDataType::Int32)], + rows: [1, 11, 101] + .into_iter() + .map(|value| Row { + values: vec![Value { + value_data: Some(ValueData::I32Value(value)), + }], + }) + .collect(), + }), + }], + }; + let result = converter.convert(requests, skip_wal).await.unwrap(); + let (selected, other) = if instant { + (result.instant_requests, result.normal_requests) + } else { + (result.normal_requests, result.instant_requests) + }; + assert!(other.requests.is_empty()); + assert_eq!(selected.requests.len(), 3); + assert!( + selected + .requests + .iter() + .all(|request| request.skip_wal == skip_wal) + ); + assert_eq!( + selected + .requests + .iter() + .map(|request| request.rows.as_ref().unwrap().rows.len()) + .sum::(), + 3 + ); + } +} diff --git a/src/operator/src/req_convert/insert/stmt_to_region.rs b/src/operator/src/req_convert/insert/stmt_to_region.rs index b2424d3a17a..e811facf74d 100644 --- a/src/operator/src/req_convert/insert/stmt_to_region.rs +++ b/src/operator/src/req_convert/insert/stmt_to_region.rs @@ -154,7 +154,7 @@ impl<'a> StatementToRegion<'a> { } let requests = Partitioner::new(self.partition_manager) - .partition_insert_requests(&table_info, Rows { schema, rows }) + .partition_insert_requests(&table_info, Rows { schema, rows }, query_ctx.skip_wal()) .await?; let requests = RegionInsertRequests { requests }; if table_info.is_ttl_instant_table() { diff --git a/src/operator/src/req_convert/insert/table_to_region.rs b/src/operator/src/req_convert/insert/table_to_region.rs index a89b4d6204d..4fbbbaa02c4 100644 --- a/src/operator/src/req_convert/insert/table_to_region.rs +++ b/src/operator/src/req_convert/insert/table_to_region.rs @@ -46,7 +46,7 @@ impl<'a> TableToRegion<'a> { let rows = Rows { schema, rows }; let requests = Partitioner::new(self.partition_manager) - .partition_insert_requests(self.table_info, rows) + .partition_insert_requests(self.table_info, rows, request.skip_wal) .await?; let requests = RegionInsertRequests { requests }; @@ -84,6 +84,11 @@ mod tests { #[tokio::test] async fn test_insert_request_table_to_region() { + check_insert_request_table_to_region(false).await; + check_insert_request_table_to_region(true).await; + } + + async fn check_insert_request_table_to_region(skip_wal: bool) { // region to datanode placement: // 1 -> 1 // 2 -> 2 @@ -100,12 +105,13 @@ mod tests { let converter = TableToRegion::new(&table_info, &partition_manager); - let table_request = build_table_request(Arc::new(Int32Vector::from(vec![ + let mut table_request = build_table_request(Arc::new(Int32Vector::from(vec![ Some(1), None, Some(11), Some(101), ]))); + table_request.skip_wal = skip_wal; let versions = partition_manager .find_physical_partition_info(1) .await @@ -127,21 +133,26 @@ mod tests { let region_request = region_id_to_region_requests.remove(®ion_id).unwrap(); assert_eq!( region_request, - build_region_request(vec![Some(101)], region_id, versions[®ion_id]) + build_region_request(vec![Some(101)], region_id, versions[®ion_id], skip_wal) ); let region_id = RegionId::new(1, 2).as_u64(); let region_request = region_id_to_region_requests.remove(®ion_id).unwrap(); assert_eq!( region_request, - build_region_request(vec![Some(11)], region_id, versions[®ion_id]) + build_region_request(vec![Some(11)], region_id, versions[®ion_id], skip_wal) ); let region_id = RegionId::new(1, 3).as_u64(); let region_request = region_id_to_region_requests.remove(®ion_id).unwrap(); assert_eq!( region_request, - build_region_request(vec![Some(1), None], region_id, versions[®ion_id]) + build_region_request( + vec![Some(1), None], + region_id, + versions[®ion_id], + skip_wal + ) ); } @@ -151,6 +162,7 @@ mod tests { schema_name: DEFAULT_SCHEMA_NAME.to_string(), table_name: "table_1".to_string(), columns_values: HashMap::from([("a".to_string(), vector)]), + skip_wal: false, } } @@ -158,8 +170,10 @@ mod tests { rows: Vec>, region_id: u64, version: Option, + skip_wal: bool, ) -> RegionInsertRequest { RegionInsertRequest { + skip_wal, region_id, rows: Some(Rows { schema: vec![tag_column_schema("a", ColumnDataType::Int32)], diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 21c3ed61464..150020bc9f3 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -63,7 +63,7 @@ use query::QueryEngineRef; use query::parser::QueryStatement; use session::context::{Channel, QueryContextBuilder, QueryContextRef}; use session::table_name::table_idents_to_full_name; -use set::{set_query_timeout, set_read_preference}; +use set::{set_query_timeout, set_read_preference, set_skip_wal}; use snafu::{OptionExt, ResultExt, ensure}; use sql::ast::ObjectNamePartExt; use sql::statements::OptionMap; @@ -534,6 +534,7 @@ impl StatementExecutor { match var_name.as_str() { "READ_PREFERENCE" => set_read_preference(set_var.value, query_ctx)?, + "SKIP_WAL" => set_skip_wal(set_var.value, query_ctx)?, "@@TIME_ZONE" | "@@SESSION.TIME_ZONE" | "TIMEZONE" | "TIME_ZONE" => { set_timezone(set_var.value, query_ctx)? diff --git a/src/operator/src/statement/copy_table_from.rs b/src/operator/src/statement/copy_table_from.rs index 39bc3dafcdb..c69d0477dbd 100644 --- a/src/operator/src/statement/copy_table_from.rs +++ b/src/operator/src/statement/copy_table_from.rs @@ -27,7 +27,7 @@ use common_datasource::file_format::json::JsonFormat; use common_datasource::file_format::orc::{ReaderAdapter, infer_orc_schema, new_orc_stream_reader}; use common_datasource::file_format::{FileFormat, Format, file_to_stream}; use common_datasource::lister::{Lister, Source}; -use common_datasource::object_store::build_backend_with_path; +use common_datasource::object_store::{LocalFileAccess, build_backend_with_path}; use common_query::{OutputCost, OutputRows}; use common_recordbatch::DfSendableRecordBatchStream; use common_recordbatch::adapter::RecordBatchStreamTypeAdapter; @@ -44,7 +44,7 @@ use datatypes::arrow::datatypes::{DataType as ArrowDataType, Schema, SchemaRef}; use datatypes::arrow::record_batch::RecordBatch; use datatypes::vectors::Helper; use futures_util::StreamExt; -use object_store::{Entry, EntryMode, ObjectStore}; +use object_store::{EntryMode, ObjectStore}; use regex::Regex; use session::context::QueryContextRef; use snafu::{ResultExt, ensure}; @@ -92,43 +92,59 @@ impl FileMetadata { } } -impl StatementExecutor { - async fn list_copy_from_entries( - &self, - req: &CopyTableRequest, - ) -> Result<(ObjectStore, Vec)> { - let backend = - build_backend_with_path(&req.location, &req.connection, &self.local_file_access) - .await - .context(error::BuildBackendSnafu)?; - let regex = req - .pattern - .as_ref() - .map(|x| Regex::new(x)) - .transpose() - .context(error::BuildRegexSnafu)?; +async fn list_copy_from_paths( + req: &CopyTableRequest, + local_file_access: &LocalFileAccess, +) -> Result<(ObjectStore, Vec)> { + let backend = build_backend_with_path(&req.location, &req.connection, local_file_access) + .await + .context(error::BuildBackendSnafu)?; + let regex = req + .pattern + .as_ref() + .map(|x| Regex::new(x)) + .transpose() + .context(error::BuildRegexSnafu)?; - let source = if let Some(filename) = backend.object_path { - Source::Filename(filename) + // Listing a known file's parent for every table makes COPY DATABASE do + // quadratic directory work. + if let Some(filename) = &backend.object_path { + let is_file = backend + .is_file(filename) + .await + .with_context(|_| common_datasource::error::ListObjectsSnafu { + path: req.location.clone(), + }) + .context(error::ListObjectsSnafu)?; + let paths = if is_file { + vec![filename.clone()] } else { - Source::Dir + vec![] }; - - let lister = Lister::new( - backend.object_store.clone(), - source.clone(), - req.location.clone(), - regex, - ); - - let entries = lister.list().await.context(error::ListObjectsSnafu)?; - debug!( - "Copy from location: {:?}, {source:?}, entries: {entries:?}", - req.location - ); - Ok((backend.object_store, entries)) + return Ok((backend.object_store, paths)); } + let lister = Lister::new( + backend.object_store.clone(), + Source::Dir, + req.location.clone(), + regex, + ); + + let entries = lister.list().await.context(error::ListObjectsSnafu)?; + debug!( + "Copy from location: {:?}, entries: {entries:?}", + req.location + ); + let paths = entries + .into_iter() + .filter(|entry| entry.metadata().mode() == EntryMode::FILE) + .map(|entry| entry.path().to_string()) + .collect(); + Ok((backend.object_store, paths)) +} + +impl StatementExecutor { async fn collect_metadata( &self, object_store: &ObjectStore, @@ -371,8 +387,8 @@ impl StatementExecutor { }; let table = self.get_table(&table_ref).await?; let format = Format::try_from(&req.with).context(error::ParseFileFormatSnafu)?; - let (object_store, entries) = self.list_copy_from_entries(&req).await?; - let mut files = Vec::with_capacity(entries.len()); + let (object_store, paths) = list_copy_from_paths(&req, &self.local_file_access).await?; + let mut files = Vec::with_capacity(paths.len()); let table_schema = table.schema().arrow_schema().clone(); let filters = table .schema() @@ -383,13 +399,9 @@ impl StatementExecutor { .into_iter() .collect::>(); - for entry in entries.iter() { - if entry.metadata().mode() != EntryMode::FILE { - continue; - } - let path = entry.path(); + for path in paths { let file_metadata = self - .collect_metadata(&object_store, format.clone(), path.to_string()) + .collect_metadata(&object_store, format.clone(), path) .await?; validate_csv_headers_if_required(&file_metadata, &table_schema)?; @@ -487,6 +499,7 @@ impl StatementExecutor { schema_name: req.schema_name.clone(), table_name: req.table_name.clone(), columns_values, + skip_wal: query_ctx.skip_wal(), }, query_ctx.clone(), )); @@ -825,6 +838,236 @@ mod tests { use super::*; + fn copy_from_request(location: &str, pattern: Option<&str>) -> CopyTableRequest { + CopyTableRequest { + catalog_name: "greptime".into(), + schema_name: "public".into(), + table_name: "test".into(), + location: location.into(), + with: HashMap::new(), + connection: HashMap::new(), + pattern: pattern.map(String::from), + direction: table::requests::CopyDirection::Import, + timestamp_range: None, + limit: None, + } + } + + #[tokio::test] + async fn test_copy_from_local_paths() { + use common_test_util::temp_dir::create_temp_dir; + + let dir = create_temp_dir("copy_from_paths"); + let nested = dir.path().join("nested"); + std::fs::create_dir_all(nested.join("data.dir")).unwrap(); + std::fs::write(nested.join("data file.parquet"), b"data").unwrap(); + std::fs::write(nested.join("other.csv"), b"other").unwrap(); + let access = LocalFileAccess::sandboxed(dir.path()).unwrap(); + for location in [ + "nested/data file.parquet".to_string(), + url::Url::from_file_path(nested.join("data file.parquet")) + .unwrap() + .to_string(), + ] { + // PATTERN applies to directories, not an explicitly selected file. + let req = copy_from_request(&location, Some("does-not-match")); + let (store, paths) = list_copy_from_paths(&req, &access).await.unwrap(); + assert_eq!(paths, ["data file.parquet"]); + assert_eq!(store.read(&paths[0]).await.unwrap().to_vec(), b"data"); + } + + let req = copy_from_request("nested/", Some("^data")); + let (_, paths) = list_copy_from_paths(&req, &access).await.unwrap(); + assert_eq!(paths, ["data file.parquet"]); + let req = copy_from_request("nested/", None); + let (_, mut paths) = list_copy_from_paths(&req, &access).await.unwrap(); + paths.sort(); + assert_eq!(paths, ["data file.parquet", "other.csv"]); + let req = copy_from_request("nested/data.dir", None); + assert!( + list_copy_from_paths(&req, &access) + .await + .unwrap() + .1 + .is_empty() + ); + let req = copy_from_request("nested/", Some("does-not-match")); + assert!( + list_copy_from_paths(&req, &access) + .await + .unwrap() + .1 + .is_empty() + ); + + let req = copy_from_request("nested/missing.parquet", None); + assert!(matches!( + list_copy_from_paths(&req, &access).await, + Err(error::Error::ListObjects { + source: common_datasource::error::Error::ListObjects { error, .. }, .. + }) if error.kind() == object_store::ErrorKind::NotFound + )); + let req = copy_from_request("nested/data file.parquet", Some("[")); + assert!(matches!( + list_copy_from_paths(&req, &access).await, + Err(error::Error::BuildRegex { .. }) + )); + let req = copy_from_request("nested/data file.parquet", None); + assert!(matches!( + list_copy_from_paths(&req, &LocalFileAccess::Disabled).await, + Err(error::Error::BuildBackend { + source: common_datasource::error::Error::LocalFileAccessDisabled { .. }, + .. + }) + )); + let req = copy_from_request("../outside.parquet", None); + assert!(matches!( + list_copy_from_paths(&req, &access).await, + Err(error::Error::BuildBackend { + source: common_datasource::error::Error::LocalFileAccessDenied { .. }, + .. + }) + )); + } + + #[cfg(unix)] + #[tokio::test] + async fn test_copy_from_local_symlinks() { + use common_test_util::temp_dir::create_temp_dir; + + let dir = create_temp_dir("copy_from_symlinks"); + std::fs::write(dir.path().join("data.parquet"), b"data").unwrap(); + std::os::unix::fs::symlink("data.parquet", dir.path().join("link.parquet")).unwrap(); + std::os::unix::fs::symlink("missing.parquet", dir.path().join("dangling.parquet")).unwrap(); + let access = LocalFileAccess::sandboxed(dir.path()).unwrap(); + + let req = copy_from_request("link.parquet", None); + assert!( + list_copy_from_paths(&req, &access) + .await + .unwrap() + .1 + .is_empty() + ); + let req = copy_from_request("./", None); + let (_, paths) = list_copy_from_paths(&req, &access).await.unwrap(); + assert_eq!(paths, ["data.parquet"]); + let req = copy_from_request("dangling.parquet", None); + assert!(matches!( + list_copy_from_paths(&req, &access).await, + Err(error::Error::ListObjects { + source: common_datasource::error::Error::ListObjects { error, .. }, .. + }) if error.kind() == object_store::ErrorKind::NotFound + )); + } + + #[tokio::test] + async fn test_copy_from_s3_paths_without_relisting() { + use std::sync::Mutex; + + use axum::Router; + use axum::http::{Method, StatusCode, Uri}; + use axum::response::IntoResponse; + + let requests = Arc::new(Mutex::new(Vec::new())); + let recorded = requests.clone(); + let app = Router::new().fallback(move |method: Method, uri: Uri| { + let recorded = recorded.clone(); + async move { + recorded.lock().unwrap().push((method.clone(), uri.clone())); + if method == Method::HEAD { + let status = if uri.path().ends_with("missing.parquet") { + StatusCode::NOT_FOUND + } else if uri.path().ends_with("denied.parquet") { + StatusCode::FORBIDDEN + } else { + StatusCode::OK + }; + return (status, [("content-length", "0")]).into_response(); + } + assert_eq!(method, Method::GET); + let query = + axum::extract::Query::>::try_from_uri(&uri).unwrap(); + assert_eq!(query.get("prefix").unwrap(), "backup/schema/1/"); + assert_eq!(query.get("list-type").unwrap(), "2"); + let files = (0..16).map(|i| format!( + "backup/schema/1/data.{i}.parquet2026-09-01T00:00:00Z0" + )).collect::(); + format!( + "false{files}\ + backup/schema/1/other.csv2026-09-01T00:00:00Z0\ + backup/schema/1/data.dir/\ + " + ) + .into_response() + } + }); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let connection = HashMap::from([ + ("endpoint".into(), endpoint), + ("region".into(), "us-east-1".into()), + ("access_key_id".into(), "test".into()), + ("secret_access_key".into(), "test".into()), + ("disable_ec2_metadata".into(), "true".into()), + ]); + for i in 0..16 { + let filename = format!("data.{i}.parquet"); + let mut req = copy_from_request( + &format!("s3://bucket/backup/schema/1/{filename}"), + Some("does-not-match"), + ); + req.connection = connection.clone(); + let (_, paths) = list_copy_from_paths(&req, &LocalFileAccess::Disabled) + .await + .unwrap(); + assert_eq!(paths, [filename]); + } + { + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 16); + assert!(requests.iter().all(|(method, uri)| method == Method::HEAD + && uri.path().starts_with("/bucket/backup/schema/1/"))); + } + let mut req = copy_from_request("s3://bucket/backup/schema/1/", Some(r"^data\.")); + req.connection = connection.clone(); + let (_, mut paths) = list_copy_from_paths(&req, &LocalFileAccess::Disabled) + .await + .unwrap(); + paths.sort(); + let mut expected = (0..16) + .map(|i| format!("data.{i}.parquet")) + .collect::>(); + expected.sort(); + assert_eq!(paths, expected); + assert_eq!( + requests + .lock() + .unwrap() + .iter() + .filter(|(method, _)| *method == Method::GET) + .count(), + 1 + ); + + for (name, kind) in [ + ("missing", object_store::ErrorKind::NotFound), + ("denied", object_store::ErrorKind::PermissionDenied), + ] { + let mut req = + copy_from_request(&format!("s3://bucket/backup/schema/1/{name}.parquet"), None); + req.connection = connection.clone(); + assert!(matches!( + list_copy_from_paths(&req, &LocalFileAccess::Disabled).await, + Err(error::Error::ListObjects { + source: common_datasource::error::Error::ListObjects { error, .. }, .. + }) if error.kind() == kind + )); + } + server.abort(); + } + fn test_schema_matches(from: (DataType, bool), to: (DataType, bool), matches: bool) { let s1 = Arc::new(Schema::new(vec![Field::new("col", from.0.clone(), from.1)])); let s2 = Arc::new(Schema::new(vec![Field::new("col", to.0.clone(), to.1)])); diff --git a/src/operator/src/statement/ddl.rs b/src/operator/src/statement/ddl.rs index 59695a6e7b3..0e54ba42440 100644 --- a/src/operator/src/statement/ddl.rs +++ b/src/operator/src/statement/ddl.rs @@ -2587,6 +2587,7 @@ fn validate_and_normalize_annotations( CheckError::ColumnNotFound { column } => ColumnNotFoundSnafu { msg: column }.build(), e @ (CheckError::UnknownKey { .. } | CheckError::InvalidValue { .. } + | CheckError::InvalidPartitionNumHint { .. } | CheckError::ColumnNotStringForm { .. }) => InvalidSqlSnafu { err_msg: e.to_string(), } @@ -3525,6 +3526,36 @@ SELECT max(c1), min(c2) FROM schema_2.table_2;"; } } + #[test] + fn test_create_table_with_repartition_partition_num_hint() { + let key = table::requests::REPARTITION_PARTITION_NUM_HINT_KEY; + for (value, expected) in [ + (" 8 ", Some("8")), + ("0", None), + ("-1", None), + ("4294967296", None), + ] { + let expr = create_expr_from_sql(&format!( + "CREATE TABLE metrics (host STRING, ts TIMESTAMP TIME INDEX, PRIMARY KEY(host)) WITH ('{key}' = '{value}')" + )); + let result = create_table_info(&expr, vec![]); + if let Some(expected) = expected { + let info = result.unwrap(); + assert_eq!( + info.meta.options.extra_options.get(key).map(String::as_str), + Some(expected) + ); + } else { + assert!( + result + .unwrap_err() + .to_string() + .contains("expects a positive integer") + ); + } + } + } + #[test] fn test_create_table_with_repartition_column_hint() { let expr = create_expr_from_sql( diff --git a/src/operator/src/statement/set.rs b/src/operator/src/statement/set.rs index b0305dbdaea..83924953e75 100644 --- a/src/operator/src/statement/set.rs +++ b/src/operator/src/statement/set.rs @@ -38,6 +38,35 @@ lazy_static! { static ref PG_TIME_INPUT_REGEX: Regex = Regex::new(r"^(\d+)(ms|s|min|h|d)$").unwrap(); } +/// Sets the session WAL policy for ordinary inserts without changing table options. +pub fn set_skip_wal(exprs: Vec, ctx: QueryContextRef) -> Result<()> { + let [Expr::Value(value)] = exprs.as_slice() else { + return NotSupportedSnafu { + feat: "SET skip_wal requires exactly one boolean value", + } + .fail(); + }; + let skip_wal = match &value.value { + Value::Boolean(value) => *value, + Value::SingleQuotedString(value) | Value::DoubleQuotedString(value) => { + value.parse::().map_err(|_| { + NotSupportedSnafu { + feat: format!("Invalid skip_wal value {value:?}: expected true or false"), + } + .build() + })? + } + _ => { + return NotSupportedSnafu { + feat: "SET skip_wal requires true or false", + } + .fail(); + } + }; + ctx.set_skip_wal(skip_wal); + Ok(()) +} + pub fn set_read_preference(exprs: Vec, ctx: QueryContextRef) -> Result<()> { let read_preference_expr = exprs.first().context(NotSupportedSnafu { feat: "No read preference find in set variable statement", @@ -381,8 +410,58 @@ fn parse_pg_query_timeout_input(input: &str) -> Result { #[cfg(test)] mod test { + use session::Session; + use session::context::{Channel, QueryContext}; + use sql::ast::{Expr, Value}; + + use super::set_skip_wal; use crate::statement::set::parse_pg_query_timeout_input; + #[test] + fn test_set_skip_wal() { + let ctx = QueryContext::arc(); + for value in [true, false] { + set_skip_wal(vec![Expr::Value(Value::Boolean(value).into())], ctx.clone()).unwrap(); + assert_eq!(ctx.skip_wal(), value); + set_skip_wal( + vec![Expr::Value( + Value::SingleQuotedString(value.to_string()).into(), + )], + ctx.clone(), + ) + .unwrap(); + assert_eq!(ctx.skip_wal(), value); + } + ctx.set_skip_wal(true); + for values in [ + vec![], + vec![Expr::Value(Value::Number("1".to_string(), false).into())], + vec![Expr::Value( + Value::SingleQuotedString("invalid".to_string()).into(), + )], + vec![Expr::Value(Value::Boolean(false).into()); 2], + ] { + assert!(set_skip_wal(values, ctx.clone()).is_err()); + assert!(ctx.skip_wal()); + } + } + + #[test] + fn test_set_skip_wal_session_isolation() { + for channel in [Channel::Mysql, Channel::Postgres] { + let session = Session::new(None, channel, Default::default(), 0); + let other = Session::new(None, channel, Default::default(), 1); + assert!(!session.new_query_context().skip_wal()); + set_skip_wal( + vec![Expr::Value(Value::Boolean(true).into())], + session.new_query_context(), + ) + .unwrap(); + assert!(session.new_query_context().skip_wal()); + assert!(!other.new_query_context().skip_wal()); + } + } + #[test] fn test_parse_pg_query_timeout_input() { assert!(parse_pg_query_timeout_input("").is_err()); diff --git a/src/promql/benches/bench_range_fn.rs b/src/promql/benches/bench_range_fn.rs index 551dd0b0bc1..b42b4f60b0f 100644 --- a/src/promql/benches/bench_range_fn.rs +++ b/src/promql/benches/bench_range_fn.rs @@ -37,13 +37,17 @@ use datatypes::arrow::datatypes::{DataType, Field}; use futures::StreamExt; use promql::extension_plan::RangeManipulate; use promql::functions::{ - Changes, Delta, IDelta, Increase, PredictLinear, QuantileOverTime, Rate, Resets, SumOverTime, + AbsentOverTime, Changes, CountOverTime, Delta, DoubleExponentialSmoothing, IDelta, Increase, + LastOverTime, PredictLinear, PresentOverTime, QuantileOverTime, Rate, Resets, SumOverTime, }; use promql::range_array::RangeArray; +/// A `window_step` below `window_size` makes consecutive windows overlap, which is the normal +/// PromQL range query shape. fn build_sliding_ranges( num_points: usize, window_size: u32, + window_step: usize, values: Vec, eval_offset_ms: i64, ) -> (RangeArray, RangeArray, Arc) { @@ -59,10 +63,12 @@ fn build_sliding_ranges( 0 }; - let ranges: Vec<(u32, u32)> = (0..num_windows).map(|i| (i as u32, window_size)).collect(); + let offsets: Vec = (0..num_windows).step_by(window_step).collect(); + let ranges: Vec<(u32, u32)> = offsets.iter().map(|&i| (i as u32, window_size)).collect(); - let eval_ts: Vec = (0..num_windows) - .map(|i| timestamps[i + window_size as usize - 1] + eval_offset_ms) + let eval_ts: Vec = offsets + .iter() + .map(|&i| timestamps[i + window_size as usize - 1] + eval_offset_ms) .collect(); let eval_ts_array = Arc::new(TimestampMillisecondArray::from(eval_ts)); @@ -121,11 +127,12 @@ fn build_changing_values(num_points: usize) -> Vec { fn make_extrapolated_rate_input( num_points: usize, window_size: u32, + window_step: usize, values: Vec, eval_offset_ms: i64, ) -> Vec { let (ts_range, val_range, eval_ts) = - build_sliding_ranges(num_points, window_size, values, eval_offset_ms); + build_sliding_ranges(num_points, window_size, window_step, values, eval_offset_ms); let range_length = window_size as i64 * 1000; vec![ ColumnarValue::Array(Arc::new(ts_range.into_dict())), @@ -192,8 +199,13 @@ fn make_delta_rate_comparison_input( } fn make_idelta_input(num_points: usize, window_size: u32) -> Vec { - let (ts_range, val_range, _) = - build_sliding_ranges(num_points, window_size, build_default_values(num_points), 0); + let (ts_range, val_range, _) = build_sliding_ranges( + num_points, + window_size, + 1, + build_default_values(num_points), + 0, + ); vec![ ColumnarValue::Array(Arc::new(ts_range.into_dict())), ColumnarValue::Array(Arc::new(val_range.into_dict())), @@ -205,7 +217,7 @@ fn make_edge_count_input( window_size: u32, values: Vec, ) -> Vec { - let (ts_range, val_range, _) = build_sliding_ranges(num_points, window_size, values, 0); + let (ts_range, val_range, _) = build_sliding_ranges(num_points, window_size, 1, values, 0); vec![ ColumnarValue::Array(Arc::new(ts_range.into_dict())), ColumnarValue::Array(Arc::new(val_range.into_dict())), @@ -229,8 +241,13 @@ fn make_edge_count_input_with_ranges( } fn make_quantile_input(num_points: usize, window_size: u32) -> Vec { - let (ts_range, val_range, _) = - build_sliding_ranges(num_points, window_size, build_default_values(num_points), 0); + let (ts_range, val_range, _) = build_sliding_ranges( + num_points, + window_size, + 1, + build_default_values(num_points), + 0, + ); vec![ ColumnarValue::Array(Arc::new(ts_range.into_dict())), ColumnarValue::Array(Arc::new(val_range.into_dict())), @@ -239,8 +256,13 @@ fn make_quantile_input(num_points: usize, window_size: u32) -> Vec Vec { - let (ts_range, val_range, _) = - build_sliding_ranges(num_points, window_size, build_default_values(num_points), 0); + let (ts_range, val_range, _) = build_sliding_ranges( + num_points, + window_size, + 1, + build_default_values(num_points), + 0, + ); vec![ ColumnarValue::Array(Arc::new(ts_range.into_dict())), ColumnarValue::Array(Arc::new(val_range.into_dict())), @@ -249,6 +271,26 @@ fn make_predict_linear_input(num_points: usize, window_size: u32) -> Vec Vec { + let (ts_range, val_range, _) = build_sliding_ranges( + num_points, + window_size, + window_step, + build_gauge_values(num_points), + 0, + ); + vec![ + ColumnarValue::Array(Arc::new(ts_range.into_dict())), + ColumnarValue::Array(Arc::new(val_range.into_dict())), + ColumnarValue::Scalar(ScalarValue::Float64(Some(0.5))), + ColumnarValue::Scalar(ScalarValue::Float64(Some(0.1))), + ] +} + struct PreparedUdfCall { args: Vec, arg_fields: Vec>, @@ -338,6 +380,43 @@ fn assert_edge_count_output( assert_eq!(actual, expected); } +fn bench_presence_range_functions(c: &mut Criterion) { + let mut group = c.benchmark_group("presence_range_fn"); + let values = build_default_values(4_096); + let overlapping = PreparedUdfCall::new(make_edge_count_input(4_096, 20, values.clone())); + let low_coverage_ranges = vec![ + (0, 4), + (512, 4), + (1_024, 4), + (1_536, 4), + (2_048, 4), + (2_560, 4), + (3_584, 4), + (4_092, 4), + ]; + let low_coverage = PreparedUdfCall::new(make_edge_count_input_with_ranges( + values, + low_coverage_ranges, + )); + let udfs = [ + ("count_over_time", CountOverTime::scalar_udf()), + ("last_over_time", LastOverTime::scalar_udf()), + ("present_over_time", PresentOverTime::scalar_udf()), + ("absent_over_time", AbsentOverTime::scalar_udf()), + ]; + + for (name, udf) in &udfs { + group.bench_with_input(BenchmarkId::new(*name, "N4096_overlap_w20"), &(), |b, _| { + b.iter(|| invoke_prepared(udf, &overlapping)) + }); + group.bench_with_input(BenchmarkId::new(*name, "N4096_windows8_w4"), &(), |b, _| { + b.iter(|| invoke_prepared(udf, &low_coverage)) + }); + } + + group.finish(); +} + fn bench_range_functions(c: &mut Criterion) { let mut group = c.benchmark_group("range_fn"); @@ -355,6 +434,7 @@ fn bench_range_functions(c: &mut Criterion) { let prepared = PreparedUdfCall::new(make_extrapolated_rate_input( n, w, + 1, build_monotonic_counter_values(n), 500, )); @@ -370,6 +450,7 @@ fn bench_range_functions(c: &mut Criterion) { let prepared = PreparedUdfCall::new(make_extrapolated_rate_input( n, w, + 1, build_resetting_counter_values(n), 500, )); @@ -386,6 +467,7 @@ fn bench_range_functions(c: &mut Criterion) { let prepared = PreparedUdfCall::new(make_extrapolated_rate_input( n, w, + 1, build_monotonic_counter_values(n), 500, )); @@ -401,6 +483,7 @@ fn bench_range_functions(c: &mut Criterion) { let prepared = PreparedUdfCall::new(make_extrapolated_rate_input( n, w, + 1, build_resetting_counter_values(n), 500, )); @@ -417,6 +500,7 @@ fn bench_range_functions(c: &mut Criterion) { let prepared = PreparedUdfCall::new(make_extrapolated_rate_input( n, w, + 1, build_gauge_values(n), 500, )); @@ -471,6 +555,26 @@ fn bench_range_functions(c: &mut Criterion) { ); } + // --- double_exponential_smoothing --- + let smoothing_udf = DoubleExponentialSmoothing::scalar_udf(); + for (window_size, window_step, case) in [ + (4, 1, "N4096_w4_overlap"), + (20, 1, "N4096_w20_overlap"), + (240, 1, "N4096_w240_overlap"), + (240, 240, "N4096_w240_nonoverlap"), + ] { + let prepared = PreparedUdfCall::new(make_double_exponential_smoothing_input( + 4_096, + window_size, + window_step, + )); + group.bench_with_input( + BenchmarkId::new("double_exponential_smoothing", case), + &(), + |b, _| b.iter(|| invoke_prepared(&smoothing_udf, &prepared)), + ); + } + // --- RangeArray: get vs get_offset_length micro-benchmark --- // Isolates the overhead of array slicing vs offset/length lookup for &(n, w) in params { @@ -552,6 +656,43 @@ fn bench_delta_rate_comparison(c: &mut Criterion) { group.finish(); } +/// Counter-reset correction is reduced per window, so its cost follows how much the windows +/// overlap and how many resets each one covers. `range_fn` fixes the query step at one sample +/// and `delta_rate_comparison` at four and twenty, so sweep both dimensions here. +fn bench_rate_window_steps(c: &mut Criterion) { + let mut group = c.benchmark_group("rate_window_steps"); + let rate_udf = Rate::scalar_udf(); + let num_points = 20_280; + let window_size = 120u32; + + // No resets, one reset per 24 window widths, and roughly three resets per window. + for reset_period in [0usize, 2_880, 37] { + let values: Vec = match reset_period { + 0 => (0..num_points).map(|i| i as f64).collect(), + period => (0..num_points).map(|i| (i % period) as f64).collect(), + }; + for window_step in [1usize, 10, 120] { + let prepared = PreparedUdfCall::new(make_extrapolated_rate_input( + num_points, + window_size, + window_step, + values.clone(), + 500, + )); + group.bench_with_input( + BenchmarkId::new( + "rate_counter", + format!("reset{reset_period}_step{window_step}"), + ), + &(), + |b, _| b.iter(|| invoke_prepared(&rate_udf, &prepared)), + ); + } + } + + group.finish(); +} + fn bench_edge_count_functions(c: &mut Criterion) { let mut group = c.benchmark_group("edge_count_fn"); let num_points = 4_096; @@ -840,6 +981,7 @@ fn bench_range_manipulate_wall_time(c: &mut Criterion) { 0, (evaluations as i64 - 1) * RANGE_MANIPULATE_CADENCE_MS, RANGE_MANIPULATE_CADENCE_MS, + 0, window_points as i64 * RANGE_MANIPULATE_CADENCE_MS, "timestamp".to_string(), field_columns, @@ -880,7 +1022,9 @@ fn bench_range_manipulate_wall_time(c: &mut Criterion) { criterion_group!( benches, bench_range_functions, + bench_presence_range_functions, bench_delta_rate_comparison, + bench_rate_window_steps, bench_edge_count_functions, bench_range_manipulate_wall_time ); diff --git a/src/promql/src/extension_plan.rs b/src/promql/src/extension_plan.rs index 29b31b7ca03..750ee41b5e8 100644 --- a/src/promql/src/extension_plan.rs +++ b/src/promql/src/extension_plan.rs @@ -29,9 +29,12 @@ pub use absent::{Absent, AbsentExec, AbsentStream}; use common_query::native_histogram::{SUM_FIELD, native_histogram_value_type}; use common_query::prometheus::is_prometheus_stale_nan; use datafusion::arrow::array::{Array, Float64Array, StructArray}; -use datafusion::arrow::datatypes::{ArrowPrimitiveType, TimestampMillisecondType}; -use datafusion::common::DFSchemaRef; +use datafusion::arrow::datatypes::{ + ArrowPrimitiveType, DataType, TimeUnit, TimestampMillisecondType, +}; +use datafusion::common::{Column, DFSchemaRef}; use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::logical_expr::{Expr, Extension, LogicalPlan}; use datatypes::data_type::DataType as _; pub use empty_metric::{EmptyMetric, EmptyMetricExec, EmptyMetricStream, build_special_time_expr}; pub use histogram_fold::{ @@ -47,6 +50,78 @@ pub use union_distinct_on::{UnionDistinctOn, UnionDistinctOnExec, UnionDistinctO pub type Millisecond = ::Native; +pub(crate) fn timestamp_unit(data_type: &DataType) -> datafusion::error::Result { + match data_type { + DataType::Timestamp(unit, _) => Ok(*unit), + _ => Err(datafusion::error::DataFusionError::Execution( + "Time index column is not a timestamp".into(), + )), + } +} + +pub(crate) fn nanoseconds_per_native_tick(unit: TimeUnit) -> i128 { + match unit { + TimeUnit::Second => 1_000_000_000, + TimeUnit::Millisecond => 1_000_000, + TimeUnit::Microsecond => 1_000, + TimeUnit::Nanosecond => 1, + } +} + +/// Recovers the offset serialized only by an immediately underlying normalize node. +/// +/// This is decode-only recovery for manipulators whose wire messages have no offset field. +/// Follow identity projections (as used by `timestamp()`), but stop at other nodes or changed +/// time columns to avoid applying an inner selector's offset again to an outer subquery. +pub(crate) fn local_offset(plan: &LogicalPlan, time_index: &str) -> Millisecond { + let Some(index) = plan.schema().index_of_column_by_name(None, time_index) else { + return 0; + }; + let (qualifier, field) = plan.schema().qualified_field(index); + let mut time_index = Column::new(qualifier.cloned(), field.name().clone()); + let mut plan = plan; + + loop { + match plan { + LogicalPlan::Extension(Extension { node }) => { + return node + .as_any() + .downcast_ref::() + .and_then(|normalize| normalize.offset_for_time_index(&time_index)) + .unwrap_or_default(); + } + LogicalPlan::Projection(projection) => { + let Some(output_index) = projection.schema.maybe_index_of_column(&time_index) + else { + return 0; + }; + let expr = &projection.expr[output_index]; + let source = match expr { + Expr::Column(column) => column, + Expr::Alias(alias) => { + let Expr::Column(column) = alias.expr.as_ref() else { + return 0; + }; + if alias.name != column.name { + return 0; + } + column + } + _ => return 0, + }; + let Some(input_index) = projection.input.schema().maybe_index_of_column(source) + else { + return 0; + }; + let (qualifier, field) = projection.input.schema().qualified_field(input_index); + time_index = Column::new(qualifier.cloned(), field.name().clone()); + plan = projection.input.as_ref(); + } + _ => return 0, + } + } +} + const METRIC_NUM_SERIES: &str = "num_series"; fn prometheus_stale_sample_column(column: &dyn Array) -> Option<(&dyn Array, &Float64Array)> { @@ -109,3 +184,87 @@ pub fn resolve_column_names( .map(|idx| resolve_column_name(*idx, schema, context, column_type)) .collect() } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit}; + use datafusion::common::ToDFSchema; + use datafusion::logical_expr::{EmptyRelation, Extension, LogicalPlan, Projection}; + use datafusion_expr::col; + + use super::*; + + fn input() -> LogicalPlan { + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::new(Schema::new(vec![ + Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new( + "other_ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])) + .to_dfschema_ref() + .unwrap(), + }) + } + + fn normalized() -> LogicalPlan { + LogicalPlan::Extension(Extension { + node: Arc::new(SeriesNormalize::new( + 1_000, + "timestamp", + false, + Vec::new(), + input(), + )), + }) + } + + #[test] + fn local_offset_tracks_identity_preserving_projections() { + let projection = + Projection::try_new(vec![col("timestamp"), col("value")], Arc::new(normalized())) + .unwrap(); + let projection = Projection::try_new( + vec![col("timestamp").alias("timestamp"), col("value")], + Arc::new(LogicalPlan::Projection(projection)), + ) + .unwrap(); + + assert_eq!( + 1_000, + local_offset(&LogicalPlan::Projection(projection), "timestamp") + ); + } + + #[test] + fn local_offset_rejects_a_different_timestamp_or_manipulator() { + let renamed = Projection::try_new( + vec![col("other_ts").alias("timestamp"), col("value")], + Arc::new(normalized()), + ) + .unwrap(); + assert_eq!( + 0, + local_offset(&LogicalPlan::Projection(renamed), "timestamp") + ); + + let divide = LogicalPlan::Extension(Extension { + node: Arc::new(SeriesDivide::new( + Vec::new(), + "timestamp".to_string(), + normalized(), + )), + }); + assert_eq!(0, local_offset(÷, "timestamp")); + } +} diff --git a/src/promql/src/extension_plan/instant_manipulate.rs b/src/promql/src/extension_plan/instant_manipulate.rs index 25d62e6bee3..8619100b4dd 100644 --- a/src/promql/src/extension_plan/instant_manipulate.rs +++ b/src/promql/src/extension_plan/instant_manipulate.rs @@ -13,7 +13,6 @@ // limitations under the License. use std::any::Any; -use std::cmp::Ordering; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -29,6 +28,7 @@ use datafusion::execution::context::TaskContext; use datafusion::logical_expr::{ EmptyRelation, Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore, }; +use datafusion::physical_expr::EquivalenceProperties; use datafusion::physical_plan::metrics::{ BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricValue, MetricsSet, }; @@ -38,6 +38,7 @@ use datafusion::physical_plan::{ }; use datafusion_expr::col; use datatypes::arrow::compute; +use datatypes::timestamp::timestamp_array_to_primitive; use futures::{Stream, StreamExt, ready}; use greptime_proto::substrait_extension as pb; use prost::Message; @@ -46,8 +47,9 @@ use snafu::ResultExt; use crate::error::{DeserializeSnafu, Result}; use crate::extension_plan::series_divide::SeriesDivide; use crate::extension_plan::{ - METRIC_NUM_SERIES, Millisecond, is_prometheus_stale_sample, prometheus_stale_sample_column, - resolve_column_name, serialize_column_index, + METRIC_NUM_SERIES, Millisecond, is_prometheus_stale_sample, local_offset, + nanoseconds_per_native_tick, prometheus_stale_sample_column, resolve_column_name, + serialize_column_index, timestamp_unit, }; use crate::metrics::PROMQL_SERIES_COUNT; @@ -67,21 +69,52 @@ fn mixed_sample_fields(field: Option<&str>) -> [Option<&str>; 2] { /// This plan will try to align the input time series, for every timestamp between /// `start` and `end` with step `interval`. Find in the `lookback` range if data /// is missing at the given timestamp. -#[derive(Debug, PartialEq, Eq, Hash, PartialOrd)] +#[derive(Debug, PartialEq, Eq, Hash)] pub struct InstantManipulate { start: Millisecond, end: Millisecond, lookback_delta: Millisecond, interval: Millisecond, + offset: Millisecond, time_index_column: String, // Planner-provided tag-column hint for execution fast paths. tag_columns: Vec, /// Primary sample column used to derive the columns checked for staleness. field_column: Option, input: LogicalPlan, + output_schema: DFSchemaRef, unfix: Option, } +impl PartialOrd for InstantManipulate { + fn partial_cmp(&self, other: &Self) -> Option { + ( + self.start, + self.end, + self.lookback_delta, + self.interval, + self.offset, + &self.time_index_column, + &self.tag_columns, + &self.field_column, + &self.input, + &self.unfix, + ) + .partial_cmp(&( + other.start, + other.end, + other.lookback_delta, + other.interval, + other.offset, + &other.time_index_column, + &other.tag_columns, + &other.field_column, + &other.input, + &other.unfix, + )) + } +} + #[derive(Debug, PartialEq, Eq, Hash, PartialOrd)] struct UnfixIndices { pub time_index_idx: u64, @@ -98,7 +131,7 @@ impl UserDefinedLogicalNodeCore for InstantManipulate { } fn schema(&self) -> &DFSchemaRef { - self.input.schema() + &self.output_schema } fn expressions(&self) -> Vec { @@ -180,6 +213,8 @@ impl UserDefinedLogicalNodeCore for InstantManipulate { end: self.end, lookback_delta: self.lookback_delta, interval: self.interval, + offset: local_offset(&input, &time_index_column), + output_schema: Self::calculate_output_schema(&input, &time_index_column)?, time_index_column, tag_columns: Self::resolve_tag_columns(&input, &self.tag_columns), field_column, @@ -192,9 +227,11 @@ impl UserDefinedLogicalNodeCore for InstantManipulate { end: self.end, lookback_delta: self.lookback_delta, interval: self.interval, + offset: self.offset, time_index_column: self.time_index_column.clone(), tag_columns: Self::resolve_tag_columns(&input, &self.tag_columns), field_column: self.field_column.clone(), + output_schema: Self::calculate_output_schema(&input, &self.time_index_column)?, input, unfix: None, }) @@ -203,12 +240,45 @@ impl UserDefinedLogicalNodeCore for InstantManipulate { } impl InstantManipulate { + fn calculate_output_schema( + input: &LogicalPlan, + time_index_column: &str, + ) -> DataFusionResult { + let input_schema = input.schema(); + let time_index = input_schema + .index_of_column_by_name(None, time_index_column) + .ok_or_else(|| { + DataFusionError::Internal(format!( + "InstantManipulate time index {time_index_column} not found" + )) + })?; + let mut fields = (0..input_schema.fields().len()) + .map(|index| { + let (qualifier, field) = input_schema.qualified_field(index); + (qualifier.cloned(), field.clone()) + }) + .collect::>(); + let (qualifier, field) = input_schema.qualified_field(time_index); + fields[time_index] = ( + qualifier.cloned(), + Arc::new(field.as_ref().clone().with_data_type(DataType::Timestamp( + datafusion::arrow::datatypes::TimeUnit::Millisecond, + None, + ))), + ); + Ok(Arc::new(DFSchema::new_with_metadata( + fields, + input_schema.metadata().clone(), + )?)) + } + #[allow(clippy::too_many_arguments)] pub fn new( start: Millisecond, end: Millisecond, lookback_delta: Millisecond, interval: Millisecond, + offset: Millisecond, time_index_column: String, tag_columns: Vec, field_column: Option, @@ -219,6 +289,9 @@ impl InstantManipulate { end, lookback_delta, interval, + offset, + output_schema: Self::calculate_output_schema(&input, &time_index_column) + .unwrap_or_else(|_| input.schema().clone()), time_index_column, tag_columns, field_column, @@ -274,7 +347,27 @@ impl InstantManipulate { pub fn to_execution_plan(&self, exec_input: Arc) -> Arc { let reuse_tsid_column = matches!(self.tag_columns.as_slice(), [tag] if tag == "__tsid"); + let mut fields = exec_input.schema().fields().to_vec(); + let time_index = exec_input + .schema() + .index_of(&self.time_index_column) + .expect("time index column not found"); + fields[time_index] = Arc::new(fields[time_index].as_ref().clone().with_data_type( + DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Millisecond, None), + )); + let output_schema = Arc::new(datafusion::arrow::datatypes::Schema::new_with_metadata( + fields, + exec_input.schema().metadata().clone(), + )); + let input_properties = exec_input.properties(); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(output_schema.clone()), + input_properties.partitioning.clone(), + input_properties.emission_type, + input_properties.boundedness, + )); Arc::new(InstantManipulateExec { + offset: self.offset, start: self.start, end: self.end, lookback_delta: self.lookback_delta, @@ -283,6 +376,8 @@ impl InstantManipulate { field_column: self.field_column.clone(), reuse_tsid_column, input: exec_input, + output_schema, + properties, metric: ExecutionPlanMetricsSet::new(), }) } @@ -311,9 +406,10 @@ impl InstantManipulate { pub fn deserialize(bytes: &[u8]) -> Result { let pb_instant_manipulate = pb::InstantManipulate::decode(bytes).context(DeserializeSnafu)?; + let empty_schema = Arc::new(DFSchema::empty()); let placeholder_plan = LogicalPlan::EmptyRelation(EmptyRelation { produce_one_row: false, - schema: Arc::new(DFSchema::empty()), + schema: empty_schema.clone(), }); let unfix = UnfixIndices { @@ -326,9 +422,11 @@ impl InstantManipulate { end: pb_instant_manipulate.end, lookback_delta: pb_instant_manipulate.lookback_delta, interval: pb_instant_manipulate.interval, + offset: 0, time_index_column: String::new(), tag_columns: Vec::new(), field_column: None, + output_schema: empty_schema, input: placeholder_plan, unfix: Some(unfix), }) @@ -337,6 +435,7 @@ impl InstantManipulate { #[derive(Debug)] pub struct InstantManipulateExec { + offset: Millisecond, start: Millisecond, end: Millisecond, lookback_delta: Millisecond, @@ -346,6 +445,8 @@ pub struct InstantManipulateExec { reuse_tsid_column: bool, input: Arc, + output_schema: SchemaRef, + properties: Arc, metric: ExecutionPlanMetricsSet, } @@ -355,11 +456,11 @@ impl ExecutionPlan for InstantManipulateExec { } fn schema(&self) -> SchemaRef { - self.input.schema() + self.output_schema.clone() } fn properties(&self) -> &Arc { - self.input.properties() + &self.properties } fn required_input_distribution(&self) -> Vec { @@ -380,7 +481,18 @@ impl ExecutionPlan for InstantManipulateExec { children: Vec>, ) -> DataFusionResult> { assert!(!children.is_empty()); + let input = children[0].clone(); + // The child may expose native timestamps, but our output schema uses ms. + // Retain its execution properties, not its schema equivalences or ordering. + let input_properties = input.properties(); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(self.output_schema.clone()), + input_properties.partitioning.clone(), + input_properties.emission_type, + input_properties.boundedness, + )); Ok(Arc::new(Self { + offset: self.offset, start: self.start, end: self.end, lookback_delta: self.lookback_delta, @@ -388,7 +500,9 @@ impl ExecutionPlan for InstantManipulateExec { time_index_column: self.time_index_column.clone(), field_column: self.field_column.clone(), reuse_tsid_column: self.reuse_tsid_column, - input: children[0].clone(), + input, + output_schema: self.output_schema.clone(), + properties, metric: self.metric.clone(), })) } @@ -413,6 +527,7 @@ impl ExecutionPlan for InstantManipulateExec { .column_with_name(&self.time_index_column) .expect("time index column not found") .0; + let time_unit = timestamp_unit(schema.field(time_index).data_type())?; let field_indices = mixed_sample_fields(self.field_column.as_deref()).map(|field| { field.and_then(|field| schema.column_with_name(field).map(|(index, _)| index)) }); @@ -421,15 +536,17 @@ impl ExecutionPlan for InstantManipulateExec { .filter(|(_, field)| field.data_type() == &DataType::UInt64) .map(|(index, _)| index); Ok(Box::pin(InstantManipulateStream { + offset: self.offset, start: self.start, end: self.end, lookback_delta: self.lookback_delta, interval: self.interval, time_index, + time_unit, field_indices, tsid_index, reuse_tsid_column: self.reuse_tsid_column && tsid_index.is_some(), - schema, + schema: self.output_schema.clone(), input, metric: baseline_metric, num_series, @@ -487,12 +604,14 @@ impl DisplayAs for InstantManipulateExec { } pub struct InstantManipulateStream { + offset: Millisecond, start: Millisecond, end: Millisecond, lookback_delta: Millisecond, interval: Millisecond, // Column index of TIME INDEX column's position in schema time_index: usize, + time_unit: datafusion::arrow::datatypes::TimeUnit, field_indices: [Option; 2], tsid_index: Option, reuse_tsid_column: bool, @@ -516,11 +635,10 @@ impl Stream for InstantManipulateStream { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let poll = match ready!(self.input.poll_next_unpin(cx)) { Some(Ok(batch)) => { - if batch.num_rows() == 0 { - return Poll::Pending; - } let timer = std::time::Instant::now(); - self.num_series.add(1); + if batch.num_rows() != 0 { + self.num_series.add(1); + } let result = Ok(batch).and_then(|batch| self.manipulate(batch)); self.metric.elapsed_compute().add_elapsed(timer); Poll::Ready(Some(result)) @@ -543,22 +661,12 @@ impl InstantManipulateStream { /// lookback window `(eval_ts - lookback_delta, eval_ts]`; a sample at exactly /// `eval_ts - lookback_delta` is too old. pub fn manipulate(&self, input: RecordBatch) -> DataFusionResult { - let ts_column = input - .column(self.time_index) - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Execution( - "Time index Column downcast to TimestampMillisecondArray failed".into(), - ) - })?; - - // Early return for empty input + let ts_column = input.column(self.time_index); if ts_column.is_empty() { - return Ok(input); + // Returning Pending after consuming a ready batch without a wake can stall the stream. + return Ok(RecordBatch::new_empty(self.schema.clone())); } - - // Field columns for staleness checks, classified once per batch. + let scale = nanoseconds_per_native_tick(self.time_unit); let stale_sample_columns = self.field_indices.map(|index| { index.and_then(|index| prometheus_stale_sample_column(input.column(index).as_ref())) }); @@ -568,101 +676,82 @@ impl InstantManipulateStream { .flatten() .any(|column| is_prometheus_stale_sample(*column, row)) }; - - // Optimize iteration range based on actual data bounds - let first_ts = ts_column.value(0); - let last_ts = ts_column.value(ts_column.len() - 1); - // A sample at `t` is eligible for eval time `eval_ts` iff: - // t > eval_ts - lookback_delta <=> eval_ts < t + lookback_delta. - // Therefore the last eval timestamp for which the last sample is still eligible is: - // last_ts + lookback_delta - 1 (millisecond granularity). - let last_useful = if self.lookback_delta > 0 { - last_ts + self.lookback_delta - 1 + let (timestamps, _) = timestamp_array_to_primitive(ts_column).ok_or_else(|| { + DataFusionError::Execution("Time index column is not a timestamp".into()) + })?; + let timestamps = timestamps.values(); + let len = timestamps.len(); + // Shift the native-tick timeline in i128 before comparing samples. Doing + // this in the Arrow storage unit can overflow even when the shifted + // PromQL millisecond evaluation time is representable. + let to_nanoseconds = + |timestamp: i64| (timestamp as i128) * scale + (self.offset as i128) * 1_000_000; + let first_ns = to_nanoseconds(timestamps[0]); + let last_ns = to_nanoseconds(timestamps[len - 1]); + // An exact sample remains useful with zero lookback. Otherwise the lower + // boundary is exclusive, so subtract one nanosecond from its final window. + let last_useful = if self.lookback_delta == 0 { + last_ns } else { - last_ts + last_ns + (self.lookback_delta as i128) * 1_000_000 - 1 + }; + let first_ms = (first_ns + 999_999).div_euclid(1_000_000); + let last_ms = last_useful.div_euclid(1_000_000); + let query_start = self.start as i128; + let query_end = self.end as i128; + let interval = self.interval as i128; + let max_start = first_ms.max(query_start); + let min_end = last_ms.min(query_end); + let (aligned_start, aligned_end) = if max_start > min_end { + (1, 0) + } else { + ( + query_start + (max_start - query_start) / interval * interval, + query_end - (query_end - min_end) / interval * interval, + ) }; - - let max_start = first_ts.max(self.start); - let min_end = last_useful.min(self.end); - - let aligned_start = self.start + (max_start - self.start) / self.interval * self.interval; - let aligned_end = self.end - (self.end - min_end) / self.interval * self.interval; - let estimated_points = if aligned_end >= aligned_start { - ((aligned_end - aligned_start) / self.interval).saturating_add(1) as usize + (aligned_end - aligned_start) / interval + 1 } else { 0 }; - if estimated_points > MAX_INSTANT_MANIPULATE_OUTPUT_POINTS { + if estimated_points > MAX_INSTANT_MANIPULATE_OUTPUT_POINTS as i128 { return Err(DataFusionError::Execution(format!( "InstantManipulate output points exceed limit: {estimated_points} > {MAX_INSTANT_MANIPULATE_OUTPUT_POINTS}" ))); } + let estimated_points = estimated_points as usize; + let aligned_start = aligned_start as i64; + let aligned_end = aligned_end as i64; let mut take_indices = Vec::with_capacity(estimated_points); - - let mut cursor = 0; - - let aligned_ts_iter = (aligned_start..=aligned_end).step_by(self.interval as usize); let mut aligned_ts = Vec::with_capacity(estimated_points); - - // calculate the offsets to take - 'next: for expected_ts in aligned_ts_iter { - // first, search toward end to see if there is matched timestamp - while cursor < ts_column.len() { - let curr = ts_column.value(cursor); - match curr.cmp(&expected_ts) { - Ordering::Equal => { - if is_stale(cursor) { - // Ignore the stale marker. - } else { - take_indices.push(cursor as u64); - aligned_ts.push(expected_ts); - } - continue 'next; - } - Ordering::Greater => break, - Ordering::Less => {} + let mut cursor = 0; + for expected_ms in (aligned_start..=aligned_end).step_by(self.interval as usize) { + let expected = (expected_ms as i128) * 1_000_000; + let mut exact_candidate = None; + while cursor < len && to_nanoseconds(timestamps[cursor]) <= expected { + if to_nanoseconds(timestamps[cursor]) == expected && exact_candidate.is_none() { + exact_candidate = Some(cursor); } cursor += 1; } - if cursor == ts_column.len() { - cursor -= 1; - // short cut this loop - if ts_column.value(cursor) + self.lookback_delta <= expected_ts { - break; - } - } - - // then examine the value - let curr_ts = ts_column.value(cursor); - if curr_ts + self.lookback_delta <= expected_ts { + // Keep the first row among exact timestamp ties; otherwise use the + // latest preceding row. Zero lookback admits only exact samples. + // Test staleness after choosing: a selected stale marker suppresses + // this evaluation rather than falling back to an older finite value. + let Some(candidate) = exact_candidate.or_else(|| cursor.checked_sub(1)) else { continue; - } - if curr_ts > expected_ts { - // exceeds current expected timestamp, examine the previous value - if let Some(prev_cursor) = cursor.checked_sub(1) { - let prev_ts = ts_column.value(prev_cursor); - if prev_ts + self.lookback_delta > expected_ts { - // only use the point in the time range - if is_stale(prev_cursor) { - // Do not use a stale marker as the newest value. - continue; - } - // use this point - take_indices.push(prev_cursor as u64); - aligned_ts.push(expected_ts); - } - } - } else if is_stale(cursor) { - // Do not use a stale marker as the newest value. - } else { - // use this point - take_indices.push(cursor as u64); - aligned_ts.push(expected_ts); + }; + let candidate_ts = to_nanoseconds(timestamps[candidate]); + let lower = expected - (self.lookback_delta as i128) * 1_000_000; + if (candidate_ts == expected || candidate_ts > lower) + && candidate_ts <= expected + && !is_stale(candidate) + { + take_indices.push(candidate as u64); + aligned_ts.push(expected_ms); } } - - // take record batch and replace the time index column self.take_record_batch_optional(input, take_indices, aligned_ts) } @@ -696,7 +785,7 @@ impl InstantManipulateStream { arrays.push(compute::take(array, indices_array, None)?); } - let result = RecordBatch::try_new(record_batch.schema(), arrays) + let result = RecordBatch::try_new(self.schema.clone(), arrays) .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; Ok(result) } @@ -718,14 +807,19 @@ fn reuse_constant_column(array: &Arc, len: usize) -> DataFusionResult mod test { use common_query::native_histogram::build_histogram_array; use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; - use datafusion::arrow::array::Float64Array; + use datafusion::arrow::array::{ + Float64Array, TimestampMicrosecondArray, TimestampNanosecondArray, TimestampSecondArray, + }; use datafusion::arrow::buffer::NullBuffer; - use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit}; use datafusion::common::ToDFSchema; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; - use datafusion::logical_expr::{EmptyRelation, LogicalPlan}; + use datafusion::logical_expr::{ + EmptyRelation, Extension, LogicalPlan, Projection, UserDefinedLogicalNodeCore, + }; use datafusion::prelude::SessionContext; + use datafusion_expr::col; use super::*; use crate::extension_plan::test_util::{ @@ -746,6 +840,7 @@ mod test { Arc::new(prepare_test_data()) }; let normalize_exec = Arc::new(InstantManipulateExec { + offset: 0, start, end, lookback_delta, @@ -753,6 +848,8 @@ mod test { time_index_column: TIME_INDEX_COLUMN.to_string(), field_column: Some("value".to_string()), reuse_tsid_column: false, + output_schema: memory_exec.schema(), + properties: memory_exec.properties().clone(), input: memory_exec, metric: ExecutionPlanMetricsSet::new(), }); @@ -767,6 +864,319 @@ mod test { assert_eq!(result_literal, expected); } + #[tokio::test] + async fn native_timestamps_select_exact_samples_and_keep_ms_output() { + for (unit, ticks_per_ms) in [ + (TimeUnit::Microsecond, 1_000_i64), + (TimeUnit::Nanosecond, 1_000_000_i64), + ] { + let lower = 1_000 * ticks_per_ms; + let upper = 1_001 * ticks_per_ms; + let stale = f64::from_bits(PROMETHEUS_STALE_NAN_BITS); + for (name, timestamps, values, expected_timestamps, expected_values) in [ + ( + "exact upper sample", + vec![lower + 1, upper], + vec![1.0, 2.0], + vec![1_001], + vec![2.0], + ), + ( + "exclusive lower boundary and future sample", + vec![lower, upper + 1], + vec![1.0, 2.0], + vec![1_000], + vec![1.0], + ), + ( + "one native tick above lower boundary", + vec![lower + 1, upper + 1], + vec![1.0, 2.0], + vec![1_001], + vec![1.0], + ), + ( + "future stale marker does not suppress", + vec![lower + 1, upper + 1], + vec![1.0, stale], + vec![1_001], + vec![1.0], + ), + ( + "latest in-window stale marker suppresses", + vec![lower + 1, lower + 2, upper + 1], + vec![1.0, stale, 3.0], + vec![], + vec![], + ), + ] { + let schema = Arc::new(Schema::new(vec![ + Field::new(TIME_INDEX_COLUMN, DataType::Timestamp(unit, None), false), + Field::new("value", DataType::Float64, true), + ])); + let time: Arc = match unit { + TimeUnit::Microsecond => Arc::new(TimestampMicrosecondArray::from(timestamps)), + TimeUnit::Nanosecond => Arc::new(TimestampNanosecondArray::from(timestamps)), + _ => unreachable!(), + }; + let batch = RecordBatch::try_new( + schema.clone(), + vec![time, Arc::new(Float64Array::from(values))], + ) + .unwrap(); + let logical_input = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: schema.clone().to_dfschema_ref().unwrap(), + }); + let plan = InstantManipulate::new( + 1_000, + 1_001, + 1, + 1, + 0, + TIME_INDEX_COLUMN.to_string(), + Vec::new(), + Some("value".to_string()), + logical_input.clone(), + ); + let output_schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])); + assert_eq!(plan.schema().as_arrow(), output_schema.as_ref()); + + let rebuilt = InstantManipulate::deserialize(&plan.serialize()) + .unwrap() + .with_exprs_and_inputs(vec![], vec![logical_input]) + .unwrap(); + assert_eq!(rebuilt.schema(), plan.schema()); + assert_eq!(rebuilt.input.schema().as_arrow(), schema.as_ref()); + + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema.clone(), None).unwrap(), + ))); + let exec = rebuilt.to_execution_plan(input); + assert_eq!(exec.schema(), output_schema); + assert_eq!(exec.children()[0].schema(), schema); + + let batches = + datafusion::physical_plan::collect(exec, SessionContext::default().task_ctx()) + .await + .unwrap(); + assert_eq!(batches.len(), 1, "{unit:?}: {name}"); + let output = &batches[0]; + assert_eq!(output.schema(), output_schema); + let timestamps = output + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let values = output + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + timestamps.values().as_ref(), + expected_timestamps.as_slice(), + "{unit:?}: {name}" + ); + assert_eq!( + values.values().as_ref(), + expected_values.as_slice(), + "{unit:?}: {name}" + ); + assert_eq!(values.null_count(), 0, "{unit:?}: {name}"); + } + } + } + + #[tokio::test] + async fn logical_normalize_offset_survives_rebuild_and_executes() { + for (name, time_unit, raw, offset, start, lookback_delta) in [ + ( + "millisecond offset", + TimeUnit::Millisecond, + 0, + 1_000, + 1_000, + 0, + ), + ( + "second timestamp with negative fractional offset", + TimeUnit::Second, + 1, + -500, + 500, + 0, + ), + ] { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(time_unit, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])); + let input = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: schema.clone().to_dfschema_ref().unwrap(), + }); + let normalize = crate::extension_plan::SeriesNormalize::new( + offset, + TIME_INDEX_COLUMN, + false, + Vec::new(), + input.clone(), + ); + let normalize = + crate::extension_plan::SeriesNormalize::deserialize(&normalize.serialize()) + .unwrap() + .with_exprs_and_inputs(vec![], vec![input.clone()]) + .unwrap(); + let normalized = LogicalPlan::Projection( + Projection::try_new( + vec![col(TIME_INDEX_COLUMN), col("value")], + Arc::new(LogicalPlan::Extension(Extension { + node: Arc::new(normalize), + })), + ) + .unwrap(), + ); + let fresh = InstantManipulate::new( + start, + start, + lookback_delta, + 1, + offset, + TIME_INDEX_COLUMN.to_string(), + Vec::new(), + Some("value".to_string()), + input.clone(), + ) + .with_exprs_and_inputs(vec![], vec![input.clone()]) + .unwrap(); + let serialized = InstantManipulate::new( + start, + start, + lookback_delta, + 1, + offset, + TIME_INDEX_COLUMN.to_string(), + Vec::new(), + Some("value".to_string()), + normalized.clone(), + ); + let decoded = InstantManipulate::deserialize(&serialized.serialize()) + .unwrap() + .with_exprs_and_inputs(vec![], vec![normalized]) + .unwrap(); + let timestamp: Arc = match time_unit { + TimeUnit::Millisecond => Arc::new(TimestampMillisecondArray::from(vec![raw])), + TimeUnit::Second => Arc::new(TimestampSecondArray::from(vec![raw])), + _ => unreachable!(), + }; + let batch = RecordBatch::try_new( + schema.clone(), + vec![timestamp, Arc::new(Float64Array::from(vec![7.0]))], + ) + .unwrap(); + for (mode, rebuilt) in [("fresh", fresh), ("decoded", decoded)] { + let rebuilt = rebuilt + .with_exprs_and_inputs(vec![], vec![input.clone()]) + .unwrap(); + assert_eq!(rebuilt.offset, offset, "{name}: {mode}"); + assert_eq!(rebuilt.input.schema(), input.schema(), "{name}: {mode}"); + + let empty_exec_input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![]], schema.clone(), None).unwrap(), + ))); + let exec_input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch.clone()]], schema.clone(), None) + .unwrap(), + ))); + let exec = rebuilt + .to_execution_plan(empty_exec_input) + .with_new_children(vec![exec_input]) + .unwrap(); + let output = + datafusion::physical_plan::collect(exec, SessionContext::default().task_ctx()) + .await + .unwrap(); + let output = &output[0]; + assert_eq!(output.num_rows(), 1, "{name}: {mode}"); + assert_eq!( + output + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + start, + "{name}: {mode}" + ); + assert_eq!( + output + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 7.0, + "{name}: {mode}" + ); + } + } + } + + #[test] + fn deserialized_ordering_preserves_column_indices() { + let mut wire = pb::InstantManipulate::default(); + let first = InstantManipulate::deserialize(&wire.encode_to_vec()).unwrap(); + wire.time_index_idx = 1; + let second = InstantManipulate::deserialize(&wire.encode_to_vec()).unwrap(); + assert_ne!(first, second); + assert_eq!(first.partial_cmp(&second), Some(std::cmp::Ordering::Less)); + wire.field_index_idx = 2; + let third = InstantManipulate::deserialize(&wire.encode_to_vec()).unwrap(); + assert_ne!(second, third); + assert_eq!(second.partial_cmp(&third), Some(std::cmp::Ordering::Less)); + + let input = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: prepare_test_data().schema().to_dfschema_ref().unwrap(), + }); + let first = InstantManipulate::new( + 0, + 0, + 0, + 0, + 0, + TIME_INDEX_COLUMN.to_string(), + Vec::new(), + Some("value".to_string()), + input.clone(), + ); + let second = InstantManipulate::new( + 0, + 0, + 0, + 0, + 1, + TIME_INDEX_COLUMN.to_string(), + Vec::new(), + Some("value".to_string()), + input, + ); + assert_ne!(first, second); + assert_eq!(first.partial_cmp(&second), Some(std::cmp::Ordering::Less)); + } + #[test] fn pruning_should_keep_time_and_field_columns_for_exec() { let df_schema = prepare_test_data().schema().to_dfschema_ref().unwrap(); @@ -779,6 +1189,7 @@ mod test { 0, 0, 0, + 0, TIME_INDEX_COLUMN.to_string(), Vec::new(), Some("value".to_string()), @@ -811,6 +1222,7 @@ mod test { 0, 0, 0, + 0, TIME_INDEX_COLUMN.to_string(), vec!["__tsid".to_string()], Some("value".to_string()), @@ -853,6 +1265,7 @@ mod test { 0, 0, 0, + 0, TIME_INDEX_COLUMN.to_string(), vec!["__tsid".to_string()], Some("value".to_string()), @@ -886,6 +1299,7 @@ mod test { 0, 0, 0, + 0, TIME_INDEX_COLUMN.to_string(), vec!["__tsid".to_string()], Some("value".to_string()), @@ -927,6 +1341,7 @@ mod test { MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), ))); let normalize_exec = Arc::new(InstantManipulateExec { + offset: 0, start: 0, end: 1_500, lookback_delta: 1_000, @@ -934,6 +1349,8 @@ mod test { time_index_column: TIME_INDEX_COLUMN.to_string(), field_column: Some("value".to_string()), reuse_tsid_column: true, + output_schema: input.schema(), + properties: input.properties().clone(), input, metric: ExecutionPlanMetricsSet::new(), }); @@ -988,6 +1405,7 @@ mod test { MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), ))); let normalize_exec = Arc::new(InstantManipulateExec { + offset: 0, start: 0, end: 1_500, lookback_delta: 1_000, @@ -995,6 +1413,8 @@ mod test { time_index_column: TIME_INDEX_COLUMN.to_string(), field_column: Some("value".to_string()), reuse_tsid_column: true, + output_schema: input.schema(), + properties: input.properties().clone(), input, metric: ExecutionPlanMetricsSet::new(), }); @@ -1042,6 +1462,7 @@ mod test { ))); let too_many_points = MAX_INSTANT_MANIPULATE_OUTPUT_POINTS as Millisecond + 1; let normalize_exec = Arc::new(InstantManipulateExec { + offset: 0, start: 0, end: too_many_points, lookback_delta: too_many_points + 1, @@ -1049,6 +1470,8 @@ mod test { time_index_column: TIME_INDEX_COLUMN.to_string(), field_column: Some("value".to_string()), reuse_tsid_column: false, + output_schema: input.schema(), + properties: input.properties().clone(), input, metric: ExecutionPlanMetricsSet::new(), }); @@ -1336,6 +1759,290 @@ mod test { .await; } + #[test] + fn exact_ties_select_first_and_lookback_uses_latest() { + for (values, expected_timestamp) in [ + (vec![42.0, f64::from_bits(PROMETHEUS_STALE_NAN_BITS)], 1_000), + (vec![f64::from_bits(PROMETHEUS_STALE_NAN_BITS), 42.0], 1_050), + ] { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])); + let input = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_000, 1_000])), + Arc::new(Float64Array::from(values)), + ], + ) + .unwrap(); + let stream = InstantManipulateStream { + offset: 0, + start: 1_000, + end: 1_050, + lookback_delta: 100, + interval: 50, + time_index: 0, + time_unit: TimeUnit::Millisecond, + field_indices: [Some(1), None], + tsid_index: None, + reuse_tsid_column: false, + schema: schema.clone(), + input: Box::pin( + datafusion::physical_plan::memory::MemoryStream::try_new(vec![], schema, None) + .unwrap(), + ), + metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + num_series: Count::new(), + }; + + let output = stream.manipulate(input).unwrap(); + let timestamps = output + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let values = output + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(timestamps.values(), &[expected_timestamp]); + assert_eq!(values.values(), &[42.0]); + } + } + + #[test] + fn empty_batches_preserve_stream_progression_and_output_schema() { + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + use futures::stream; + use futures::task::noop_waker_ref; + + let input_schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Second, None), + false, + )])); + let output_schema = Arc::new(Schema::new(vec![Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + )])); + let empty = RecordBatch::new_empty(input_schema.clone()); + let valid = RecordBatch::try_new( + input_schema.clone(), + vec![Arc::new(TimestampSecondArray::from(vec![1]))], + ) + .unwrap(); + let input = RecordBatchStreamAdapter::new( + input_schema, + stream::iter(vec![ + Ok(empty.clone()), + Ok(empty.clone()), + Ok(valid), + Ok(empty), + Err(DataFusionError::Execution("injected input error".into())), + ]), + ); + let mut stream = InstantManipulateStream { + offset: 0, + start: 1_000, + end: 1_000, + lookback_delta: 0, + interval: 1, + time_index: 0, + time_unit: TimeUnit::Second, + field_indices: [None, None], + tsid_index: None, + reuse_tsid_column: false, + schema: output_schema.clone(), + input: Box::pin(input), + metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + num_series: Count::new(), + }; + let waker = noop_waker_ref(); + let mut cx = Context::from_waker(waker); + + for _ in 0..2 { + let Poll::Ready(Some(Ok(batch))) = Pin::new(&mut stream).poll_next(&mut cx) else { + panic!("empty batch must be returned immediately"); + }; + assert_eq!(batch.num_rows(), 0); + assert_eq!(batch.schema(), output_schema); + } + + let Poll::Ready(Some(Ok(valid))) = Pin::new(&mut stream).poll_next(&mut cx) else { + panic!("valid batch must follow empty batches"); + }; + assert_eq!(valid.schema(), output_schema); + assert_eq!( + valid.column(0).data_type(), + &DataType::Timestamp(TimeUnit::Millisecond, None) + ); + assert_eq!( + valid + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values(), + &[1_000] + ); + + let Poll::Ready(Some(Ok(empty))) = Pin::new(&mut stream).poll_next(&mut cx) else { + panic!("empty batch after valid batch must be returned immediately"); + }; + assert_eq!(empty.num_rows(), 0); + assert_eq!(empty.schema(), output_schema); + + let Poll::Ready(Some(Err(error))) = Pin::new(&mut stream).poll_next(&mut cx) else { + panic!("input error must propagate"); + }; + assert!(error.to_string().contains("injected input error")); + assert!(matches!( + Pin::new(&mut stream).poll_next(&mut cx), + Poll::Ready(None) + )); + assert_eq!(stream.num_series.value(), 1); + } + + #[test] + fn extreme_alignment_retains_exact_sample() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])); + let input = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![i64::MAX])), + Arc::new(Float64Array::from(vec![7.0])), + ], + ) + .unwrap(); + let stream = InstantManipulateStream { + offset: 0, + start: i64::MIN + 1, + end: i64::MAX, + lookback_delta: 0, + interval: i64::MAX, + time_index: 0, + time_unit: TimeUnit::Millisecond, + field_indices: [Some(1), None], + tsid_index: None, + reuse_tsid_column: false, + schema: schema.clone(), + input: Box::pin( + datafusion::physical_plan::memory::MemoryStream::try_new(vec![], schema, None) + .unwrap(), + ), + metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + num_series: Count::new(), + }; + + let output = stream.manipulate(input).unwrap(); + let timestamps = output + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let values = output + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(timestamps.values(), &[i64::MAX]); + assert_eq!(values.values(), &[7.0]); + } + + #[test] + fn native_nanosecond_offset_uses_wide_shifted_timeline() { + for (raw, offset, eval) in [ + ( + 9_223_112_837_000_000_000_i64, + 259_200_000, + 9_223_372_037_000, + ), + ( + -9_223_112_837_000_000_000_i64, + -259_200_000, + -9_223_372_037_000, + ), + ] { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])); + let input = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampNanosecondArray::from(vec![raw])), + Arc::new(Float64Array::from(vec![7.0])), + ], + ) + .unwrap(); + let stream = InstantManipulateStream { + offset, + start: eval, + end: eval, + lookback_delta: 300_000, + interval: 1, + time_index: 0, + time_unit: TimeUnit::Nanosecond, + field_indices: [Some(1), None], + tsid_index: None, + reuse_tsid_column: false, + schema: Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])), + input: Box::pin( + datafusion::physical_plan::memory::MemoryStream::try_new(vec![], schema, None) + .unwrap(), + ), + metric: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + num_series: Count::new(), + }; + let output = stream.manipulate(input).unwrap(); + assert_eq!(output.num_rows(), 1); + assert_eq!( + output + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + eval + ); + assert_eq!( + output + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 7.0 + ); + } + } + #[tokio::test] async fn ordinary_nan_is_selected_for_exact_and_lookback() { let schema = Arc::new(Schema::new(vec![ @@ -1358,6 +2065,7 @@ mod test { MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), ))); let exec = Arc::new(InstantManipulateExec { + offset: 0, start: 1_000, end: 1_500, lookback_delta: 1_000, @@ -1365,6 +2073,8 @@ mod test { time_index_column: TIME_INDEX_COLUMN.to_string(), field_column: Some("value".to_string()), reuse_tsid_column: false, + output_schema: input.schema(), + properties: input.properties().clone(), input, metric: ExecutionPlanMetricsSet::new(), }); @@ -1437,6 +2147,7 @@ mod test { MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), ))); let exec = Arc::new(InstantManipulateExec { + offset: 0, start: 750, end: 1_500, lookback_delta: 1_001, @@ -1444,6 +2155,8 @@ mod test { time_index_column: TIME_INDEX_COLUMN.to_string(), field_column: Some("value".to_string()), reuse_tsid_column: false, + output_schema: input.schema(), + properties: input.properties().clone(), input, metric: ExecutionPlanMetricsSet::new(), }); @@ -1500,6 +2213,7 @@ mod test { MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), ))); let exec = Arc::new(InstantManipulateExec { + offset: 0, start: 1_000, end: 1_500, lookback_delta: 1_001, @@ -1507,6 +2221,8 @@ mod test { time_index_column: TIME_INDEX_COLUMN.to_string(), field_column: Some("value".to_string()), reuse_tsid_column: false, + output_schema: input.schema(), + properties: input.properties().clone(), input, metric: ExecutionPlanMetricsSet::new(), }); @@ -1547,6 +2263,7 @@ mod test { MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), ))); let exec = Arc::new(InstantManipulateExec { + offset: 0, start: 1_000, end: 1_500, lookback_delta: 1_000, @@ -1554,6 +2271,8 @@ mod test { time_index_column: TIME_INDEX_COLUMN.to_string(), field_column: Some("value".to_string()), reuse_tsid_column: false, + output_schema: input.schema(), + properties: input.properties().clone(), input, metric: ExecutionPlanMetricsSet::new(), }); diff --git a/src/promql/src/extension_plan/normalize.rs b/src/promql/src/extension_plan/normalize.rs index e3410be9267..c5f8f073695 100644 --- a/src/promql/src/extension_plan/normalize.rs +++ b/src/promql/src/extension_plan/normalize.rs @@ -20,7 +20,7 @@ use std::task::{Context, Poll}; use common_query::native_histogram::{START_TIMESTAMP_FIELD, native_histogram_arrow_type}; use datafusion::arrow::array::{Array, BooleanArray, StructArray}; use datafusion::arrow::compute; -use datafusion::common::{DFSchema, DFSchemaRef, Result as DataFusionResult, Statistics}; +use datafusion::common::{Column, DFSchema, DFSchemaRef, Result as DataFusionResult, Statistics}; use datafusion::error::DataFusionError; use datafusion::execution::context::TaskContext; use datafusion::logical_expr::{EmptyRelation, Expr, LogicalPlan, UserDefinedLogicalNodeCore}; @@ -48,13 +48,13 @@ use crate::extension_plan::{ }; use crate::metrics::PROMQL_SERIES_COUNT; -/// Normalize the input record batch. Notice that for simplicity, this method assumes -/// the input batch only contains sample points from one time series. +/// Normalizes a single-series input batch and optionally removes Prometheus stale markers. /// -/// Roughly speaking, this method does these things: -/// - bias sample and native histogram start timestamps by offset -/// - sort the record batch based on timestamp column -/// - remove Prometheus stale markers (optional) +/// This node remains the serialized carrier of the selector offset. Native sample timestamps +/// stay raw: applying an offset can overflow native `i64` ticks even when evaluation is valid. +/// Manipulators instead apply it in `i128` during selection and produce millisecond outputs. +/// Histogram start timestamps are already millisecond payloads used for rate/reset, so their +/// offsets are applied here, preserving unknown zero values and nulls. #[derive(Debug, PartialEq, Eq, Hash, PartialOrd)] pub struct SeriesNormalize { offset: Millisecond, @@ -179,6 +179,14 @@ impl UserDefinedLogicalNodeCore for SeriesNormalize { } impl SeriesNormalize { + pub(crate) fn offset_for_time_index(&self, time_index: &Column) -> Option { + let index = self.input.schema().maybe_index_of_column(time_index)?; + let (qualifier, field) = self.input.schema().qualified_field(index); + (field.name() == &self.time_index_column_name + && time_index == &Column::new(qualifier.cloned(), field.name().clone())) + .then_some(self.offset) + } + pub fn new>( offset: Millisecond, time_index_column_name: N, @@ -334,13 +342,8 @@ impl ExecutionPlan for SeriesNormalizeExec { let input = self.input.execute(partition, context)?; let schema = input.schema(); - let time_index = schema - .column_with_name(&self.time_index_column_name) - .expect("time index column not found") - .0; Ok(Box::pin(SeriesNormalizeStream { offset: self.offset, - time_index, filter_stale_markers: self.filter_stale_markers, schema, input, @@ -380,8 +383,6 @@ impl DisplayAs for SeriesNormalizeExec { pub struct SeriesNormalizeStream { offset: Millisecond, - // Column index of TIME INDEX column's position in schema - time_index: usize, filter_stale_markers: bool, schema: SchemaRef, @@ -393,33 +394,12 @@ pub struct SeriesNormalizeStream { impl SeriesNormalizeStream { pub fn normalize(&self, input: RecordBatch) -> DataFusionResult { - let ts_column = input - .column(self.time_index) - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Execution( - "Time index Column downcast to TimestampMillisecondArray failed".into(), - ) - })?; - - let bias_timestamp = |timestamp: i64| { - timestamp.checked_add(self.offset).ok_or_else(|| { - DataFusionError::Execution("SeriesNormalize: timestamp offset overflow".into()) - }) - }; - - // bias the timestamp column by offset - let ts_column_biased = if self.offset == 0 { - Arc::new(ts_column.clone()) as _ - } else { - Arc::new(ts_column.try_unary::<_, TimestampMillisecondType, _>(&bias_timestamp)?) - }; + // Native sample timestamps remain raw. Manipulators apply the selector offset + // in wide nanosecond arithmetic, avoiding overflow in native Arrow storage. let mut columns = input.columns().to_vec(); - columns[self.time_index] = ts_column_biased; - // Offset selectors move samples into the evaluation timeline. Keep native histogram - // start timestamps on the same timeline for rate and reset calculations. + // Offset selectors move native histogram start timestamps onto the evaluation + // timeline for rate and reset calculations. These payloads are milliseconds. if self.offset != 0 { let native_histogram_type = native_histogram_arrow_type(); for column in &mut columns { @@ -444,11 +424,17 @@ impl SeriesNormalizeStream { if timestamp == 0 { Ok(0) } else { - bias_timestamp(timestamp) + timestamp.checked_add(self.offset).ok_or_else(|| { + DataFusionError::Execution( + "SeriesNormalize: histogram timestamp offset overflow".into(), + ) + }) } })?; - // Replace only the start timestamp child to preserve the histogram payload and - // null bitmap. + // Struct arrays are immutable, so rebuild the physical histogram payload with + // only its start-timestamp child replaced. Its logical schema stays unchanged: + // the selector offset affects histogram reset/rate metadata, not the native + // sample timestamp column consumed by later manipulators. let mut children = histograms.columns().to_vec(); children[start_timestamp_index] = Arc::new(start_timestamps); *column = Arc::new(StructArray::new( @@ -516,10 +502,12 @@ impl Stream for SeriesNormalizeStream { mod test { use common_query::native_histogram::{build_histogram_array, read_histogram}; use common_query::prometheus::PROMETHEUS_STALE_NAN_BITS; - use datafusion::arrow::array::Float64Array; + use datafusion::arrow::array::{ + DictionaryArray, Float64Array, TimestampMicrosecondArray, TimestampNanosecondArray, + }; use datafusion::arrow::buffer::NullBuffer; use datafusion::arrow::datatypes::{ - ArrowPrimitiveType, DataType, Field, Schema, TimestampMillisecondType, + ArrowPrimitiveType, DataType, Field, Int64Type, Schema, TimeUnit, TimestampMillisecondType, }; use datafusion::common::ToDFSchema; use datafusion::datasource::memory::MemorySourceConfig; @@ -530,7 +518,9 @@ mod test { use datatypes::arrow_array::StringArray; use super::*; + use crate::extension_plan::RangeManipulate; use crate::extension_plan::test_util::native_histogram; + use crate::range_array::RangeArray; const TIME_INDEX_COLUMN: &str = "timestamp"; @@ -630,11 +620,11 @@ mod test { "+---------------------+--------+------+\ \n| timestamp | value | path |\ \n+---------------------+--------+------+\ - \n| 1970-01-01T00:01:01 | 0.0 | foo |\ - \n| 1970-01-01T00:02:01 | 1.0 | foo |\ - \n| 1970-01-01T00:00:01 | 10.0 | foo |\ - \n| 1970-01-01T00:00:31 | 100.0 | foo |\ - \n| 1970-01-01T00:01:31 | 1000.0 | foo |\ + \n| 1970-01-01T00:01:00 | 0.0 | foo |\ + \n| 1970-01-01T00:02:00 | 1.0 | foo |\ + \n| 1970-01-01T00:00:00 | 10.0 | foo |\ + \n| 1970-01-01T00:00:30 | 100.0 | foo |\ + \n| 1970-01-01T00:01:30 | 1000.0 | foo |\ \n+---------------------+--------+------+", ); @@ -720,12 +710,104 @@ mod test { regular.start_timestamp = Some(500); let mut ordinary_nan = native_histogram(f64::NAN); ordinary_nan.start_timestamp = Some(0); + let mut unknown_start = native_histogram(7.0); + unknown_start.start_timestamp = None; let histograms = build_histogram_array(&[ Some(regular), Some(native_histogram(f64::from_bits(PROMETHEUS_STALE_NAN_BITS))), Some(ordinary_nan), + Some(unknown_start), None, ]); + for (unit, ticks_per_ms) in [ + (TimeUnit::Millisecond, 1_i64), + (TimeUnit::Microsecond, 1_000), + (TimeUnit::Nanosecond, 1_000_000), + ] { + let timestamp_array = |values: Vec| -> Arc { + match unit { + TimeUnit::Millisecond => Arc::new(TimestampMillisecondArray::from(values)), + TimeUnit::Microsecond => Arc::new(TimestampMicrosecondArray::from(values)), + TimeUnit::Nanosecond => Arc::new(TimestampNanosecondArray::from(values)), + TimeUnit::Second => unreachable!(), + } + }; + for offset in [-1_i64, 1] { + let timestamps = timestamp_array( + [1_000, 2_000, 3_000, 4_000, 5_000] + .into_iter() + .map(|timestamp| timestamp * ticks_per_ms) + .collect(), + ); + let schema = Arc::new(Schema::new(vec![ + Field::new(TIME_INDEX_COLUMN, timestamps.data_type().clone(), false), + Field::new("value", histograms.data_type().clone(), true), + ])); + let batch = + RecordBatch::try_new(schema.clone(), vec![timestamps, histograms.clone()]) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), + ))); + let exec = Arc::new(SeriesNormalizeExec { + offset, + time_index_column_name: TIME_INDEX_COLUMN.to_string(), + filter_stale_markers: true, + tag_columns: Vec::new(), + input, + metric: ExecutionPlanMetricsSet::new(), + }); + let context = SessionContext::default(); + let batches = datafusion::physical_plan::collect(exec, context.task_ctx()) + .await + .unwrap(); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 4, + "unit={unit:?}, offset={offset}" + ); + let batch = batches.iter().find(|batch| batch.num_rows() == 4).unwrap(); + let expected_timestamps = timestamp_array( + [1_000, 3_000, 4_000, 5_000] + .into_iter() + .map(|timestamp| timestamp * ticks_per_ms) + .collect(), + ); + assert_eq!( + batch.column(0).to_data(), + expected_timestamps.to_data(), + "unit={unit:?}, offset={offset}" + ); + let values = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let regular = read_histogram(values, 0).unwrap().unwrap(); + assert_eq!( + (regular.sum, regular.start_timestamp), + (42.0, Some(500 + offset)), + "unit={unit:?}, offset={offset}" + ); + let ordinary_nan = read_histogram(values, 1).unwrap().unwrap(); + assert!(ordinary_nan.sum.is_nan()); + assert_eq!(ordinary_nan.start_timestamp, Some(0)); + let unknown_start = read_histogram(values, 2).unwrap().unwrap(); + assert_eq!( + (unknown_start.sum, unknown_start.start_timestamp), + (7.0, None) + ); + assert!(read_histogram(values, 3).unwrap().is_none()); + } + } + + let mut known_start = native_histogram(42.0); + known_start.start_timestamp = Some(500); + let mut sentinel_start = native_histogram(8.0); + sentinel_start.start_timestamp = Some(0); + let unknown_start = native_histogram(7.0); + let histograms = + build_histogram_array(&[Some(known_start), Some(sentinel_start), Some(unknown_start)]); let schema = Arc::new(Schema::new(vec![ Field::new( TIME_INDEX_COLUMN, @@ -737,47 +819,85 @@ mod test { let batch = RecordBatch::try_new( schema.clone(), vec![ - Arc::new(TimestampMillisecondArray::from(vec![ - 1_000, 2_000, 3_000, 4_000, - ])), + Arc::new(TimestampMillisecondArray::from(vec![1_000; 3])), histograms, ], ) .unwrap(); + let logical_input = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: schema.clone().to_dfschema_ref().unwrap(), + }); + let normalized = + SeriesNormalize::new(1_000, TIME_INDEX_COLUMN, false, Vec::new(), logical_input); + let range = RangeManipulate::new( + 2_000, + 2_000, + 1, + 1_000, + 1, + TIME_INDEX_COLUMN.to_string(), + vec!["value".to_string()], + LogicalPlan::Extension(datafusion::logical_expr::Extension { + node: Arc::new(normalized), + }), + ) + .unwrap(); let input = Arc::new(DataSourceExec::new(Arc::new( MemorySourceConfig::try_new(&[vec![batch]], schema, None).unwrap(), ))); - let exec = Arc::new(SeriesNormalizeExec { + let normalized_input = Arc::new(SeriesNormalizeExec { offset: 1_000, time_index_column_name: TIME_INDEX_COLUMN.to_string(), - filter_stale_markers: true, + filter_stale_markers: false, tag_columns: Vec::new(), input, metric: ExecutionPlanMetricsSet::new(), }); - - let context = SessionContext::default(); - let batches = datafusion::physical_plan::collect(exec, context.task_ctx()) - .await - .unwrap(); - let batch = batches.iter().find(|batch| batch.num_rows() == 3).unwrap(); - let values = batch - .column(1) + let output = datafusion::physical_plan::collect( + range.to_execution_plan(normalized_input), + SessionContext::default().task_ctx(), + ) + .await + .unwrap(); + let values = RangeArray::try_new( + output[0] + .column(1) + .as_any() + .downcast_ref::>() + .unwrap() + .clone(), + ) + .unwrap(); + let values = values.get(0).unwrap(); + let values = values .as_any() .downcast_ref::() .unwrap(); - - let timestamps = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(timestamps.values(), &[2_000, 4_000, 5_000]); - let regular = read_histogram(values, 0).unwrap().unwrap(); - assert_eq!((regular.sum, regular.start_timestamp), (42.0, Some(1_500))); - let ordinary_nan = read_histogram(values, 1).unwrap().unwrap(); - assert!(ordinary_nan.sum.is_nan()); - assert_eq!(ordinary_nan.start_timestamp, Some(0)); - assert!(read_histogram(values, 2).unwrap().is_none()); + assert_eq!( + read_histogram(values, 0).unwrap().unwrap().start_timestamp, + Some(1_500) + ); + assert_eq!( + read_histogram(values, 1).unwrap().unwrap().start_timestamp, + Some(0) + ); + assert_eq!( + read_histogram(values, 2).unwrap().unwrap().start_timestamp, + None + ); + let timestamps = RangeArray::try_new( + output[0] + .column(2) + .as_any() + .downcast_ref::>() + .unwrap() + .clone(), + ) + .unwrap(); + assert_eq!( + timestamps.get(0).unwrap().to_data(), + TimestampMillisecondArray::from(vec![2_000; 3]).to_data() + ); } } diff --git a/src/promql/src/extension_plan/range_manipulate.rs b/src/promql/src/extension_plan/range_manipulate.rs index 44c5f490947..9bce8e05ad2 100644 --- a/src/promql/src/extension_plan/range_manipulate.rs +++ b/src/promql/src/extension_plan/range_manipulate.rs @@ -21,7 +21,7 @@ use std::task::{Context, Poll}; use common_telemetry::{debug, warn}; use datafusion::arrow::array::{Array, ArrayRef, Int64Array, TimestampMillisecondArray}; use datafusion::arrow::compute; -use datafusion::arrow::datatypes::{Field, SchemaRef}; +use datafusion::arrow::datatypes::{DataType, Field, SchemaRef, TimeUnit}; use datafusion::arrow::error::ArrowError; use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::stats::Precision; @@ -39,6 +39,7 @@ use datafusion::physical_plan::{ }; use datafusion::sql::TableReference; use datafusion_expr::col; +use datatypes::timestamp::timestamp_array_to_primitive; use futures::{Stream, StreamExt, ready}; use greptime_proto::substrait_extension as pb; use prost::Message; @@ -46,7 +47,8 @@ use snafu::ResultExt; use crate::error::{DeserializeSnafu, Result}; use crate::extension_plan::{ - METRIC_NUM_SERIES, Millisecond, resolve_column_name, serialize_column_index, + METRIC_NUM_SERIES, Millisecond, local_offset, nanoseconds_per_native_tick, resolve_column_name, + serialize_column_index, timestamp_unit, }; use crate::metrics::PROMQL_SERIES_COUNT; use crate::range_array::RangeArray; @@ -66,6 +68,7 @@ pub struct RangeManipulate { end: Millisecond, interval: Millisecond, range: Millisecond, + offset: Millisecond, time_index: String, field_columns: Vec, input: LogicalPlan, @@ -80,10 +83,12 @@ struct UnfixIndices { } impl RangeManipulate { + #[allow(clippy::too_many_arguments)] pub fn new( start: Millisecond, end: Millisecond, interval: Millisecond, + offset: Millisecond, range: Millisecond, time_index: String, field_columns: Vec, @@ -96,6 +101,7 @@ impl RangeManipulate { end, interval, range, + offset, time_index, field_columns, input, @@ -142,9 +148,21 @@ impl RangeManipulate { )); }; let ts_col_field = &columns[ts_col_index]; + let output_time_field = Arc::new( + ts_col_field + .as_ref() + .clone() + .with_data_type(DataType::Timestamp(TimeUnit::Millisecond, None)), + ); + new_columns[ts_col_index] = ( + input_schema.qualified_field(ts_col_index).0.cloned(), + output_time_field.clone(), + ); let timestamp_range_field = Field::new( Self::build_timestamp_range_name(time_index), - RangeArray::convert_field(ts_col_field).data_type().clone(), + RangeArray::convert_field(output_time_field.as_ref()) + .data_type() + .clone(), ts_col_field.is_nullable(), ); new_columns.push((None, Arc::new(timestamp_range_field))); @@ -177,6 +195,7 @@ impl RangeManipulate { properties.boundedness, )); Arc::new(RangeManipulateExec { + offset: self.offset, start: self.start, end: self.end, interval: self.interval, @@ -235,6 +254,7 @@ impl RangeManipulate { end: pb_range_manipulate.end, interval: pb_range_manipulate.interval, range: pb_range_manipulate.range, + offset: 0, time_index: String::new(), field_columns: Vec::new(), input: placeholder_plan, @@ -263,6 +283,10 @@ impl PartialOrd for RangeManipulate { Some(core::cmp::Ordering::Equal) => {} ord => return ord, } + match self.offset.partial_cmp(&other.offset) { + Some(core::cmp::Ordering::Equal) => {} + ord => return ord, + } match self.time_index.partial_cmp(&other.time_index) { Some(core::cmp::Ordering::Equal) => {} ord => return ord, @@ -383,6 +407,7 @@ impl UserDefinedLogicalNodeCore for RangeManipulate { end: self.end, interval: self.interval, range: self.range, + offset: local_offset(&input, &time_index), time_index, field_columns, input, @@ -398,6 +423,7 @@ impl UserDefinedLogicalNodeCore for RangeManipulate { end: self.end, interval: self.interval, range: self.range, + offset: self.offset, time_index: self.time_index.clone(), field_columns: self.field_columns.clone(), input, @@ -410,6 +436,7 @@ impl UserDefinedLogicalNodeCore for RangeManipulate { #[derive(Debug)] pub struct RangeManipulateExec { + offset: Millisecond, start: Millisecond, end: Millisecond, interval: Millisecond, @@ -470,6 +497,7 @@ impl ExecutionPlan for RangeManipulateExec { properties.boundedness, )); Ok(Arc::new(Self { + offset: self.offset, start: self.start, end: self.end, interval: self.interval, @@ -515,14 +543,17 @@ impl ExecutionPlan for RangeManipulateExec { .0 }) .collect(); + let time_unit = timestamp_unit(schema.field(time_index).data_type())?; let aligned_ts_array = RangeManipulateStream::build_aligned_ts_array(self.start, self.end, self.interval); Ok(Box::pin(RangeManipulateStream { + offset: self.offset, start: self.start, end: self.end, interval: self.interval, range: self.range, time_index, + time_unit, field_columns, aligned_ts_array, output_schema: self.output_schema.clone(), @@ -579,11 +610,13 @@ impl DisplayAs for RangeManipulateExec { } pub struct RangeManipulateStream { + offset: Millisecond, start: Millisecond, end: Millisecond, interval: Millisecond, range: Millisecond, time_index: usize, + time_unit: TimeUnit, field_columns: Vec, aligned_ts_array: ArrayRef, @@ -655,11 +688,35 @@ impl RangeManipulateStream { new_columns[*index] = new_column; } - // push timestamp range column - let ts_range_column = - RangeArray::from_ranges(input.column(self.time_index).clone(), ranges.clone()) - .map_err(|e| ArrowError::InvalidArgumentError(e.to_string()))? - .into_dict(); + // The timestamp range payload is always millisecond ABI. Shift in wide + // native precision before truncating toward zero, preserving null validity. + let scale = nanoseconds_per_native_tick(self.time_unit); + let (timestamps, _) = timestamp_array_to_primitive(input.column(self.time_index)) + .ok_or_else(|| { + DataFusionError::Execution("Time index column is not a timestamp".into()) + })?; + let timestamp_values = timestamps + .values() + .iter() + .enumerate() + .map(|(index, timestamp)| { + if !input.column(self.time_index).is_valid(index) { + return Ok(None); + } + let shifted_ns = (*timestamp as i128) * scale + (self.offset as i128) * 1_000_000; + i64::try_from(shifted_ns / 1_000_000) + .map(Some) + .map_err(|_| { + ArrowError::ComputeError( + "RangeManipulate timestamp payload overflow".into(), + ) + }) + }) + .collect::, _>>()?; + let timestamp_values = TimestampMillisecondArray::from(timestamp_values); + let ts_range_column = RangeArray::from_ranges(Arc::new(timestamp_values), ranges.clone()) + .map_err(|e| ArrowError::InvalidArgumentError(e.to_string()))? + .into_dict(); new_columns.push(Arc::new(ts_range_column)); // truncate other columns @@ -694,52 +751,58 @@ impl RangeManipulateStream { &self, input: &RecordBatch, ) -> DataFusionResult<(Vec<(u32, u32)>, (i64, i64))> { - let ts_column = input - .column(self.time_index) - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Execution( - "Time index Column downcast to TimestampMillisecondArray failed".into(), - ) - })?; - - let len = ts_column.len(); + let ts_column = input.column(self.time_index); + let scale = nanoseconds_per_native_tick(self.time_unit); + let (timestamps, _) = timestamp_array_to_primitive(ts_column).ok_or_else(|| { + DataFusionError::Execution("Time index column is not a timestamp".into()) + })?; + let timestamps = timestamps.values(); + let timestamp = + |index| (timestamps[index] as i128) * scale + (self.offset as i128) * 1_000_000; + let len = timestamps.len(); if len == 0 { return Ok((vec![], (self.start, self.end))); } - // shorten the range to calculate - let first_ts = ts_column.value(0); - // Preserve the query's alignment pattern when optimizing start time - let remainder = (first_ts - self.start).rem_euclid(self.interval); - let first_ts_aligned = if remainder == 0 { - first_ts - } else { - first_ts + (self.interval - remainder) - }; - let last_ts = ts_column.value(ts_column.len() - 1); - let last_ts_with_range = last_ts + self.range; - let remainder = (last_ts_with_range - self.start).rem_euclid(self.interval); + // Shorten the range using wide arithmetic so timestamps near the native + // type limits retain every query-aligned evaluation point. + let query_start = self.start as i128; + let query_end = self.end as i128; + let interval = self.interval as i128; + let first_ts = timestamp(0).div_euclid(1_000_000); + // Preserve the query's alignment pattern when optimizing start time. + let remainder = (first_ts - query_start).rem_euclid(interval); + let first_ts_aligned = first_ts + (interval - remainder).rem_euclid(interval); + let last_ts_with_range = + (timestamp(len - 1) + (self.range as i128) * 1_000_000).div_euclid(1_000_000); + let remainder = (last_ts_with_range - query_start).rem_euclid(interval); let last_ts_aligned = last_ts_with_range - remainder; - let start = self.start.max(first_ts_aligned); - let end = self.end.min(last_ts_aligned); + let start = query_start.max(first_ts_aligned); + let end = query_end.min(last_ts_aligned); if start > end { - return Ok((vec![], (start, end))); + return Ok((vec![], (self.start, self.end))); } - let mut ranges = Vec::with_capacity(((self.end - self.start) / self.interval + 1) as usize); + // The intersection is within the declared i64 query bounds. + let start = start as i64; + let end = end as i64; + let mut ranges = Vec::new(); - // calculate for every aligned timestamp (`curr_ts`), assume the ts column is ordered. + // Range membership is decided on shifted native ticks, before the + // timestamp-range payload is converted to its millisecond ABI. This + // keeps sub-millisecond samples distinct in a range; equal millisecond + // payload values are not a reason to deduplicate input samples. + // + // Calculate for every aligned timestamp (`curr_ts`), assuming ordered timestamps. let mut left = 0usize; let mut right = 0usize; for curr_ts in (start..=end).step_by(self.interval as _) { - let start_ts = curr_ts - self.range; + let start_ts = (curr_ts as i128) * 1_000_000 - (self.range as i128) * 1_000_000; - while left < len && ts_column.value(left) <= start_ts { + while left < len && timestamp(left) <= start_ts { left += 1; } right = right.max(left); - while right < len && ts_column.value(right) <= curr_ts { + while right < len && timestamp(right) <= (curr_ts as i128) * 1_000_000 { right += 1; } @@ -756,14 +819,20 @@ impl RangeManipulateStream { #[cfg(test)] mod test { - use datafusion::arrow::array::{ArrayRef, DictionaryArray, Float64Array, StringArray}; + use datafusion::arrow::array::{ + ArrayRef, DictionaryArray, Float64Array, StringArray, TimestampMicrosecondArray, + TimestampNanosecondArray, TimestampSecondArray, + }; + use datafusion::arrow::buffer::NullBuffer; use datafusion::arrow::datatypes::{ ArrowPrimitiveType, DataType, Field, Int64Type, Schema, TimestampMillisecondType, }; use datafusion::common::ToDFSchema; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; - use datafusion::logical_expr::{EmptyRelation, LogicalPlan}; + use datafusion::logical_expr::{ + EmptyRelation, Extension, LogicalPlan, UserDefinedLogicalNodeCore, + }; use datafusion::physical_expr::Partitioning; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::memory::MemoryStream; @@ -845,6 +914,7 @@ mod test { Boundedness::Bounded, )); let normalize_exec = Arc::new(RangeManipulateExec { + offset: 0, start, end, interval, @@ -888,6 +958,432 @@ mod test { assert_eq!(result_literal, expected); } + #[tokio::test] + async fn native_timestamps_preserve_range_membership_and_ms_payload() { + for (unit, ticks_per_ms) in [ + (TimeUnit::Microsecond, 1_000_i64), + (TimeUnit::Nanosecond, 1_000_000_i64), + ] { + let lower = 1_000 * ticks_per_ms; + let upper = 1_001 * ticks_per_ms; + // Exclude the lower boundary and future sample; retain both native + // samples in the same millisecond bucket and the exact upper sample. + let timestamps = vec![lower, lower + 1, lower + 2, upper, upper + 1]; + let time: ArrayRef = match unit { + TimeUnit::Microsecond => Arc::new(TimestampMicrosecondArray::from(timestamps)), + TimeUnit::Nanosecond => Arc::new(TimestampNanosecondArray::from(timestamps)), + _ => unreachable!(), + }; + let schema = Arc::new(Schema::new(vec![ + Field::new(TIME_INDEX_COLUMN, DataType::Timestamp(unit, None), false), + Field::new("value", DataType::Float64, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + time, + Arc::new(Float64Array::from(vec![10.0, 20.0, 30.0, 40.0, 50.0])), + ], + ) + .unwrap(); + let logical_input = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: schema.clone().to_dfschema_ref().unwrap(), + }); + let plan = RangeManipulate::new( + 1_001, + 1_001, + 1, + 0, + 1, + TIME_INDEX_COLUMN.to_string(), + vec!["value".to_string()], + logical_input.clone(), + ) + .unwrap(); + let output_time = Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ); + let output_schema = Arc::new(Schema::new(vec![ + output_time.clone(), + RangeArray::convert_field(&Field::new("value", DataType::Float64, true)), + Field::new( + RangeManipulate::build_timestamp_range_name(TIME_INDEX_COLUMN), + RangeArray::convert_field(&output_time).data_type().clone(), + false, + ), + ])); + assert_eq!(plan.schema().as_arrow(), output_schema.as_ref()); + + let rebuilt = RangeManipulate::deserialize(&plan.serialize()) + .unwrap() + .with_exprs_and_inputs(vec![], vec![logical_input]) + .unwrap(); + assert_eq!(rebuilt.schema(), plan.schema()); + assert_eq!(rebuilt.input.schema().as_arrow(), schema.as_ref()); + + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema.clone(), None).unwrap(), + ))); + let exec = rebuilt.to_execution_plan(input); + assert_eq!(exec.schema(), output_schema); + assert_eq!(exec.children()[0].schema(), schema); + + let batches = + datafusion::physical_plan::collect(exec, SessionContext::default().task_ctx()) + .await + .unwrap(); + assert_eq!(batches.len(), 1, "{unit:?}"); + let output = &batches[0]; + assert_eq!(output.schema(), output_schema); + assert_eq!(output.num_rows(), 1); + assert_eq!( + output + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .as_ref(), + &[1_001] + ); + + // RangeArray packs offset/length into dictionary keys; Arrow dictionary + // equality treats those packed keys as indices and cannot compare them. + let values = RangeArray::try_new( + output + .column(1) + .as_any() + .downcast_ref::>() + .unwrap() + .clone(), + ) + .unwrap(); + assert_eq!(values.get_offset_length(0), Some((1, 3))); + assert_eq!( + values.get(0).unwrap().to_data(), + Float64Array::from(vec![20.0, 30.0, 40.0]).to_data() + ); + let timestamps = RangeArray::try_new( + output + .column(2) + .as_any() + .downcast_ref::>() + .unwrap() + .clone(), + ) + .unwrap(); + assert_eq!(timestamps.get_offset_length(0), Some((1, 3))); + assert_eq!( + timestamps.get(0).unwrap().to_data(), + TimestampMillisecondArray::from(vec![1_000, 1_000, 1_001]).to_data() + ); + } + } + + #[test] + fn logical_offset_participates_in_ordering() { + let input = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: prepare_test_data().schema().to_dfschema_ref().unwrap(), + }); + let first = RangeManipulate::new( + 0, + 0, + 0, + 0, + 0, + TIME_INDEX_COLUMN.to_string(), + vec!["value_1".to_string()], + input.clone(), + ) + .unwrap(); + let second = RangeManipulate::new( + 0, + 0, + 0, + 1, + 0, + TIME_INDEX_COLUMN.to_string(), + vec!["value_1".to_string()], + input, + ) + .unwrap(); + assert_ne!(first, second); + assert_eq!(first.partial_cmp(&second), Some(std::cmp::Ordering::Less)); + } + + #[tokio::test] + async fn logical_normalize_offset_survives_rebuild_and_executes() { + for (name, time_unit, raw, offset, start, range, expected_payload) in [ + ( + "millisecond offset", + TimeUnit::Millisecond, + 0, + 1_000, + 1_000, + 1_000, + 1_000, + ), + ( + "negative native lower limit with positive window", + TimeUnit::Nanosecond, + -9_223_112_837_000_000_000, + -259_200_000, + -9_223_372_037_000, + 300_000, + -9_223_372_037_000, + ), + ( + "second timestamp with negative fractional offset", + TimeUnit::Second, + 1, + -500, + 1_000, + 1_000, + 500, + ), + ] { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + DataType::Timestamp(time_unit, None), + false, + ), + Field::new("value", DataType::Float64, true), + ])); + let input = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: schema.clone().to_dfschema_ref().unwrap(), + }); + let normalize = crate::extension_plan::SeriesNormalize::new( + offset, + TIME_INDEX_COLUMN, + false, + Vec::new(), + input.clone(), + ); + let normalize = + crate::extension_plan::SeriesNormalize::deserialize(&normalize.serialize()) + .unwrap() + .with_exprs_and_inputs(vec![], vec![input.clone()]) + .unwrap(); + let normalized = LogicalPlan::Extension(Extension { + node: Arc::new(normalize), + }); + let fresh = RangeManipulate::new( + start, + start, + 1, + offset, + range, + TIME_INDEX_COLUMN.to_string(), + vec!["value".to_string()], + input.clone(), + ) + .unwrap() + .with_exprs_and_inputs(vec![], vec![input.clone()]) + .unwrap(); + let serialized = RangeManipulate::new( + start, + start, + 1, + offset, + range, + TIME_INDEX_COLUMN.to_string(), + vec!["value".to_string()], + normalized.clone(), + ) + .unwrap(); + let decoded = RangeManipulate::deserialize(&serialized.serialize()) + .unwrap() + .with_exprs_and_inputs(vec![], vec![normalized]) + .unwrap(); + let timestamp: ArrayRef = match time_unit { + TimeUnit::Millisecond => Arc::new(TimestampMillisecondArray::from(vec![raw])), + TimeUnit::Nanosecond => Arc::new(TimestampNanosecondArray::from(vec![raw])), + TimeUnit::Second => Arc::new(TimestampSecondArray::from(vec![raw])), + _ => unreachable!(), + }; + let batch = RecordBatch::try_new( + schema.clone(), + vec![timestamp, Arc::new(Float64Array::from(vec![7.0]))], + ) + .unwrap(); + for (mode, rebuilt) in [("fresh", fresh), ("decoded", decoded)] { + let rebuilt = rebuilt + .with_exprs_and_inputs(vec![], vec![input.clone()]) + .unwrap(); + assert_eq!(rebuilt.offset, offset, "{name}: {mode}"); + assert_eq!(rebuilt.input.schema(), input.schema(), "{name}: {mode}"); + + let empty_exec_input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![]], schema.clone(), None).unwrap(), + ))); + let exec_input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch.clone()]], schema.clone(), None) + .unwrap(), + ))); + let exec = rebuilt + .to_execution_plan(empty_exec_input) + .with_new_children(vec![exec_input]) + .unwrap(); + let output = + datafusion::physical_plan::collect(exec, SessionContext::default().task_ctx()) + .await + .unwrap(); + assert_eq!(output.len(), 1, "{name}: {mode}"); + let output = &output[0]; + assert_eq!(output.num_rows(), 1, "{name}: {mode}"); + assert_eq!( + output + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + start, + "{name}: {mode}" + ); + let values = RangeArray::try_new( + output + .column(1) + .as_any() + .downcast_ref::>() + .unwrap() + .clone(), + ) + .unwrap(); + assert_eq!(values.get_offset_length(0), Some((0, 1)), "{name}: {mode}"); + assert_eq!( + values.get(0).unwrap().to_data(), + Float64Array::from(vec![7.0]).to_data(), + "{name}: {mode}" + ); + let timestamps = RangeArray::try_new( + output + .column(2) + .as_any() + .downcast_ref::>() + .unwrap() + .clone(), + ) + .unwrap(); + assert_eq!( + timestamps.get_offset_length(0), + Some((0, 1)), + "{name}: {mode}" + ); + assert_eq!( + timestamps.get(0).unwrap().to_data(), + TimestampMillisecondArray::from(vec![expected_payload]).to_data(), + "{name}: {mode}" + ); + } + } + } + + #[tokio::test] + async fn range_payload_preserves_null_timestamp_and_rejects_offset_overflow() { + let schema = Arc::new(Schema::new(vec![ + Field::new(TIME_INDEX_COLUMN, TimestampMillisecondType::DATA_TYPE, true), + Field::new("value", DataType::Float64, true), + ])); + let null_timestamp = + TimestampMillisecondArray::new(vec![1_000].into(), Some(NullBuffer::from(vec![false]))); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(null_timestamp), + Arc::new(Float64Array::from(vec![7.0])), + ], + ) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema.clone(), None).unwrap(), + ))); + let plan = RangeManipulate::new( + 1_000, + 1_000, + 1, + 0, + 1, + TIME_INDEX_COLUMN.to_string(), + vec!["value".to_string()], + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: schema.clone().to_dfschema_ref().unwrap(), + }), + ) + .unwrap(); + let output = datafusion::physical_plan::collect( + plan.to_execution_plan(input), + SessionContext::default().task_ctx(), + ) + .await + .unwrap(); + let timestamps = RangeArray::try_new( + output[0] + .column(2) + .as_any() + .downcast_ref::>() + .unwrap() + .clone(), + ) + .unwrap(); + let payload = timestamps.get(0).unwrap(); + let payload = payload + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(payload.len(), 1); + assert!(!payload.is_valid(0)); + + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![0, i64::MAX])), + Arc::new(Float64Array::from(vec![7.0, 8.0])), + ], + ) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![batch]], schema.clone(), None).unwrap(), + ))); + let normalized = crate::extension_plan::SeriesNormalize::new( + 1, + TIME_INDEX_COLUMN, + false, + Vec::new(), + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: schema.to_dfschema_ref().unwrap(), + }), + ); + let plan = RangeManipulate::new( + 1, + 1, + 1, + 1, + 1, + TIME_INDEX_COLUMN.to_string(), + vec!["value".to_string()], + LogicalPlan::Extension(Extension { + node: Arc::new(normalized), + }), + ) + .unwrap(); + let error = datafusion::physical_plan::collect( + plan.to_execution_plan(input), + SessionContext::default().task_ctx(), + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("timestamp payload overflow")); + } + #[tokio::test] async fn pruning_should_keep_time_and_value_columns_for_exec() { let schema = Arc::new(Schema::new(vec![ @@ -905,6 +1401,7 @@ mod test { 0, 310_000, 30_000, + 0, 90_000, TIME_INDEX_COLUMN.to_string(), vec!["value_1".to_string(), "value_2".to_string()], @@ -1042,11 +1539,13 @@ mod test { let empty_stream = MemoryStream::try_new(vec![], schema.clone(), None).unwrap(); let stream = RangeManipulateStream { + offset: 0, start: 1758093274000, // ends in 4000 end: 1758093334000, // ends in 4000 interval: 30000, // 30s step range: 60000, // 60s lookback time_index: 0, + time_unit: TimeUnit::Millisecond, field_columns: vec![], aligned_ts_array: Arc::new(TimestampMillisecondArray::from(vec![0i64; 0])), output_schema: schema.clone(), @@ -1092,6 +1591,69 @@ mod test { } } + #[tokio::test] + async fn no_intersection_batch_is_skipped_and_stream_continues() { + let schema = Arc::new(Schema::new(vec![ + Field::new( + TIME_INDEX_COLUMN, + TimestampMillisecondType::DATA_TYPE, + false, + ), + Field::new("value", DataType::Float64, false), + ])); + let input = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: schema.clone().to_dfschema_ref().unwrap(), + }); + let plan = RangeManipulate::new( + 0, + 50, + 10, + 0, + 1, + TIME_INDEX_COLUMN.to_string(), + vec!["value".to_string()], + input, + ) + .unwrap(); + let no_intersection = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![100])), + Arc::new(Float64Array::from(vec![1.0])), + ], + ) + .unwrap(); + let intersection = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![20])), + Arc::new(Float64Array::from(vec![2.0])), + ], + ) + .unwrap(); + let input = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![no_intersection, intersection]], schema, None) + .unwrap(), + ))); + + let batches = datafusion::physical_plan::collect( + plan.to_execution_plan(input), + SessionContext::default().task_ctx(), + ) + .await + .unwrap(); + + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 1); + let timestamps = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(timestamps.value(0), 20); + } + fn calculate_range_for_test( query_start: i64, query_end: i64, @@ -1106,11 +1668,13 @@ mod test { )])); let empty_stream = MemoryStream::try_new(vec![], schema.clone(), None).unwrap(); let stream = RangeManipulateStream { + offset: 0, start: query_start, end: query_end, interval, range, time_index: 0, + time_unit: TimeUnit::Millisecond, field_columns: vec![], aligned_ts_array: Arc::new(TimestampMillisecondArray::from(vec![0i64; 0])), output_schema: schema.clone(), @@ -1216,7 +1780,7 @@ mod test { 10, 0, vec![100], - (100, 50), + (0, 50), ), ]; @@ -1227,6 +1791,15 @@ mod test { } } + #[test] + fn calculate_range_keeps_extreme_range_tail() { + let (ranges, bounds) = + calculate_range_for_test(i64::MAX - 1, i64::MAX, 1, i64::MAX, &[i64::MAX]); + + assert_eq!(bounds, (i64::MAX, i64::MAX)); + assert_eq!(ranges, vec![(0, 1)]); + } + #[test] fn calculate_range_matches_bruteforce_oracle_for_deterministic_cases() { let cases = vec![ @@ -1345,6 +1918,17 @@ mod test { let (actual, (start, end)) = calculate_range_for_test(query_start, query_end, interval, range, ×tamps); let expected = calculate_range_oracle(×tamps, start, end, interval, range); + let expected = if actual.is_empty() && !expected.is_empty() { + assert!( + expected.iter().all(|(_, len)| *len == 0), + "case={case}, timestamps={timestamps:?}, query=({query_start}, {query_end}), \ + interval={interval}, range={range}, bounds=({start}, {end}): \ + no-intersection output must have no selected samples" + ); + vec![] + } else { + expected + }; assert_eq!( actual, expected, "case={case}, timestamps={timestamps:?}, query=({query_start}, {query_end}), \ diff --git a/src/promql/src/functions/aggr_over_time.rs b/src/promql/src/functions/aggr_over_time.rs index 05b2b1b2e2d..3eeae47e61f 100644 --- a/src/promql/src/functions/aggr_over_time.rs +++ b/src/promql/src/functions/aggr_over_time.rs @@ -26,6 +26,129 @@ use datatypes::arrow::datatypes::DataType; use crate::functions::{compensated_sum_inc, extract_array}; use crate::range_array::RangeArray; +#[derive(Clone, Copy)] +enum PresenceEvaluator { + Count, + Last, + Absent, + Present, +} + +fn count_over_time_evaluator( + input: &[ColumnarValue], + name: &str, +) -> Result { + evaluate_presence(input, name, PresenceEvaluator::Count) +} + +fn last_over_time_evaluator( + input: &[ColumnarValue], + name: &str, +) -> Result { + evaluate_presence(input, name, PresenceEvaluator::Last) +} + +fn absent_over_time_evaluator( + input: &[ColumnarValue], + name: &str, +) -> Result { + evaluate_presence(input, name, PresenceEvaluator::Absent) +} + +fn present_over_time_evaluator( + input: &[ColumnarValue], + name: &str, +) -> Result { + evaluate_presence(input, name, PresenceEvaluator::Present) +} + +fn evaluate_presence( + input: &[ColumnarValue], + name: &str, + operation: PresenceEvaluator, +) -> Result { + assert_eq!(input.len(), 2); + + let timestamp_ranges = RangeArray::try_new(extract_array(&input[0])?.to_data().into())?; + let value_ranges = RangeArray::try_new(extract_array(&input[1])?.to_data().into())?; + let len = timestamp_ranges.len(); + if len != value_ranges.len() { + return Err(DataFusionError::Execution(format!( + "RangeArray have different lengths in PromQL function {name}: array1={len}, array2={}", + value_ranges.len() + ))); + } + + if timestamp_ranges.is_empty() { + return Ok(ColumnarValue::Array(Arc::new(Float64Array::from_iter( + std::iter::empty::>(), + )))); + } + + // The generic range-function wrapper downcasts both arrays for each window. + // Do it once after retaining its zero-window behavior above. + assert!( + timestamp_ranges + .values() + .as_any() + .is::() + ); + let values = value_ranges + .values() + .as_any() + .downcast_ref::() + .unwrap(); + let evaluator: fn(&Float64Array, usize, usize) -> Option = if values.null_count() == 0 { + match operation { + PresenceEvaluator::Count => |_, _, length| (length != 0).then_some(length as f64), + PresenceEvaluator::Last => { + |values, offset, length| (length != 0).then(|| values.value(offset + length - 1)) + } + PresenceEvaluator::Absent => |_, _, length| (length == 0).then_some(1.0), + PresenceEvaluator::Present => |_, _, length| (length != 0).then_some(1.0), + } + } else { + match operation { + PresenceEvaluator::Count => |values, offset, length| { + let count = (offset..offset + length) + .filter(|&index| values.is_valid(index)) + .count(); + (count != 0).then_some(count as f64) + }, + PresenceEvaluator::Last => |values, offset, length| { + (offset..offset + length) + .rev() + .find(|&index| values.is_valid(index)) + .map(|index| values.value(index)) + }, + PresenceEvaluator::Absent => |values, offset, length| { + (!(offset..offset + length).any(|index| values.is_valid(index))).then_some(1.0) + }, + PresenceEvaluator::Present => |values, offset, length| { + (offset..offset + length) + .any(|index| values.is_valid(index)) + .then_some(1.0) + }, + } + }; + + let mut result = Vec::with_capacity(len); + for index in 0..len { + let (_, timestamp_length) = timestamp_ranges.get_offset_length(index).unwrap(); + let (value_offset, value_length) = value_ranges.get_offset_length(index).unwrap(); + if timestamp_length != value_length { + return Err(DataFusionError::Execution(format!( + "RangeArray's element {index} have different lengths in PromQL function {name}: array1={timestamp_length}, array2={value_length}" + ))); + } + result.push(evaluator(values, value_offset, value_length)); + } + + Ok(ColumnarValue::Array(Arc::new(Float64Array::from_iter( + result, + )))) +} + /// The average value of all points in the specified interval. #[range_fn( name = AvgOverTime, @@ -33,7 +156,10 @@ use crate::range_array::RangeArray; display_name = prom_avg_over_time )] pub fn avg_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> Option { - compute::sum(values).map(|result| result / values.len() as f64) + // `sum` already skips null slots and yields `None` for an all-null window, so only the + // divisor needs to count samples instead of slots. + let sample_count = values.len() - values.null_count(); + compute::sum(values).map(|result| result / sample_count as f64) } /// The minimum value of all points in the specified interval. @@ -84,24 +210,23 @@ pub fn sum_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> Op #[range_fn( name = CountOverTime, ret = Float64Array, - display_name = prom_count_over_time + display_name = prom_count_over_time, + evaluator = count_over_time_evaluator )] pub fn count_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> Option { - if values.is_empty() { - None - } else { - Some(values.len() as f64) - } + let count = values.iter().flatten().count(); + (count != 0).then_some(count as f64) } /// The most recent point value in specified interval. #[range_fn( name = LastOverTime, ret = Float64Array, - display_name = prom_last_over_time + display_name = prom_last_over_time, + evaluator = last_over_time_evaluator )] pub fn last_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> Option { - values.values().last().copied() + values.iter().flatten().last() } /// absent_over_time returns an empty vector if the range vector passed to it has any @@ -110,20 +235,22 @@ pub fn last_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> O #[range_fn( name = AbsentOverTime, ret = Float64Array, - display_name = prom_absent_over_time + display_name = prom_absent_over_time, + evaluator = absent_over_time_evaluator )] pub fn absent_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> Option { - if values.is_empty() { Some(1.0) } else { None } + values.iter().flatten().next().is_none().then_some(1.0) } /// the value 1 for any series in the specified interval. #[range_fn( name = PresentOverTime, ret = Float64Array, - display_name = prom_present_over_time + display_name = prom_present_over_time, + evaluator = present_over_time_evaluator )] pub fn present_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> Option { - if values.is_empty() { None } else { Some(1.0) } + values.iter().flatten().next().is_some().then_some(1.0) } /// the population standard variance of the values in the specified interval. @@ -135,26 +262,21 @@ pub fn present_over_time(_: &TimestampMillisecondArray, values: &Float64Array) - display_name = prom_stdvar_over_time )] pub fn stdvar_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> Option { - if values.is_empty() { - None - } else { - let mut count = 0; - let mut mean: f64 = 0.0; - let mut result: f64 = 0.0; - for value in values { - let value = value.unwrap(); - let new_count = count + 1; - let delta1 = value - mean; - let new_mean = delta1 / new_count as f64 + mean; - let delta2 = value - new_mean; - let new_result = result + delta1 * delta2; + let mut count = 0; + let mut mean: f64 = 0.0; + let mut result: f64 = 0.0; + for value in values.iter().flatten() { + let new_count = count + 1; + let delta1 = value - mean; + let new_mean = delta1 / new_count as f64 + mean; + let delta2 = value - new_mean; + let new_result = result + delta1 * delta2; - count += 1; - mean = new_mean; - result = new_result; - } - Some(result / count as f64) + count = new_count; + mean = new_mean; + result = new_result; } + (count > 0).then(|| result / count as f64) } /// the population standard deviation of the values in the specified interval. @@ -165,35 +287,32 @@ pub fn stdvar_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> display_name = prom_stddev_over_time )] pub fn stddev_over_time(_: &TimestampMillisecondArray, values: &Float64Array) -> Option { - if values.is_empty() { - None - } else { - let mut count = 0.0; - let mut mean = 0.0; - let mut comp_mean = 0.0; - let mut deviations_sum_sq = 0.0; - let mut comp_deviations_sum_sq = 0.0; - for v in values { - count += 1.0; - let current_value = v.unwrap(); - let delta = current_value - (mean + comp_mean); - let (new_mean, new_comp_mean) = compensated_sum_inc(delta / count, mean, comp_mean); - mean = new_mean; - comp_mean = new_comp_mean; - let (new_deviations_sum_sq, new_comp_deviations_sum_sq) = compensated_sum_inc( - delta * (current_value - (mean + comp_mean)), - deviations_sum_sq, - comp_deviations_sum_sq, - ); - deviations_sum_sq = new_deviations_sum_sq; - comp_deviations_sum_sq = new_comp_deviations_sum_sq; - } - Some(((deviations_sum_sq + comp_deviations_sum_sq) / count).sqrt()) + let mut count = 0.0; + let mut mean = 0.0; + let mut comp_mean = 0.0; + let mut deviations_sum_sq = 0.0; + let mut comp_deviations_sum_sq = 0.0; + for current_value in values.iter().flatten() { + count += 1.0; + let delta = current_value - (mean + comp_mean); + let (new_mean, new_comp_mean) = compensated_sum_inc(delta / count, mean, comp_mean); + mean = new_mean; + comp_mean = new_comp_mean; + let (new_deviations_sum_sq, new_comp_deviations_sum_sq) = compensated_sum_inc( + delta * (current_value - (mean + comp_mean)), + deviations_sum_sq, + comp_deviations_sum_sq, + ); + deviations_sum_sq = new_deviations_sum_sq; + comp_deviations_sum_sq = new_comp_deviations_sum_sq; } + (count > 0.0).then(|| ((deviations_sum_sq + comp_deviations_sum_sq) / count).sqrt()) } #[cfg(test)] mod test { + use datafusion::arrow::buffer::NullBuffer; + use super::*; use crate::functions::test_util::simple_range_udf_runner; @@ -218,6 +337,315 @@ mod test { assert_over_time_value(max_over_time(×tamps, &values), expected_max); } + fn special_ranges() -> (RangeArray, RangeArray) { + use datafusion::arrow::buffer::NullBuffer; + + let timestamps = Arc::new(TimestampMillisecondArray::from_iter_values(0..10)).slice(1, 8); + let values = Arc::new(Float64Array::new( + vec![ + 99.0, + -0.0, + 0.0, + f64::INFINITY, + f64::NEG_INFINITY, + f64::from_bits(0x7ff8_0000_0000_0042), + f64::from_bits(0x7ff8_0000_0000_0066), + 3.0, + -7.0, + 42.0, + ] + .into(), + Some(NullBuffer::from(vec![ + true, true, true, true, true, false, true, true, true, true, + ])), + )) + .slice(1, 8); + // Empty, singleton, overlapping, disjoint, repeated, and backward windows use + // independent timestamp and value offsets. + let timestamp_ranges = [ + (0, 0), + (1, 1), + (2, 1), + (3, 1), + (4, 1), + (5, 1), + (6, 1), + (0, 3), + (3, 2), + (4, 2), + (6, 2), + (5, 1), + (2, 1), + ]; + let value_ranges = [ + (7, 0), + (0, 1), + (1, 1), + (2, 1), + (3, 1), + (4, 1), + (5, 1), + (0, 3), + (3, 2), + (4, 2), + (1, 2), + (5, 1), + (1, 1), + ]; + + let timestamp_ranges = RangeArray::from_ranges( + Arc::new(timestamps), + std::iter::once((0, 0)).chain(timestamp_ranges), + ) + .unwrap() + .into_dict() + .slice(1, timestamp_ranges.len()); + let value_ranges = RangeArray::from_ranges( + Arc::new(values), + std::iter::once((0, 0)).chain(value_ranges), + ) + .unwrap() + .into_dict() + .slice(1, value_ranges.len()); + + ( + RangeArray::try_new(timestamp_ranges).unwrap(), + RangeArray::try_new(value_ranges).unwrap(), + ) + } + + fn assert_specialized_matches_oracle( + udf: ScalarUDF, + kernel: fn(&TimestampMillisecondArray, &Float64Array) -> Option, + ) { + use crate::functions::test_util::invoke_range_udf; + + let (timestamps, values) = special_ranges(); + let expected = (0..timestamps.len()) + .map(|index| { + let timestamps = timestamps.get(index).unwrap(); + let values = values.get(index).unwrap(); + kernel( + timestamps + .as_any() + .downcast_ref::() + .unwrap(), + values.as_any().downcast_ref::().unwrap(), + ) + }) + .collect::>(); + let (timestamps, values) = special_ranges(); + let output = invoke_range_udf(udf, timestamps, values).unwrap(); + let output_array = extract_array(&output).unwrap(); + let output = output_array + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(output.len(), expected.len()); + for (index, expected) in expected.into_iter().enumerate() { + assert_eq!(output.is_null(index), expected.is_none()); + if let Some(expected) = expected { + assert_eq!(output.value(index).to_bits(), expected.to_bits()); + } + } + } + + #[test] + fn specialized_presence_range_udfs_match_slice_oracles() { + assert_specialized_matches_oracle(CountOverTime::scalar_udf(), count_over_time); + assert_specialized_matches_oracle(LastOverTime::scalar_udf(), last_over_time); + assert_specialized_matches_oracle(AbsentOverTime::scalar_udf(), absent_over_time); + assert_specialized_matches_oracle(PresentOverTime::scalar_udf(), present_over_time); + } + + #[test] + fn specialized_presence_range_udfs_ignore_null_samples() { + use datafusion::arrow::buffer::NullBuffer; + + let make_ranges = || { + let timestamps = + Arc::new(TimestampMillisecondArray::from_iter_values(0..9)).slice(1, 7); + let values = Arc::new(Float64Array::new( + vec![99.0, 1.0, 0.0, 0.0, 4.0, f64::NAN, 0.0, -0.0, 42.0].into(), + Some(NullBuffer::from(vec![ + true, true, false, false, true, true, true, true, true, + ])), + )) + .slice(1, 7); + let ranges = [(0, 4), (0, 3), (1, 2), (4, 1), (5, 2), (7, 0)]; + + ( + RangeArray::from_ranges(Arc::new(timestamps), ranges).unwrap(), + RangeArray::from_ranges(Arc::new(values), ranges).unwrap(), + ) + }; + // The first window is [1, NULL, NULL, 4]; the old physical-length evaluator + // incorrectly returned 4 for count_over_time. + let cases = [ + ( + CountOverTime::scalar_udf(), + [Some(2.0), Some(1.0), None, Some(1.0), Some(2.0), None], + ), + ( + LastOverTime::scalar_udf(), + [Some(4.0), Some(1.0), None, Some(f64::NAN), Some(-0.0), None], + ), + ( + AbsentOverTime::scalar_udf(), + [None, None, Some(1.0), None, None, Some(1.0)], + ), + ( + PresentOverTime::scalar_udf(), + [Some(1.0), Some(1.0), None, Some(1.0), Some(1.0), None], + ), + ]; + + for (udf, expected) in cases { + let (timestamps, values) = make_ranges(); + let output = + crate::functions::test_util::invoke_range_udf(udf, timestamps, values).unwrap(); + let output_array = extract_array(&output).unwrap(); + let output = output_array + .as_any() + .downcast_ref::() + .unwrap(); + + for (index, expected) in expected.into_iter().enumerate() { + assert_eq!(output.is_valid(index), expected.is_some()); + if let Some(expected) = expected { + assert_eq!(output.value(index).to_bits(), expected.to_bits()); + } + } + } + } + + #[test] + fn specialized_presence_range_udfs_preserve_errors_and_metadata() { + use datafusion::arrow::array::{DictionaryArray, Int64Array}; + use datafusion::arrow::datatypes::{Field, Int64Type}; + use datafusion_common::config::ConfigOptions; + use datafusion_expr::ScalarFunctionArgs; + + use crate::functions::test_util::{assert_execution_error, invoke_range_udf}; + + for (udf, name) in [ + (CountOverTime::scalar_udf(), "prom_count_over_time"), + (LastOverTime::scalar_udf(), "prom_last_over_time"), + (AbsentOverTime::scalar_udf(), "prom_absent_over_time"), + (PresentOverTime::scalar_udf(), "prom_present_over_time"), + ] { + assert_eq!(udf.name(), name); + assert_eq!(udf.signature().volatility, Volatility::Volatile); + } + + let timestamps = Arc::new(TimestampMillisecondArray::from_iter_values(0..3)); + let values = Arc::new(Float64Array::from_iter_values([1.0, 2.0, 3.0])); + let error = invoke_range_udf( + CountOverTime::scalar_udf(), + RangeArray::from_ranges(timestamps.clone(), [(0, 1), (1, 1)]).unwrap(), + RangeArray::from_ranges(values.clone(), [(0, 1)]).unwrap(), + ) + .unwrap_err(); + assert_execution_error( + error, + "RangeArray have different lengths in PromQL function prom_count_over_time: array1=2, array2=1", + ); + + let error = invoke_range_udf( + CountOverTime::scalar_udf(), + RangeArray::from_ranges(timestamps.clone(), [(0, 1), (1, 2)]).unwrap(), + RangeArray::from_ranges(values.clone(), [(0, 1), (1, 1)]).unwrap(), + ) + .unwrap_err(); + assert_execution_error( + error, + "RangeArray's element 1 have different lengths in PromQL function prom_count_over_time: array1=2, array2=1", + ); + + let invoke_dict = |timestamps: DictionaryArray, + values: DictionaryArray| { + let args = vec![ + ColumnarValue::Array(Arc::new(timestamps)), + ColumnarValue::Array(Arc::new(values)), + ]; + CountOverTime::scalar_udf().invoke_with_args(ScalarFunctionArgs { + arg_fields: args + .iter() + .enumerate() + .map(|(index, value)| { + Arc::new(Field::new( + format!("c{index}"), + value.data_type().clone(), + true, + )) + }) + .collect(), + args, + number_rows: 1, + return_field: Arc::new(Field::new("out", DataType::Float64, true)), + config_options: Arc::new(ConfigOptions::default()), + }) + }; + let empty = invoke_dict( + DictionaryArray::new( + Int64Array::from(Vec::::new()), + Arc::new(Float64Array::from_iter_values([1.0])), + ), + DictionaryArray::new( + Int64Array::from(Vec::::new()), + Arc::new(TimestampMillisecondArray::from_iter_values([0])), + ), + ) + .unwrap(); + assert!(extract_array(&empty).unwrap().is_empty()); + + let null_keys = DictionaryArray::new( + Int64Array::from(vec![None]), + Arc::new(TimestampMillisecondArray::from_iter_values([0])), + ); + let values_range = RangeArray::from_ranges(values.clone(), [(0, 1)]).unwrap(); + assert_eq!( + invoke_dict(null_keys, values_range.into_dict()) + .unwrap_err() + .to_string(), + "External error: Empty range is not expected" + ); + + let timestamps_range = RangeArray::from_ranges(timestamps, [(0, 1)]).unwrap(); + let invalid_values = unsafe { RangeArray::from_ranges_unchecked(values, [(2, 2)]) }; + assert_eq!( + invoke_dict(timestamps_range.into_dict(), invalid_values.into_dict()) + .unwrap_err() + .to_string(), + "External error: Illegal range: offset 2, length 2, array len 3" + ); + + let output = invoke_range_udf( + SumOverTime::scalar_udf(), + RangeArray::from_ranges( + Arc::new(TimestampMillisecondArray::from_iter_values([0, 1])), + [(0, 2)], + ) + .unwrap(), + RangeArray::from_ranges( + Arc::new(Float64Array::from_iter_values([1.0, 2.0])), + [(0, 2)], + ) + .unwrap(), + ) + .unwrap(); + assert_eq!( + extract_array(&output) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 3.0 + ); + } + #[test] fn min_max_over_time_ignore_ordinary_nan_when_finite_values_exist() { let ordinary_nan = f64::from_bits(0x7ff8_0000_0000_0000); @@ -563,4 +991,57 @@ mod test { vec![Some(0.0), Some(3.249615361854384)], ); } + + /// Timestamps and value ranges shared by the null-sample assertions below. + fn null_sample_range_arrays() -> (RangeArray, RangeArray) { + let ts_array = Arc::new(TimestampMillisecondArray::from_iter_values([ + 0i64, 1000, 2000, 3000, + ])); + // Samples are 2.0@0 and 8.0@3000; the null slots keep a payload that would skew every + // aggregate if it were read. + let values_array = Arc::new(Float64Array::new( + vec![2.0, 1000.0, -1000.0, 8.0].into(), + Some(NullBuffer::from_iter([true, false, false, true])), + )); + // The second window holds no sample at all. + let ranges = [(0, 4), (1, 2)]; + + ( + RangeArray::from_ranges(ts_array, ranges).unwrap(), + RangeArray::from_ranges(values_array, ranges).unwrap(), + ) + } + + #[test] + fn avg_over_time_divides_by_sample_count() { + let (ts_array, value_array) = null_sample_range_arrays(); + simple_range_udf_runner( + AvgOverTime::scalar_udf(), + ts_array, + value_array, + vec![], + vec![Some(5.0), None], + ); + } + + #[test] + fn stdvar_and_stddev_over_time_skip_null_samples() { + let (ts_array, value_array) = null_sample_range_arrays(); + simple_range_udf_runner( + StdvarOverTime::scalar_udf(), + ts_array, + value_array, + vec![], + vec![Some(9.0), None], + ); + + let (ts_array, value_array) = null_sample_range_arrays(); + simple_range_udf_runner( + StddevOverTime::scalar_udf(), + ts_array, + value_array, + vec![], + vec![Some(3.0), None], + ); + } } diff --git a/src/promql/src/functions/changes.rs b/src/promql/src/functions/changes.rs index c143d29bef7..8630c914a8f 100644 --- a/src/promql/src/functions/changes.rs +++ b/src/promql/src/functions/changes.rs @@ -43,8 +43,8 @@ mod test { use super::*; use crate::functions::test_util::{ - self, STALE_NAN, TinyPrng, assert_execution_error, build_test_range_arrays, - invoke_range_udf, simple_range_udf_runner, + self, STALE_NAN, assert_execution_error, build_test_range_arrays, invoke_range_udf, + simple_range_udf_runner, }; use crate::range_array::RangeArray; @@ -151,43 +151,17 @@ mod test { changes_oracle, Changes::scalar_udf(), ); - assert_eq!(expected, vec![Some(2.0), Some(0.0)]); + assert_eq!(expected, vec![Some(0.0), None]); } #[test] fn changes_range_array_seeded_differential() { - let mut prng = TinyPrng(0x2f6e_2b1d_834a_90c5); - let raw_values = (0..48) - .map(|_| match prng.next_index(12) { - 0 => -0.0, - 1 => 0.0, - 2 => -2.0, - 3 => -1.0, - 4 => 1.0, - 5 => 2.0, - 6 => f64::INFINITY, - 7 => f64::NEG_INFINITY, - 8 | 9 => f64::NAN, - _ => STALE_NAN, - }) - .collect::>(); - let values = raw_values.iter().copied().map(Some).collect(); - let mut timestamp_ranges = Vec::new(); - let mut value_ranges = Vec::new(); - for _ in 0..32 { - let length = prng.next_index(13) as u32; - timestamp_ranges.push((prng.next_index(65 - length as usize) as u32, length)); - value_ranges.push((prng.next_index(49 - length as usize) as u32, length)); - } + test_util::run_seeded_differential(changes_oracle, Changes::scalar_udf(), false); + } - test_util::run_oracle_ranges( - values, - raw_values, - timestamp_ranges, - value_ranges, - changes_oracle, - Changes::scalar_udf(), - ); + #[test] + fn changes_range_array_seeded_differential_with_nulls() { + test_util::run_seeded_differential(changes_oracle, Changes::scalar_udf(), true); } #[test] @@ -263,6 +237,8 @@ mod test { vec![Some(0.0), Some(0.0), Some(0.0)], ); + // The raw payload under the null slot would look like two changes; the sample sequence + // is 10 -> 10, which is none. The last window holds no sample at all. let values = Arc::new(Float64Array::new( vec![10.0, 7.0, 10.0].into(), Some(NullBuffer::from(vec![true, false, true])), @@ -276,10 +252,10 @@ mod test { ])); simple_range_udf_runner( Changes::scalar_udf(), - RangeArray::from_ranges(timestamps, [(0, 3)]).unwrap(), - RangeArray::from_ranges(values, [(0, 3)]).unwrap(), + RangeArray::from_ranges(timestamps, [(0, 3), (1, 1)]).unwrap(), + RangeArray::from_ranges(values, [(0, 3), (1, 1)]).unwrap(), vec![], - vec![Some(2.0)], + vec![Some(0.0), None], ); } } diff --git a/src/promql/src/functions/double_exponential_smoothing.rs b/src/promql/src/functions/double_exponential_smoothing.rs index b6768d47a11..2f8e0235d53 100644 --- a/src/promql/src/functions/double_exponential_smoothing.rs +++ b/src/promql/src/functions/double_exponential_smoothing.rs @@ -240,8 +240,6 @@ fn double_exponential_smoothing_impl(values: &[f64], sf: f64, tf: f64) -> Option return Some(f64::NAN); } - let values = values.to_vec(); - let mut s0 = 0.0; let mut s1 = values[0]; let mut b = values[1] - values[0]; @@ -353,6 +351,109 @@ mod tests { ); } + #[test] + fn test_double_exponential_smoothing_impl_copy_oracle() { + let normal_values = (0..240) + .map(|i| (i as f64 - 120.0) * 0.25) + .collect::>(); + let special_values = (0..240) + .map(|i| match i % 8 { + 0 => 0.0, + 1 => -0.0, + 2 => f64::INFINITY, + 3 => f64::NEG_INFINITY, + 4 => f64::NAN, + 5 => f64::from_bits(0x7ff8_0000_0000_0001), + 6 => 42.5, + _ => -42.5, + }) + .collect::>(); + let factors = [ + (0.0, 0.0), + (-0.0, 1.0), + (0.5, 0.1), + (1.0, 1.0), + (-0.5, 0.5), + (0.5, -0.5), + (1.5, 0.5), + (0.5, 1.5), + (f64::NAN, 0.5), + (0.5, f64::NAN), + (f64::INFINITY, 0.5), + (0.5, f64::INFINITY), + (f64::NEG_INFINITY, 0.5), + (0.5, f64::NEG_INFINITY), + ]; + + for (values_name, values) in [ + ("normal", normal_values.as_slice()), + ("special", special_values.as_slice()), + ] { + for len in [0, 1, 2, 3, 20, 240] { + let values = &values[..len]; + for (sf, tf) in factors { + let old = double_exponential_smoothing_impl_with_copy(values, sf, tf).unwrap(); + let new = double_exponential_smoothing_impl(values, sf, tf).unwrap(); + let case = format!("values={values_name}, len={len}, sf={sf:?}, tf={tf:?}"); + + if old.is_nan() || new.is_nan() { + assert!( + old.is_nan() && new.is_nan(), + "NaN mismatch for {case}: old={old:?}, new={new:?}" + ); + assert_eq!( + old.to_bits(), + new.to_bits(), + "NaN bit difference for {case}: old={:#018x}, new={:#018x}", + old.to_bits(), + new.to_bits(), + ); + } else { + assert_eq!( + old.to_bits(), + new.to_bits(), + "non-NaN bit difference for {case}: old={old:?}, new={new:?}" + ); + } + } + } + } + } + + fn double_exponential_smoothing_impl_with_copy( + values: &[f64], + sf: f64, + tf: f64, + ) -> Option { + if sf.is_nan() || tf.is_nan() || values.is_empty() { + return Some(f64::NAN); + } + if sf < 0.0 || tf < 0.0 { + return Some(f64::NEG_INFINITY); + } + if sf > 1.0 || tf > 1.0 { + return Some(f64::INFINITY); + } + + if values.len() <= 2 { + return Some(f64::NAN); + } + + let values = values.to_vec(); + let mut s0 = 0.0; + let mut s1 = values[0]; + let mut b = values[1] - values[0]; + + for (i, value) in values.iter().enumerate().skip(1) { + let x = sf * value; + b = calc_trend_value(i - 1, tf, s0, s1, b); + let y = (1.0 - sf) * (s1 + b); + s0 = s1; + s1 = x + y; + } + Some(s1) + } + #[test] fn test_prom_double_exponential_smoothing_monotonic() { let ranges = [(0, 5)]; @@ -450,7 +551,7 @@ mod tests { (ts_range_array, value_range_array) } - /// Converts a prometheus functions test series into a vector of f64 element with respect to resets and trend direction + /// Converts a prometheus functions test series into a vector of f64 element with respect to resets and trend direction /// The input example: "0+10x1000 100+30x1000" fn create_test_range_from_promql_series(input: &str) -> Vec { input.split(' ').map(parse_promql_series_entry).fold( diff --git a/src/promql/src/functions/edge_count.rs b/src/promql/src/functions/edge_count.rs index 9109c742073..7cd5907a5ee 100644 --- a/src/promql/src/functions/edge_count.rs +++ b/src/promql/src/functions/edge_count.rs @@ -81,7 +81,11 @@ fn calc( .unwrap(); let requested_edges = validate_windows(×tamp_ranges, &value_ranges, name)?; let raw_values = values.values(); - let direct = should_scan_direct(requested_edges, raw_values.len()); + // A NULL field value means the series has no sample at that timestamp. The prefix sums + // encode edges between physically adjacent slots, which no longer holds once nulls are + // skipped, so a null-bearing input falls back to scanning each window. + let has_nulls = values.null_count() > 0; + let direct = has_nulls || should_scan_direct(requested_edges, raw_values.len()); let prefix = (!direct).then(|| build_prefix(raw_values.as_ref(), kind)); let mut result = Vec::with_capacity(value_ranges.len()); @@ -90,6 +94,7 @@ fn calc( let end = checked_end(offset, len, index, name)?; let count = match len { 0 => None, + _ if has_nulls => count_edges_skipping_nulls(values, offset, end, kind), 1 => Some(0), _ if direct => Some(count_edges(raw_values.as_ref(), offset, end, kind)), _ => { @@ -172,6 +177,29 @@ fn count_edges(values: &[f64], offset: usize, end: usize, kind: EdgeKind) -> u64 count } +/// Counts edges between consecutive samples in `[offset, end)`, treating null slots as +/// absent. Returns `None` when the window holds no sample. +fn count_edges_skipping_nulls( + values: &Float64Array, + offset: usize, + end: usize, + kind: EdgeKind, +) -> Option { + let raw_values = values.values(); + let mut count = 0; + let mut previous = None; + for index in offset..end { + if values.is_null(index) { + continue; + } + let current = raw_values[index]; + if let Some(previous) = previous.replace(current) { + count += u64::from(is_edge(previous, current, kind)); + } + } + previous.is_some().then_some(count) +} + fn is_edge(previous: f64, current: f64, kind: EdgeKind) -> bool { match kind { EdgeKind::Changes => previous != current && !(previous.is_nan() && current.is_nan()), diff --git a/src/promql/src/functions/extrapolate_rate.rs b/src/promql/src/functions/extrapolate_rate.rs index af0714a8f79..263ae6337ba 100644 --- a/src/promql/src/functions/extrapolate_rate.rs +++ b/src/promql/src/functions/extrapolate_rate.rs @@ -30,6 +30,7 @@ //! Implementations of `rate`, `increase` and `delta` functions in PromQL. use std::fmt::Display; +use std::ops::Range; use std::sync::Arc; use datafusion::arrow::array::{Float64Array, Float64Builder, TimestampMillisecondArray}; @@ -180,99 +181,120 @@ impl ExtrapolatedRate() .expect("validated by extract_range_dict") .values(); - let all_values = value_dict + let value_array = value_dict .values() .as_any() .downcast_ref::() - .expect("validated by extract_range_dict") - .values(); + .expect("validated by extract_range_dict"); + // A NULL field value means the series has no sample at that timestamp, so the padding + // under a null slot must not be read. Skip the per-window null scan when the whole + // backing array is null-free, which is the common case. + let has_nulls = value_array.null_count() > 0; + let all_values = value_array.values(); let eval_ts = eval_ts_array.values(); let mut result_builder = Float64Builder::with_capacity(num_windows); let range_length = self.range_length; let range_length_secs = range_length as f64 / 1000.0; - let mut counter_correction = 0.0; - let mut prev_offset = usize::MAX; - let mut prev_length = 0usize; + // Range windows normally overlap heavily, so scanning every one for resets costs far + // more than a single pass over the values. Index the reset positions once that is the + // cheaper side, and stop counting as soon as the requested pairs pass that budget, + // which heavy overlap does within the first few windows. A short lookback with a long + // step is the shape that never reaches it, and there the per-window scans do win. + let mut reset_index = if IS_COUNTER && !has_nulls { + let budget = all_values.len().saturating_sub(1); + let mut scanned_pairs = 0usize; + keys.iter() + .any(|&key| { + scanned_pairs = + scanned_pairs.saturating_add(unpack(key).1.saturating_sub(1) as usize); + scanned_pairs > budget + }) + .then(|| CounterResetIndex::new(all_values)) + } else { + None + }; for index in 0..num_windows { let (raw_offset, raw_length) = unpack(keys[index]); let offset = raw_offset as usize; let length = raw_length as usize; - if length < 2 { + let end = offset + length; + let (first_index, last_index, sample_count) = if has_nulls { + match valid_window_bounds(value_array, offset, length) { + Some(bounds) => bounds, + None => { + result_builder.append_null(); + continue; + } + } + } else { + (offset, end.saturating_sub(1), length) + }; + + if sample_count < 2 { result_builder.append_null(); - prev_offset = usize::MAX; continue; } - let end = offset + length; - let first_value = all_values[offset]; - let last_value = all_values[end - 1]; + let first_value = all_values[first_index]; + let last_value = all_values[last_index]; - let result_value = if IS_COUNTER { - // Adjacent normalized windows usually slide forward by one sample. Reuse the - // previous window's accumulated reset correction and adjust only the dropped and - // newly added edges, falling back to a full scan when the layout changes. - if prev_offset != usize::MAX && offset == prev_offset + 1 && length == prev_length { - if all_values[prev_offset + 1] < all_values[prev_offset] { - counter_correction -= all_values[prev_offset]; - } - if all_values[end - 1] < all_values[end - 2] { - counter_correction += all_values[end - 2]; - } + let mut result_value = last_value - first_value; + if IS_COUNTER { + result_value = if has_nulls { + add_counter_resets_between_samples( + result_value, + value_array, + first_index, + last_index, + ) } else { - counter_correction = 0.0; - for pair in all_values[offset..end].windows(2) { - if pair[1] < pair[0] { - counter_correction += pair[0]; - } + match &mut reset_index { + Some(reset_index) => reset_index.add_resets(result_value, offset, end), + None => add_counter_resets(result_value, &all_values[offset..end]), } - } - last_value - first_value + counter_correction - } else { - last_value - first_value - }; + }; + } - prev_offset = offset; - prev_length = length; - - let first_ts = all_timestamps[offset]; - let last_ts = all_timestamps[end - 1]; + let first_ts = all_timestamps[first_index]; + let last_ts = all_timestamps[last_index]; let range_end = eval_ts[index]; let range_start = range_end - range_length; let sampled_interval_ms = (last_ts - first_ts) as f64; - let average_interval_ms = sampled_interval_ms / (length - 1) as f64; + let average_interval_ms = sampled_interval_ms / (sample_count - 1) as f64; let mut duration_to_start_ms = (first_ts - range_start) as f64; - let duration_to_end_ms = (range_end - last_ts) as f64; + let mut duration_to_end_ms = (range_end - last_ts) as f64; + let extrapolation_threshold = average_interval_ms * 1.1; - // Counters cannot be negative, so Prometheus allows the extrapolation window to snap - // back to the inferred zero point instead of extending into negative values. + // Mirror Prometheus extrapolation: extend to the real range boundary when a sample is + // close enough, otherwise only half an average sampling interval, which is the guess + // for where the series actually starts or ends. + if duration_to_start_ms >= extrapolation_threshold { + duration_to_start_ms = average_interval_ms / 2.0; + } + // Counters cannot be negative, so the extrapolation can snap back to the inferred + // zero point instead of extending into negative values. Prometheus applies this + // after the threshold clamp, so it can only shorten the leading extrapolation. if IS_COUNTER && result_value > 0.0 && first_value >= 0.0 { let duration_to_zero = sampled_interval_ms * (first_value / result_value); if duration_to_zero < duration_to_start_ms { duration_to_start_ms = duration_to_zero; } } - - let extrapolation_threshold = average_interval_ms * 1.1; - let mut extrapolated_interval_ms = sampled_interval_ms; - - // Mirror Prometheus extrapolation: extend to the real range boundary when a sample is - // close enough, otherwise add half an average sampling interval on that side. - if duration_to_start_ms < extrapolation_threshold { - extrapolated_interval_ms += duration_to_start_ms; - } else { - extrapolated_interval_ms += average_interval_ms / 2.0; - } - if duration_to_end_ms < extrapolation_threshold { - extrapolated_interval_ms += duration_to_end_ms; - } else { - extrapolated_interval_ms += average_interval_ms / 2.0; + if duration_to_end_ms >= extrapolation_threshold { + duration_to_end_ms = average_interval_ms / 2.0; } - let mut factor = extrapolated_interval_ms / sampled_interval_ms; + // Samples sharing one timestamp leave nothing to extrapolate over. + let mut factor = if sampled_interval_ms == 0.0 { + 1.0 + } else { + (sampled_interval_ms + duration_to_start_ms + duration_to_end_ms) + / sampled_interval_ms + }; if IS_RATE { factor /= range_length_secs; @@ -286,6 +308,151 @@ impl ExtrapolatedRate f64 { + values + .windows(2) + .filter(|pair| pair[1] < pair[0]) + .fold(result, |result, pair| result + pair[0]) +} + +/// Positions of the counter resets in a value array, so that a window can accumulate the +/// resets it contains instead of scanning all of its samples. +struct CounterResetIndex<'a> { + values: &'a [f64], + /// Ascending indices `i` where `values[i] < values[i - 1]`. + positions: Vec, + /// Slice of `positions` covered by the last window. + active: Range, + /// That window, so the next one can tell whether it advanced. + previous: Range, + /// `positions[active.start]`: the reset a later `start` would drop. `usize::MAX` when the + /// active slice reaches the end of `positions`. + drops_at: usize, + /// `positions[active.end]`: the reset a later `end` would gain, saturated the same way. + gains_at: usize, +} + +impl<'a> CounterResetIndex<'a> { + fn new(values: &'a [f64]) -> Self { + let positions: Vec = (1..values.len()) + .filter(|&i| values[i] < values[i - 1]) + .collect(); + let first = positions.first().copied().unwrap_or(usize::MAX); + Self { + values, + positions, + active: 0..0, + previous: 0..0, + drops_at: first, + gains_at: first, + } + } + + /// Same additions [`add_counter_resets`] performs over `values[start..end]`, in the same + /// order, reached through the index instead of by scanning the window. + #[inline] + fn add_resets(&mut self, result: f64, start: usize, end: usize) -> f64 { + // The active slice only stays put if the window advanced without reaching either of + // the resets that bound it. + if start < self.previous.start + || end < self.previous.end + || start >= self.drops_at + || end > self.gains_at + { + self.locate(start, end); + } + self.previous = start..end; + + if self.active.start == self.active.end { + // A counter that has not reset inside this window, which is the normal case, would + // otherwise pay a range bounds check and an empty iterator for nothing. + return result; + } + + let values = self.values; + self.positions[self.active.start..self.active.end] + .iter() + .fold(result, |result, &i| result + values[i - 1]) + } + + fn locate(&mut self, start: usize, end: usize) { + // Walk the bounds forward from the previous window and only search when they move + // back. On a series that resets often the searches cost more than the additions they + // locate, because they run deep and a window holds a handful of resets. + let (left, right) = if start < self.previous.start || end < self.previous.end { + ( + self.positions.partition_point(|&i| i <= start), + self.positions.partition_point(|&i| i < end), + ) + } else { + let mut left = self.active.start; + while left < self.positions.len() && self.positions[left] <= start { + left += 1; + } + let mut right = self.active.end.max(left); + while right < self.positions.len() && self.positions[right] < end { + right += 1; + } + (left, right) + }; + self.active = left..right; + self.drops_at = self.positions.get(left).copied().unwrap_or(usize::MAX); + self.gains_at = self.positions.get(right).copied().unwrap_or(usize::MAX); + } +} + +/// Same additions [`add_counter_resets`] performs, over the samples in `[first, last]` instead +/// of over every slot. +fn add_counter_resets_between_samples( + result: f64, + values: &Float64Array, + first: usize, + last: usize, +) -> f64 { + let raw_values = values.values(); + let mut result = result; + let mut previous = raw_values[first]; + for index in first + 1..=last { + if values.is_null(index) { + continue; + } + let current = raw_values[index]; + if current < previous { + result += previous; + } + previous = current; + } + result +} + +/// Locates the samples inside `[offset, offset + length)`, returning the first and last +/// non-null index together with the number of non-null slots. Returns `None` when the +/// window holds no sample. +fn valid_window_bounds( + values: &Float64Array, + offset: usize, + length: usize, +) -> Option<(usize, usize, usize)> { + let mut first = None; + let mut last = 0; + let mut count = 0; + for index in offset..offset + length { + if values.is_null(index) { + continue; + } + first.get_or_insert(index); + last = index; + count += 1; + } + first.map(|first| (first, last, count)) +} + fn extract_eval_timestamps( columnar_value: &ColumnarValue, func_name: &str, @@ -358,9 +525,11 @@ impl Display for ExtrapolatedRate { mod test { use datafusion::arrow::array::ArrayRef; + use datafusion::arrow::buffer::NullBuffer; use datafusion_common::ScalarValue; use super::*; + use crate::functions::test_util::TinyPrng; /// Range length is fixed to 5 fn extrapolated_rate_runner( @@ -409,6 +578,366 @@ mod test { ) } + /// Evaluates `ranges` as one batch and asserts every window is bit-identical to evaluating + /// that window on its own, which always takes the direct per-window reduction. + fn assert_counter_windows_match_single(values: &[f64], ranges: &[(u32, u32)]) { + let timestamps = Arc::new(TimestampMillisecondArray::from_iter_values( + (0..values.len()).map(|i| i as i64 * 30_000 + (i % 5) as i64 * 1_000), + )); + let values = Arc::new(Float64Array::from(values.to_vec())); + let evaluate = |ranges: &[(u32, u32)]| { + let eval_ts = Arc::new(TimestampMillisecondArray::from_iter_values( + ranges.iter().map(|&(offset, length)| { + timestamps.value((offset + length.saturating_sub(1)) as usize) + 5_000 + }), + )); + let input = [ + ColumnarValue::Array(Arc::new( + RangeArray::from_ranges(timestamps.clone(), ranges.iter().copied()) + .unwrap() + .into_dict(), + )), + ColumnarValue::Array(Arc::new( + RangeArray::from_ranges(values.clone(), ranges.iter().copied()) + .unwrap() + .into_dict(), + )), + ColumnarValue::Array(eval_ts), + ColumnarValue::Scalar(ScalarValue::Int64(Some(3_600_000))), + ]; + extract_array(&Rate::new(3_600_000).calc(&input).unwrap()) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .collect::>() + }; + + for (range, batched) in ranges.iter().zip(evaluate(ranges)) { + let single = evaluate(std::slice::from_ref(range))[0]; + match (batched, single) { + (None, None) => {} + (Some(batched), Some(single)) => assert!( + batched.to_bits() == single.to_bits() || (batched.is_nan() && single.is_nan()), + "range {range:?}: batched {batched} != single {single}" + ), + _ => panic!("range {range:?}: batched {batched:?} != single {single:?}"), + } + } + } + + #[test] + fn counter_resets_accumulate_into_the_running_result() { + // Summed on their own, 1e16 and 1.0 round to 1e16, which then cancels against the + // first sample and reports no increase at all. Folding each reset into `last - first` + // as Prometheus does keeps the 1.0. Both paths detect the same two resets. + let ts_array = Arc::new(TimestampMillisecondArray::from_iter( + [1, 2, 3, 4].into_iter().map(Some), + )); + let values_array = Arc::new(Float64Array::from_iter([1e16, 1.0, 0.0, 1.0])); + let ranges = [(0, 4)]; + let ts_range = RangeArray::from_ranges(ts_array, ranges).unwrap(); + let value_range = RangeArray::from_ranges(values_array, ranges).unwrap(); + let timestamps = Arc::new(TimestampMillisecondArray::from_iter([Some(4)])) as _; + + extrapolated_rate_runner::( + ts_range, + value_range, + timestamps, + vec![1.1666666666666667], + ); + } + + #[test] + fn counter_correction_survives_huge_and_infinite_resets() { + // Both series reset twice inside the first window and once inside the second, but the + // first reset is large enough to swallow the second one when they are summed together. + for values in [ + vec![1e16, 1.0, 0.0, 1.0, 2.0], + vec![f64::INFINITY, 1.0, 0.0, 1.0, 2.0], + ] { + assert_counter_windows_match_single(&values, &[(0, 4), (1, 4)]); + } + } + + #[test] + fn counter_correction_matches_single_window_on_irregular_layouts() { + let mut values: Vec = (0..512).map(|i| (i % 37) as f64 * 0.25).collect(); + values[20] = f64::NAN; + values[70] = f64::INFINITY; + values[140] = f64::NEG_INFINITY; + values[220] = 1e300; + values[221] = 1e-200; + + // A stride of one walks both bounds across the reset positions, which sit every 37 + // samples: `start` reaches one when that reset has to leave the window, `end` when the + // next one has to stay out of it, both while the cached bounds are live. + let mut ranges: Vec<(u32, u32)> = (0..390).map(|i| (i, 120)).collect(); + // Empty, too-short, backward and disjoint windows all break a forward-only slide. + ranges.extend([(400, 0), (400, 1), (2, 20), (450, 30), (0, 120)]); + ranges.extend((0..390).rev().step_by(10).map(|i| (i, 120))); + + assert_counter_windows_match_single(&values, &ranges); + } + + /// Builds a value array whose null slots keep a distinguishable raw payload, so a + /// function that reads the padding instead of the samples produces a different result. + fn values_with_nulls(values: Vec>, padding: f64) -> Arc { + let raw = values + .iter() + .map(|value| value.unwrap_or(padding)) + .collect::>(); + Arc::new(Float64Array::new( + raw.into(), + Some(NullBuffer::from_iter( + values.iter().map(|value| value.is_some()), + )), + )) + } + + fn nullable_rate_runner( + timestamps: Vec, + values: Arc, + ranges: Vec<(u32, u32)>, + eval_timestamps: Vec, + range_length: i64, + ) -> Vec> { + let ts_array = Arc::new(TimestampMillisecondArray::from_iter_values(timestamps)); + let ts_range = RangeArray::from_ranges(ts_array, ranges.clone()).unwrap(); + let value_range = RangeArray::from_ranges(values, ranges).unwrap(); + let input = vec![ + ColumnarValue::Array(Arc::new(ts_range.into_dict())), + ColumnarValue::Array(Arc::new(value_range.into_dict())), + ColumnarValue::Array(Arc::new(TimestampMillisecondArray::from_iter_values( + eval_timestamps, + ))), + ColumnarValue::Array(Arc::new(Int64Array::from(vec![range_length]))), + ]; + let output = extract_array( + &ExtrapolatedRate::::new(range_length) + .calc(&input) + .unwrap(), + ) + .unwrap(); + let output = output.as_any().downcast_ref::().unwrap(); + output.iter().collect() + } + + #[test] + fn rate_uses_samples_not_null_padding() { + // Samples are 1.0@0 and 4.0@3000; the padding under the null slots would add two more. + let output = nullable_rate_runner::( + vec![0, 1000, 2000, 3000], + values_with_nulls(vec![Some(1.0), None, None, Some(4.0)], 99.0), + vec![(0, 4)], + vec![3000], + 4000, + ); + + assert_eq!(output, vec![Some(1.0)]); + } + + #[test] + fn rate_returns_null_for_windows_without_enough_samples() { + let output = nullable_rate_runner::( + vec![0, 1000, 2000], + values_with_nulls(vec![None, Some(2.0), None], 7.0), + vec![(0, 3), (0, 2), (2, 1)], + vec![2000, 2000, 2000], + 4000, + ); + + assert_eq!(output, vec![None, None, None]); + } + + #[test] + fn increase_corrects_counter_reset_between_samples() { + // The sample sequence is 5.0 -> 3.0, one reset. Reading the padding would see + // 5.0 -> 100.0 -> 3.0 and charge the correction against the wrong value. + let output = nullable_rate_runner::( + vec![0, 1000, 2000], + values_with_nulls(vec![Some(5.0), None, Some(3.0)], 100.0), + vec![(0, 3)], + vec![2000], + 2000, + ); + + assert_eq!(output, vec![Some(3.0)]); + } + + #[test] + fn delta_extrapolates_from_sample_timestamps() { + // The window spans (-1000, 3000] but its samples only cover 1000..2000, so the + // extrapolation adds half an average interval on the leading side. + let output = nullable_rate_runner::( + vec![0, 1000, 2000, 3000], + values_with_nulls(vec![None, Some(2.0), Some(5.0), None], 42.0), + vec![(0, 4)], + vec![3000], + 4000, + ); + + assert_eq!(output, vec![Some(7.5)]); + } + + /// Line-by-line port of Prometheus `extrapolatedRate` (promql/functions.go), float path + /// without start timestamps. Kept as a second implementation so that the order of the + /// threshold clamp and the counter zero-snap stays pinned to the upstream one. + fn prometheus_extrapolated_rate( + timestamps: &[i64], + values: &[f64], + eval_ts: i64, + range_ms: i64, + is_counter: bool, + is_rate: bool, + ) -> Option { + if values.len() < 2 { + return None; + } + let num_samples_minus_one = values.len() - 1; + let first_t = timestamps[0]; + let last_t = timestamps[num_samples_minus_one]; + let mut result = values[num_samples_minus_one] - values[0]; + if is_counter { + for index in 1..values.len() { + if values[index] < values[index - 1] { + result += values[index - 1]; + } + } + } + + let range_start = eval_ts - range_ms; + let mut duration_to_start = (first_t - range_start) as f64 / 1000.0; + let mut duration_to_end = (eval_ts - last_t) as f64 / 1000.0; + let sampled_interval = (last_t - first_t) as f64 / 1000.0; + let average_duration_between_samples = sampled_interval / num_samples_minus_one as f64; + let extrapolation_threshold = average_duration_between_samples * 1.1; + + if duration_to_start >= extrapolation_threshold { + duration_to_start = average_duration_between_samples / 2.0; + } + if is_counter { + let mut duration_to_zero = duration_to_start; + if result > 0.0 && values[0] >= 0.0 { + duration_to_zero = sampled_interval * (values[0] / result); + } + if duration_to_zero < duration_to_start { + duration_to_start = duration_to_zero; + } + } + if duration_to_end >= extrapolation_threshold { + duration_to_end = average_duration_between_samples / 2.0; + } + + let mut factor = 1.0; + if sampled_interval != 0.0 { + factor = (sampled_interval + duration_to_start + duration_to_end) / sampled_interval; + } + if is_rate { + factor /= range_ms as f64 / 1000.0; + } + Some(result * factor) + } + + fn assert_matches_prometheus( + timestamps: &[i64], + values: &[f64], + ranges: &[(u32, u32)], + eval_timestamps: &[i64], + range_ms: i64, + ) { + let actual = nullable_rate_runner::( + timestamps.to_vec(), + Arc::new(Float64Array::from(values.to_vec())), + ranges.to_vec(), + eval_timestamps.to_vec(), + range_ms, + ); + + for (index, ((offset, length), eval_ts)) in ranges.iter().zip(eval_timestamps).enumerate() { + let window = *offset as usize..(*offset + *length) as usize; + let expected = prometheus_extrapolated_rate( + ×tamps[window.clone()], + &values[window], + *eval_ts, + range_ms, + IS_COUNTER, + IS_RATE, + ); + match (actual[index], expected) { + (None, None) => {} + (Some(actual), Some(expected)) => assert!( + (actual - expected).abs() <= expected.abs() * 1e-9, + "window {index} {:?}: got {actual}, Prometheus gives {expected}", + ranges[index] + ), + (actual, expected) => { + panic!("window {index}: got {actual:?}, Prometheus gives {expected:?}") + } + } + } + } + + #[test] + fn extrapolation_matches_prometheus_on_seeded_windows() { + let mut prng = TinyPrng(0x51ed_270b_8f26_1a37); + // Uneven spacing so the average interval, and with it the extrapolation threshold, + // differs from window to window. + let timestamps = (0..48) + .scan(0i64, |clock, _| { + *clock += 1_000 + prng.next_index(4) as i64 * 500; + Some(*clock) + }) + .collect::>(); + // A counter that resets a few times, so the zero-snap branch is reached with both + // small and large leading values. + let values = (0..48) + .scan(0.0f64, |counter, _| { + *counter = match prng.next_index(8) { + 0 => 0.0, + 1 => *counter / 2.0, + _ => *counter + prng.next_index(50) as f64, + }; + Some(*counter) + }) + .collect::>(); + + let mut ranges = Vec::new(); + let mut eval_timestamps = Vec::new(); + for _ in 0..32 { + let length = 2 + prng.next_index(10) as u32; + let offset = prng.next_index(48 - length as usize) as u32; + ranges.push((offset, length)); + // Land the range boundary at varying distances from the samples, so the clamp + // fires on neither, one, or both sides. + let last = timestamps[(offset + length - 1) as usize]; + eval_timestamps.push(last + prng.next_index(5) as i64 * 500); + } + + assert_matches_prometheus::( + ×tamps, + &values, + &ranges, + &eval_timestamps, + 20_000, + ); + assert_matches_prometheus::( + ×tamps, + &values, + &ranges, + &eval_timestamps, + 20_000, + ); + assert_matches_prometheus::( + ×tamps, + &values, + &ranges, + &eval_timestamps, + 20_000, + ); + } + #[test] fn rate_rejects_wrong_input_arity() { let err = ExtrapolatedRate::::new(5) @@ -486,7 +1015,7 @@ mod test { ts_range, value_range, timestamps, - vec![2.0, 5.0, 0.0, 2.5, 0.0, 0.0], + vec![1.5, 5.0, 0.0, 2.5, 0.0, 0.0], ); } @@ -517,8 +1046,7 @@ mod test { ts_range, value_range, timestamps, - // `2.0` is because that `duration_to_zero` less than `extrapolation_threshold` - vec![2.0, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5], + vec![1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5], ); } @@ -591,7 +1119,7 @@ mod test { // that two `2.0` is because `duration_to_start` are shrunk to // `duration_to_zero`, and causes `duration_to_zero` less than // `extrapolation_threshold`. - vec![2.0, 1.5, 1.5, 1.5, 2.0, 1.5, 1.5, 1.5], + vec![1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5], ); } @@ -611,7 +1139,7 @@ mod test { ts_range, value_range, timestamps, - vec![4.0, 3.5, 3.5, 4.0], + vec![3.5, 3.5, 3.5, 3.5], ); } @@ -786,7 +1314,7 @@ mod test { ts_range, value_range, timestamps, - vec![400.0, 300.0, 300.0, 300.0, 400.0, 300.0, 300.0, 300.0], + vec![300.0, 300.0, 300.0, 300.0, 300.0, 300.0, 300.0, 300.0], ); } @@ -817,7 +1345,7 @@ mod test { ts_range, value_range, timestamps, - vec![400.0, 300.0, 300.0, 300.0, 300.0, 300.0, 300.0, 300.0], + vec![300.0, 300.0, 300.0, 300.0, 300.0, 300.0, 300.0, 300.0], ); } diff --git a/src/promql/src/functions/idelta.rs b/src/promql/src/functions/idelta.rs index 0772b0bf1ee..35ce44ba270 100644 --- a/src/promql/src/functions/idelta.rs +++ b/src/promql/src/functions/idelta.rs @@ -101,12 +101,12 @@ impl IDelta { .unwrap() .values(); - let value_values = value_range.values(); - let value_values = value_values - .as_any() - .downcast_ref::() - .unwrap() - .values(); + let value_array = value_range.values(); + let value_array = value_array.as_any().downcast_ref::().unwrap(); + // A NULL field value means the series has no sample at that timestamp, so the last two + // samples are not necessarily the last two slots. + let has_nulls = value_array.null_count() > 0; + let value_values = value_array.values(); let mut result_builder = Float64Builder::with_capacity(ts_range.len()); @@ -122,20 +122,29 @@ impl IDelta { value_len )), )?; - if len < 2 { - result_builder.append_null(); - continue; - } + let (last_position, prev_position) = if has_nulls { + match last_two_samples(value_array, value_offset, len) { + Some(positions) => positions, + None => { + result_builder.append_null(); + continue; + } + } + } else { + if len < 2 { + result_builder.append_null(); + continue; + } + (len - 1, len - 2) + }; - let last_offset = ts_offset + len - 1; - let prev_offset = last_offset - 1; + let last_offset = ts_offset + last_position; + let prev_offset = ts_offset + prev_position; let sampled_interval = (ts_values[last_offset] - ts_values[prev_offset]) as f64 / 1000.0; - let last_value_offset = value_offset + len - 1; - let prev_value_offset = last_value_offset - 1; - let last_value = value_values[last_value_offset]; - let prev_value = value_values[prev_value_offset]; + let last_value = value_values[value_offset + last_position]; + let prev_value = value_values[value_offset + prev_position]; if !IS_RATE { result_builder.append_value(last_value - prev_value); @@ -157,6 +166,22 @@ impl IDelta { } } +/// Locates the last two samples inside `[offset, offset + len)`, returning their positions +/// relative to `offset`. Returns `None` when the window holds fewer than two samples. +fn last_two_samples(values: &Float64Array, offset: usize, len: usize) -> Option<(usize, usize)> { + let mut last = None; + for position in (0..len).rev() { + if values.is_null(offset + position) { + continue; + } + match last { + None => last = Some(position), + Some(last) => return Some((last, position)), + } + } + None +} + impl Display for IDelta { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "PromQL Idelta Function (is_rate: {IS_RATE})",) @@ -166,6 +191,8 @@ impl Display for IDelta { #[cfg(test)] mod test { + use datafusion::arrow::buffer::NullBuffer; + use super::*; use crate::functions::test_util::simple_range_udf_runner; @@ -207,4 +234,33 @@ mod test { vec![Some(0.5), Some(0.0), None, Some(3.0), None, None], ); } + + #[test] + fn idelta_uses_last_two_samples_not_last_two_slots() { + let ts_array = Arc::new(TimestampMillisecondArray::from_iter_values([ + 0i64, 1000, 2000, 3000, + ])); + // Samples are 1.0@0 and 4.0@1000; the trailing slots only carry padding. + let values_array = Arc::new(Float64Array::new( + vec![1.0, 4.0, 100.0, 200.0].into(), + Some(NullBuffer::from_iter([true, true, false, false])), + )); + let ranges = [(0, 4), (2, 2), (3, 1)]; + + simple_range_udf_runner( + IDelta::::scalar_udf(), + RangeArray::from_ranges(ts_array.clone(), ranges).unwrap(), + RangeArray::from_ranges(values_array.clone(), ranges).unwrap(), + vec![], + vec![Some(3.0), None, None], + ); + + simple_range_udf_runner( + IDelta::::scalar_udf(), + RangeArray::from_ranges(ts_array, ranges).unwrap(), + RangeArray::from_ranges(values_array, ranges).unwrap(), + vec![], + vec![Some(3.0), None, None], + ); + } } diff --git a/src/promql/src/functions/quantile.rs b/src/promql/src/functions/quantile.rs index 35f0c1aa81b..3794057da5e 100644 --- a/src/promql/src/functions/quantile.rs +++ b/src/promql/src/functions/quantile.rs @@ -93,14 +93,14 @@ impl QuantileOverTime { )), )?; - let all_values = value_range - .values() - .as_any() - .downcast_ref::() - .unwrap() - .values(); + let value_array = value_range.values(); + let value_array = value_array.as_any().downcast_ref::().unwrap(); + // A NULL field value means the series has no sample at that timestamp, so a window's + // samples are not simply its slots. + let has_nulls = value_array.null_count() > 0; let mut result_builder = Float64Builder::with_capacity(ts_range.len()); let mut scratch = Vec::new(); + let mut samples = Vec::new(); match quantile_col { ColumnarValue::Scalar(quantile_scalar) => { @@ -125,11 +125,14 @@ impl QuantileOverTime { )), )?; - match quantile_with_scratch( - &all_values[value_offset..value_offset + value_len], - quantile, - &mut scratch, - ) { + let window = window_samples( + value_array, + has_nulls, + value_offset, + value_len, + &mut samples, + ); + match window_quantile(window, quantile, &mut scratch) { Some(value) => result_builder.append_value(value), None => result_builder.append_null(), } @@ -173,11 +176,14 @@ impl QuantileOverTime { } else { quantile_array.value(index) }; - match quantile_with_scratch( - &all_values[value_offset..value_offset + value_len], - quantile, - &mut scratch, - ) { + let window = window_samples( + value_array, + has_nulls, + value_offset, + value_len, + &mut samples, + ); + match window_quantile(window, quantile, &mut scratch) { Some(value) => result_builder.append_value(value), None => result_builder.append_null(), } @@ -190,6 +196,40 @@ impl QuantileOverTime { } } +/// Returns the samples of the window `[offset, offset + len)`, collecting the non-null ones +/// into `samples` when the backing array has nulls and borrowing the slice otherwise. +fn window_samples<'a>( + values: &'a Float64Array, + has_nulls: bool, + offset: usize, + len: usize, + samples: &'a mut Vec, +) -> &'a [f64] { + let raw_values = values.values(); + if !has_nulls { + return &raw_values[offset..offset + len]; + } + samples.clear(); + samples.extend( + (offset..offset + len) + .filter(|index| values.is_valid(*index)) + .map(|index| raw_values[index]), + ); + samples +} + +/// Quantile of one range window, or `None` when the window holds no sample. +/// +/// Prometheus returns an empty vector for a range without float samples rather than the NaN +/// that [`quantile_impl`] yields for an empty slice, so the emptiness check belongs here and +/// not in the shared kernel. +fn window_quantile(values: &[f64], quantile: f64, scratch: &mut Vec) -> Option { + if values.is_empty() { + return None; + } + quantile_with_scratch(values, quantile, scratch) +} + /// Refer to pub(crate) fn quantile_impl(values: &[f64], quantile: f64) -> Option { let mut scratch = Vec::new(); @@ -226,6 +266,9 @@ fn quantile_with_scratch(values: &[f64], quantile: f64, scratch: &mut Vec) #[cfg(test)] mod tests { + use datafusion::arrow::array::TimestampMillisecondArray; + use datafusion::arrow::buffer::NullBuffer; + use super::*; #[test] @@ -276,4 +319,64 @@ mod tests { let q = 0.25; assert_eq!(quantile_impl(values, q).unwrap(), 2.0); } + + #[test] + fn quantile_over_time_ranks_samples_only() { + let ts_array = Arc::new(TimestampMillisecondArray::from_iter_values([ + 0i64, 1000, 2000, + ])); + // Samples are 1.0 and 4.0; ranking the padding too would pull the median down. + let values_array = Arc::new(Float64Array::new( + vec![1.0, -100.0, 4.0].into(), + Some(NullBuffer::from_iter([true, false, true])), + )); + // The second window holds no sample, the third holds no slot at all. + let ranges = [(0, 3), (1, 1), (3, 0)]; + + let input = vec![ + ColumnarValue::Array(Arc::new( + RangeArray::from_ranges(ts_array, ranges) + .unwrap() + .into_dict(), + )), + ColumnarValue::Array(Arc::new( + RangeArray::from_ranges(values_array, ranges) + .unwrap() + .into_dict(), + )), + ColumnarValue::Scalar(ScalarValue::Float64(Some(0.5))), + ]; + let output = extract_array(&QuantileOverTime::quantile_over_time(&input).unwrap()).unwrap(); + let output = output.as_any().downcast_ref::().unwrap(); + + assert_eq!( + output.iter().collect::>(), + vec![Some(2.5), None, None] + ); + } + + #[test] + fn quantile_over_time_keeps_nan_for_an_invalid_quantile() { + let ts_array = Arc::new(TimestampMillisecondArray::from_iter_values([0i64, 1000])); + let values_array = Arc::new(Float64Array::from_iter_values([1.0, 4.0])); + let ranges = [(0, 2)]; + + let input = vec![ + ColumnarValue::Array(Arc::new( + RangeArray::from_ranges(ts_array, ranges) + .unwrap() + .into_dict(), + )), + ColumnarValue::Array(Arc::new( + RangeArray::from_ranges(values_array, ranges) + .unwrap() + .into_dict(), + )), + ColumnarValue::Scalar(ScalarValue::Float64(None)), + ]; + let output = extract_array(&QuantileOverTime::quantile_over_time(&input).unwrap()).unwrap(); + let output = output.as_any().downcast_ref::().unwrap(); + + assert!(output.value(0).is_nan()); + } } diff --git a/src/promql/src/functions/resets.rs b/src/promql/src/functions/resets.rs index 66e83cde311..7a5c0db71c1 100644 --- a/src/promql/src/functions/resets.rs +++ b/src/promql/src/functions/resets.rs @@ -43,8 +43,8 @@ mod test { use super::*; use crate::functions::test_util::{ - self, STALE_NAN, TinyPrng, assert_execution_error, build_test_range_arrays, - invoke_range_udf, simple_range_udf_runner, + self, STALE_NAN, assert_execution_error, build_test_range_arrays, invoke_range_udf, + simple_range_udf_runner, }; use crate::range_array::RangeArray; @@ -151,43 +151,17 @@ mod test { resets_oracle, Resets::scalar_udf(), ); - assert_eq!(expected, vec![Some(1.0), Some(0.0)]); + assert_eq!(expected, vec![Some(0.0), None]); } #[test] fn resets_range_array_seeded_differential() { - let mut prng = TinyPrng(0x2f6e_2b1d_834a_90c5); - let raw_values = (0..48) - .map(|_| match prng.next_index(12) { - 0 => -0.0, - 1 => 0.0, - 2 => -2.0, - 3 => -1.0, - 4 => 1.0, - 5 => 2.0, - 6 => f64::INFINITY, - 7 => f64::NEG_INFINITY, - 8 | 9 => f64::NAN, - _ => STALE_NAN, - }) - .collect::>(); - let values = raw_values.iter().copied().map(Some).collect(); - let mut timestamp_ranges = Vec::new(); - let mut value_ranges = Vec::new(); - for _ in 0..32 { - let length = prng.next_index(13) as u32; - timestamp_ranges.push((prng.next_index(65 - length as usize) as u32, length)); - value_ranges.push((prng.next_index(49 - length as usize) as u32, length)); - } + test_util::run_seeded_differential(resets_oracle, Resets::scalar_udf(), false); + } - test_util::run_oracle_ranges( - values, - raw_values, - timestamp_ranges, - value_ranges, - resets_oracle, - Resets::scalar_udf(), - ); + #[test] + fn resets_range_array_seeded_differential_with_nulls() { + test_util::run_seeded_differential(resets_oracle, Resets::scalar_udf(), true); } #[test] @@ -263,6 +237,8 @@ mod test { vec![Some(0.0), Some(0.0), Some(0.0)], ); + // The raw payload under the null slot would look like a reset; the sample sequence is + // 10 -> 10, which is none. The last window holds no sample at all. let values = Arc::new(Float64Array::new( vec![10.0, 7.0, 10.0].into(), Some(NullBuffer::from(vec![true, false, true])), @@ -276,10 +252,10 @@ mod test { ])); simple_range_udf_runner( Resets::scalar_udf(), - RangeArray::from_ranges(timestamps, [(0, 3)]).unwrap(), - RangeArray::from_ranges(values, [(0, 3)]).unwrap(), + RangeArray::from_ranges(timestamps, [(0, 3), (1, 1)]).unwrap(), + RangeArray::from_ranges(values, [(0, 3), (1, 1)]).unwrap(), vec![], - vec![Some(1.0)], + vec![Some(0.0), None], ); } } diff --git a/src/promql/src/functions/test_util.rs b/src/promql/src/functions/test_util.rs index e52f52e3bbb..6148e06c391 100644 --- a/src/promql/src/functions/test_util.rs +++ b/src/promql/src/functions/test_util.rs @@ -15,6 +15,7 @@ use std::sync::Arc; use datafusion::arrow::array::{Float64Array, TimestampMillisecondArray}; +use datafusion::arrow::buffer::NullBuffer; use datafusion::common::DataFusionError; use datafusion::logical_expr::ScalarUDF; use datafusion::physical_plan::ColumnarValue; @@ -144,12 +145,54 @@ impl TinyPrng { } } -/// Run the oracle-based differential test: build range arrays, verify raw bits and -/// null validity, invoke the UDF via [`simple_range_udf_runner`], and return the -/// expected values for further assertions. +/// Run [`run_oracle_ranges`] over a seeded mix of signed zeros, infinities, NaNs and stale +/// markers spread across random windows. `nullable` marks roughly a quarter of the slots as +/// null while keeping their raw payload, which moves the UDF off its null-free fast path. +pub fn run_seeded_differential(oracle: fn(&[f64]) -> Option, udf: ScalarUDF, nullable: bool) { + let mut prng = TinyPrng(0x2f6e_2b1d_834a_90c5); + let raw_values = (0..48) + .map(|_| match prng.next_index(12) { + 0 => -0.0, + 1 => 0.0, + 2 => -2.0, + 3 => -1.0, + 4 => 1.0, + 5 => 2.0, + 6 => f64::INFINITY, + 7 => f64::NEG_INFINITY, + 8 | 9 => f64::NAN, + _ => STALE_NAN, + }) + .collect::>(); + let values = raw_values + .iter() + .map(|value| (!nullable || prng.next_index(4) != 0).then_some(*value)) + .collect::>(); + let mut timestamp_ranges = Vec::new(); + let mut value_ranges = Vec::new(); + for _ in 0..32 { + let length = prng.next_index(13) as u32; + timestamp_ranges.push((prng.next_index(65 - length as usize) as u32, length)); + value_ranges.push((prng.next_index(49 - length as usize) as u32, length)); + } + + run_oracle_ranges( + values, + raw_values, + timestamp_ranges, + value_ranges, + oracle, + udf, + ); +} + +/// Run the oracle-based differential test: build range arrays, invoke the UDF via +/// [`simple_range_udf_runner`], and return the expected values for further assertions. /// /// `oracle` is the behavior-specific function (e.g. `changes_oracle` or `resets_oracle`) -/// that computes the expected count for a slice of raw f64 values. +/// that computes the expected count for a window's samples. Null slots are dropped before +/// the oracle runs: a NULL field value means the series has no sample at that timestamp, +/// so the raw payload underneath it is not a sample. pub fn run_oracle_ranges( values: Vec>, raw_values: Vec, @@ -166,29 +209,37 @@ pub fn run_oracle_ranges( .all(|((_, timestamp_length), (_, value_length))| timestamp_length == value_length) ); assert_eq!(values.len(), raw_values.len()); + // `values` marks the null slots, `raw_values` carries the payload of every slot including + // the null ones. The two must agree wherever a sample exists. + for (index, (value, raw)) in values.iter().zip(&raw_values).enumerate() { + assert!( + value.is_none_or(|value| value.to_bits() == raw.to_bits()), + "values[{index}] and raw_values[{index}] disagree on the sample" + ); + } let nulls = values.iter().map(Option::is_none).collect::>(); let expected = value_ranges .iter() - .map(|(offset, length)| oracle(&raw_values[*offset as usize..(*offset + *length) as usize])) + .map(|(offset, length)| { + let samples = (*offset as usize..(*offset + *length) as usize) + .filter(|index| !nulls[*index]) + .map(|index| raw_values[index]) + .collect::>(); + oracle(&samples) + }) .collect::>(); let timestamp_values = (0..64) .map(|value| Some(i64::from(value) * 1_000)) .collect::>(); let timestamp_array = Arc::new(TimestampMillisecondArray::from_iter(timestamp_values)); - let value_array = Arc::new(Float64Array::from_iter(values)); - for (index, ((actual, expected), is_null)) in value_array - .values() - .iter() - .zip(&raw_values) - .zip(nulls) - .enumerate() - { - assert_eq!(actual.to_bits(), expected.to_bits()); - assert_eq!(value_array.is_null(index), is_null); - } - + // Build from the raw payload plus a null bitmap instead of `from_iter`, which would zero + // the null slots and hide a function that reads them. + let value_array = Arc::new(Float64Array::new( + raw_values.clone().into(), + Some(NullBuffer::from_iter(nulls.iter().map(|null| !null))), + )); let timestamp_ranges = RangeArray::from_ranges(timestamp_array, timestamp_ranges).unwrap(); let value_ranges = RangeArray::from_ranges(value_array, value_ranges).unwrap(); simple_range_udf_runner( diff --git a/src/query/src/datafusion.rs b/src/query/src/datafusion.rs index f9d87782a03..28644b8ccc9 100644 --- a/src/query/src/datafusion.rs +++ b/src/query/src/datafusion.rs @@ -392,6 +392,7 @@ impl DatafusionQueryEngine { schema_name, table_name, columns_values: column_vectors, + skip_wal: query_ctx.skip_wal(), }; self.state diff --git a/src/query/src/dist_plan/analyzer.rs b/src/query/src/dist_plan/analyzer.rs index 433b66519b9..2cd32801d6c 100644 --- a/src/query/src/dist_plan/analyzer.rs +++ b/src/query/src/dist_plan/analyzer.rs @@ -248,7 +248,7 @@ impl DistPlannerAnalyzer { /// Try push down as many nodes as possible fn try_push_down(&self, plan: LogicalPlan) -> DfResult { let plan = plan.transform(&Self::inspect_plan_with_subquery)?; - let mut rewriter = PlanRewriter::default(); + let mut rewriter = PlanRewriter::new(&plan.data); let result = plan.data.rewrite(&mut rewriter)?.data; Self::assign_merge_scan_remote_dyn_filter_producer_ids(result) } @@ -302,7 +302,7 @@ impl DistPlannerAnalyzer { } fn handle_subquery(subquery: Subquery) -> DfResult { - let mut rewriter = PlanRewriter::default(); + let mut rewriter = PlanRewriter::new(&subquery.subquery); let mut rewrote_subquery = subquery .subquery .as_ref() @@ -391,6 +391,9 @@ enum RewriterStatus { #[derive(Debug, Default)] struct PlanRewriter { + /// Whether the whole plan this rewriter walks can be encoded to Substrait. + /// Used by [`PlanRewriter::should_expand`] to skip the per-node encoding check. + whole_plan_encodable: bool, /// Current level in the tree level: usize, /// Simulated stack for the `rewrite` recursion @@ -447,6 +450,17 @@ struct PlanRewriter { } impl PlanRewriter { + /// `plan` must be the root of the tree that is about to be rewritten, otherwise + /// [`PlanRewriter::should_expand`] may skip a check it has to perform. + fn new(plan: &LogicalPlan) -> Self { + Self { + whole_plan_encodable: DFLogicalSubstraitConvertor + .encode(plan, DefaultSerializer) + .is_ok(), + ..Default::default() + } + } + fn get_parent(&self) -> Option<&LogicalPlan> { // level starts from 1, it's safe to minus by 1 self.stack @@ -467,7 +481,16 @@ impl PlanRewriter { .collect::>() .join("\n"), ); - if let Err(e) = DFLogicalSubstraitConvertor.encode(plan, DefaultSerializer) { + // Substrait encoding recurses from the root to the leaves, so a root that encodes + // proves that every node below it encodes as well. `plan` here always comes from + // `self.stack`, which holds untouched sub-trees of that root, so one check on the + // root covers the whole descent. Each check builds a fresh `SessionState`, which + // is the dominant cost on deep PromQL plans. + // When the root does not encode, the plan holds at least one node that has to stay + // on the frontend and only the per-node check locates it. + if !self.whole_plan_encodable + && let Err(e) = DFLogicalSubstraitConvertor.encode(plan, DefaultSerializer) + { debug!( "PlanRewriter: plan cannot be converted to substrait with error={e:?}, expanding now: {plan}" ); diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index c42eb0e285a..5309d93a579 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/src/query/src/dist_plan/analyzer/test.rs @@ -28,6 +28,7 @@ use datafusion::datasource::DefaultTableSource; use datafusion::execution::SessionState; use datafusion::functions_aggregate::expr_fn::avg; use datafusion::functions_aggregate::min_max::{max, min}; +use datafusion::functions_nested::expr_fn::make_array; use datafusion::prelude::SessionContext; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ExprSchema, JoinType, ScalarValue}; @@ -367,6 +368,47 @@ fn frontend_only_histogram_folds_stay_above_merge_scan() { } } +/// `Categorizer` calls `Unnest` commutative, so the Substrait check in `should_expand` is +/// the only thing keeping it on the frontend. It also sits above an encodable projection, +/// which pins the rewriter to the whole-plan check rather than to the root node alone. +#[test] +fn unencodable_unnest_stays_above_merge_scan() { + let table = TestTable::table_with_name(0, "t".to_string()); + let table_source = Arc::new(DefaultTableSource::new(Arc::new( + DfTableProviderAdapter::new(table), + ))); + let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![]) + .unwrap() + .project(vec![ + col("pk1"), + make_array(vec![col("number")]).alias("numbers"), + ]) + .unwrap() + .unnest_column("numbers") + .unwrap() + .build() + .unwrap(); + assert!( + DFLogicalSubstraitConvertor + .encode(&plan, DefaultSerializer) + .is_err() + ); + + let result = DistPlannerAnalyzer {} + .analyze(plan, &ConfigOptions::default()) + .unwrap(); + let result_text = result.to_string(); + assert!(result_text.contains("Unnest:"), "{result_text}"); + + let remote_input = find_merge_scan(&result).unwrap().input(); + let remote_text = remote_input.to_string(); + assert!(!remote_text.contains("Unnest:"), "{remote_text}"); + assert!(remote_text.contains("make_array"), "{remote_text}"); + DFLogicalSubstraitConvertor + .encode(remote_input, DefaultSerializer) + .unwrap(); +} + #[test] fn dictionary_literals_are_unwrapped_before_pushdown() { let table = Arc::new(Table::new( diff --git a/src/query/src/optimizer/scan_hint.rs b/src/query/src/optimizer/scan_hint.rs index ce1009f8811..9cb50696ee8 100644 --- a/src/query/src/optimizer/scan_hint.rs +++ b/src/query/src/optimizer/scan_hint.rs @@ -19,7 +19,6 @@ use arrow_schema::SortOptions; use common_function::aggrs::aggr_wrapper::aggr_state_func_name; use common_recordbatch::OrderOption; use common_recordbatch::filter::SimpleFilterEvaluator; -use common_time::timestamp::TimeUnit; use datafusion::datasource::DefaultTableSource; use datafusion_common::tree_node::{Transformed, TreeNodeRewriter}; use datafusion_common::{Column, Result}; @@ -139,8 +138,7 @@ impl ScanHintRule { /// predicate later rejects that row. Only recognized tag/time predicates are /// allowed: tags select whole series, and supported time predicates constrain /// the scan window before row selection. Field or unrecognized predicates are - /// conservatively rejected. Finer-than-millisecond timestamps are also excluded - /// because instant evaluation can conflate distinct samples at that precision. + /// conservatively rejected. /// /// This checks only attached predicates; the path allowlist separately rejects /// residual Filter nodes between InstantManipulate and the scan. @@ -149,14 +147,6 @@ impl ScanHintRule { provider: &DummyTableProvider, ) -> bool { let metadata = provider.region_metadata(); - // Instant evaluation is millisecond-based, so finer time units can - // conflate timestamps and must not use the LastRow hint. - if !matches!( - metadata.time_index_type().unit(), - TimeUnit::Second | TimeUnit::Millisecond - ) { - return false; - } for filter in &table_scan.filters { let Some(filter) = SimpleFilterEvaluator::try_new(filter) else { return false; @@ -447,6 +437,10 @@ fn single_evaluation_node_allowed(node: &LogicalPlan) -> bool { /// This whitelist assumes the planner preserves time-index and series identity; /// it is not a proof that an arbitrary plan does so. +/// +/// Identity projections preserve the selected samples. Only the planner's named +/// seconds/milliseconds-to-milliseconds casts are accepted; a microsecond or +/// nanosecond cast could collapse a future sample onto the evaluation boundary. fn single_evaluation_projection_expr_allowed( expr: &Expr, projection: &datafusion_expr::logical_plan::Projection, @@ -635,6 +629,7 @@ mod test { 1000, 1000, 1000, + 0, "ts".to_string(), vec![], Some("v0".to_string()), @@ -682,26 +677,49 @@ mod test { } fn last_value_aggregate(input: LogicalPlan) -> LogicalPlan { - LogicalPlanBuilder::from(input) + let aggregate = LogicalPlanBuilder::from(input) .aggregate( vec![col("k0")], - vec![Expr::AggregateFunction(AggregateFunction { - func: last_value_udaf(), - params: AggregateFunctionParams { - args: vec![col("v0")], - distinct: false, - filter: None, - order_by: vec![Sort { - expr: col("ts"), - asc: true, - nulls_first: true, - }], - null_treatment: None, - }, - })], + vec![ + Expr::AggregateFunction(AggregateFunction { + func: last_value_udaf(), + params: AggregateFunctionParams { + args: vec![col("v0")], + distinct: false, + filter: None, + order_by: vec![Sort { + expr: col("ts"), + asc: true, + nulls_first: true, + }], + null_treatment: None, + }, + }), + Expr::AggregateFunction(AggregateFunction { + func: last_value_udaf(), + params: AggregateFunctionParams { + args: vec![col("ts")], + distinct: false, + filter: None, + order_by: vec![Sort { + expr: col("ts"), + asc: true, + nulls_first: true, + }], + null_treatment: None, + }, + }), + ], ) .unwrap() .build() + .unwrap(); + let timestamp = aggregate.schema().field(2).name().clone(); + + LogicalPlanBuilder::from(aggregate) + .project(vec![col("k0"), col(timestamp).alias("ts")]) + .unwrap() + .build() .unwrap() } @@ -712,6 +730,7 @@ mod test { 1000, 1000, 1000, + 0, "ts".to_string(), vec![], Some("v0".to_string()), @@ -741,6 +760,7 @@ mod test { outer_end, 1000, 1000, + 0, "ts".to_string(), vec![], Some("v0".to_string()), @@ -768,6 +788,7 @@ mod test { end, 1000, 1000, + 0, "ts".to_string(), vec![], Some("v0".to_string()), @@ -928,21 +949,39 @@ mod test { None, ) .unwrap() + .project(vec![ + Expr::Column(Column::new(Some("left"), "ts")), + Expr::Column(Column::new(Some("left"), "v0")), + ]) + .unwrap() .build() .unwrap(); let nonlast_aggregate = LogicalPlanBuilder::from(scan_plan(provider(), "aggregate")) .aggregate( vec![col("k0")], - vec![Expr::AggregateFunction(AggregateFunction { - func: max_udaf(), - params: AggregateFunctionParams { - args: vec![col("v0")], - distinct: false, - filter: None, - order_by: vec![], - null_treatment: None, - }, - })], + vec![ + Expr::AggregateFunction(AggregateFunction { + func: max_udaf(), + params: AggregateFunctionParams { + args: vec![col("v0")], + distinct: false, + filter: None, + order_by: vec![], + null_treatment: None, + }, + }), + Expr::AggregateFunction(AggregateFunction { + func: max_udaf(), + params: AggregateFunctionParams { + args: vec![col("ts")], + distinct: false, + filter: None, + order_by: vec![], + null_treatment: None, + }, + }) + .alias("ts"), + ], ) .unwrap() .build() @@ -953,6 +992,7 @@ mod test { 1000, 1000, 1000, + 0, 1000, "ts".to_string(), vec!["v0".to_string()], @@ -1023,7 +1063,49 @@ mod test { } #[test] - fn single_evaluation_rejects_microsecond_and_nanosecond_time_index_casts() { + fn single_evaluation_uses_last_row_for_microsecond_and_nanosecond_time_indexes() { + for timestamp_type in [ + ConcreteDataType::timestamp_microsecond_datatype(), + ConcreteDataType::timestamp_nanosecond_datatype(), + ] { + let direct_provider = Arc::new(mock_table_provider_with_timestamp( + RegionId::new(1, 1), + timestamp_type.clone(), + )); + let direct = ScanHintRule + .rewrite( + single_evaluation(scan_plan(direct_provider, "direct")), + &OptimizerContext::default(), + ) + .unwrap() + .data; + assert_eq!( + scan_requests(&direct)[0].series_row_selector, + Some(TimeSeriesRowSelector::LastRow { after_merge: true }) + ); + + let projection_provider = Arc::new(mock_table_provider_with_timestamp( + RegionId::new(1, 1), + timestamp_type, + )); + let projection = LogicalPlanBuilder::from(scan_plan(projection_provider, "projection")) + .project(vec![col("ts")]) + .unwrap() + .build() + .unwrap(); + let projected = ScanHintRule + .rewrite(single_evaluation(projection), &OptimizerContext::default()) + .unwrap() + .data; + assert_eq!( + scan_requests(&projected)[0].series_row_selector, + Some(TimeSeriesRowSelector::LastRow { after_merge: true }) + ); + } + } + + #[test] + fn single_evaluation_rejects_lossy_microsecond_and_nanosecond_time_index_casts() { for timestamp_type in [ ConcreteDataType::timestamp_microsecond_datatype(), ConcreteDataType::timestamp_nanosecond_datatype(), @@ -1081,7 +1163,7 @@ mod test { #[test] fn single_evaluation_rejects_projection_expressions_that_change_rows() { let invalid_projections = [ - vec![col("ts").alias("renamed")], + vec![col("ts").alias("renamed"), col("ts")], vec![ Expr::BinaryExpr(datafusion_expr::expr::BinaryExpr::new( Box::new(col("v0")), @@ -1089,6 +1171,7 @@ mod test { Box::new(lit(1.0_f64)), )) .alias("v0"), + col("ts"), ], vec![ Expr::Cast(Cast::new( @@ -1097,7 +1180,10 @@ mod test { )) .alias("ts"), ], - vec![Expr::Cast(Cast::new(Box::new(col("v0")), DataType::Int64)).alias("v0")], + vec![ + Expr::Cast(Cast::new(Box::new(col("v0")), DataType::Int64)).alias("v0"), + col("ts"), + ], vec![ Expr::Cast(Cast::new( Box::new(col("ts")), @@ -1417,6 +1503,7 @@ mod test { 1000, 1000, 1000, + 0, "ts".to_string(), vec![], Some("v0".to_string()), @@ -1441,6 +1528,7 @@ mod test { 2000, 1000, 1000, + 0, "ts".to_string(), vec![], Some("v0".to_string()), diff --git a/src/query/src/optimizer/scan_hint/vector_search.rs b/src/query/src/optimizer/scan_hint/vector_search.rs index 1a4d274ccaa..1f7f7b5a412 100644 --- a/src/query/src/optimizer/scan_hint/vector_search.rs +++ b/src/query/src/optimizer/scan_hint/vector_search.rs @@ -620,6 +620,7 @@ mod tests { 1000, 1000, 1000, + 0, "ts".to_string(), vec![], Some("v".to_string()), @@ -670,6 +671,7 @@ mod tests { 1000, 1000, 1000, + 0, "ts".to_string(), vec![], Some("v".to_string()), diff --git a/src/query/src/promql/planner.rs b/src/query/src/promql/planner.rs index 7f1d5d35fbe..6ca04ffca75 100644 --- a/src/query/src/promql/planner.rs +++ b/src/query/src/promql/planner.rs @@ -587,6 +587,7 @@ impl PromPlanner { self.ctx.start, self.ctx.end, self.ctx.interval, + 0, range_ms, time_index_column, self.ctx.field_columns.clone(), @@ -1943,6 +1944,11 @@ impl PromPlanner { if let Some(empty_plan) = self.setup_context().await? { return Ok(empty_plan); } + let offset_ms = match offset { + Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond, + Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond), + None => 0, + }; let normalize = self .selector_to_series_normalize_plan(offset, matchers, false) .await?; @@ -1974,8 +1980,48 @@ impl PromPlanner { DfExpr::Column(Column::new(qualifier.cloned(), field.name().clone())) }) .collect::>(); - project_exprs - .push(build_special_time_expr(&time_index_column).alias(×tamp_value_column)); + // `timestamp()` preserves the shifted selector timeline even though + // SeriesNormalize now retains raw native timestamp storage. Decimal + // arithmetic shifts before truncating to milliseconds. + let unit_factor = match col(&time_index_column) + .get_type(normalize.schema()) + .context(DataFusionPlanningSnafu)? + { + ArrowDataType::Timestamp(ArrowTimeUnit::Second, _) => (1_000_i128, 4, 0), + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, _) => (1, 1, 0), + ArrowDataType::Timestamp(ArrowTimeUnit::Microsecond, _) => (1, 4, 3), + ArrowDataType::Timestamp(ArrowTimeUnit::Nanosecond, _) => (1, 7, 6), + _ => unreachable!("time index is a timestamp"), + }; + let sample_time = col(&time_index_column) + .cast_to(&ArrowDataType::Int64, normalize.schema()) + .context(DataFusionPlanningSnafu)? + .cast_to(&ArrowDataType::Decimal128(19, 0), normalize.schema()) + .context(DataFusionPlanningSnafu)?; + let sample_time = DfExpr::BinaryExpr(BinaryExpr { + left: Box::new(sample_time), + op: Operator::Multiply, + right: Box::new(lit(ScalarValue::Decimal128( + Some(unit_factor.0), + unit_factor.1, + unit_factor.2, + ))), + }); + let sample_time = DfExpr::BinaryExpr(BinaryExpr { + left: Box::new(sample_time), + op: Operator::Plus, + right: Box::new(lit(ScalarValue::Decimal128(Some(offset_ms as i128), 19, 0))), + }) + .cast_to(&ArrowDataType::Int64, normalize.schema()) + .context(DataFusionPlanningSnafu)? + .cast_to(&ArrowDataType::Float64, normalize.schema()) + .context(DataFusionPlanningSnafu)?; + let sample_time = DfExpr::BinaryExpr(BinaryExpr { + left: Box::new(sample_time), + op: Operator::Divide, + right: Box::new(lit(1000.0)), + }); + project_exprs.push(sample_time.alias(×tamp_value_column)); let normalize = LogicalPlanBuilder::from(normalize) .project(project_exprs) .context(DataFusionPlanningSnafu)? @@ -1992,6 +2038,7 @@ impl PromPlanner { self.ctx.end, self.ctx.lookback_delta, self.ctx.interval, + offset_ms, time_index_column, if self.ctx.use_tsid { vec![DATA_SCHEMA_TSID_COLUMN_NAME.to_string()] @@ -2067,6 +2114,11 @@ impl PromPlanner { ensure!(!range.is_zero(), ZeroRangeSelectorSnafu); let range_ms = range.as_millis() as _; self.ctx.range = Some(range_ms); + let offset_ms = match offset { + Some(Offset::Pos(duration)) => duration.as_millis() as Millisecond, + Some(Offset::Neg(duration)) => -(duration.as_millis() as Millisecond), + None => 0, + }; // Some functions like rate may require special fields in the RangeManipulate plan // so we can't skip RangeManipulate. @@ -2081,6 +2133,7 @@ impl PromPlanner { self.ctx.start, self.ctx.end, self.ctx.interval, + offset_ms, // TODO(ruihang): convert via Timestamp datatypes to support different time units range_ms, self.ctx @@ -2142,8 +2195,6 @@ impl PromPlanner { ), }) }; - let preserve_any_value = - Self::field_columns_are_alternative_samples(input.schema(), &self.ctx.field_columns); let (mut func_exprs, new_tags) = self.create_function_expr( func, args.literals.clone(), @@ -2158,10 +2209,15 @@ impl PromPlanner { func_exprs.push(tsid_col); } + // A row survives as long as one field column produced a sample, and the fields without + // one stay NULL, which is the shape a selector already emits. Requiring every field to + // be non-NULL would drop one field's samples because another field has none in the same + // window — the reason alternative float/histogram columns already needed this form. A + // single field column reduces to the same predicate either way. let builder = LogicalPlanBuilder::from(input) .project(func_exprs) .context(DataFusionPlanningSnafu)? - .filter(self.create_empty_values_filter_expr(preserve_any_value)?) + .filter(self.create_empty_values_filter_expr(true)?) .context(DataFusionPlanningSnafu)?; let builder = match func.name { @@ -2339,14 +2395,18 @@ impl PromPlanner { None => 0, }; let mut scan_filters = Self::matchers_to_expr(label_matchers.clone(), table_schema)?; - if let Some(time_index_filter) = self.build_time_index_filter(offset_duration)? { + if let Some(time_index_filter) = + self.build_time_index_filter(offset_duration, table_schema)? + { scan_filters.push(time_index_filter); } - table_scan = LogicalPlanBuilder::from(table_scan) - .filter(conjunction(scan_filters).unwrap()) // Safety: `scan_filters` is not empty. - .context(DataFusionPlanningSnafu)? - .build() - .context(DataFusionPlanningSnafu)?; + if let Some(filter) = conjunction(scan_filters) { + table_scan = LogicalPlanBuilder::from(table_scan) + .filter(filter) + .context(DataFusionPlanningSnafu)? + .build() + .context(DataFusionPlanningSnafu)?; + } // make a projection plan if there is any `__field__` matcher if let Some(field_matchers) = &self.ctx.field_column_matcher { @@ -2718,72 +2778,99 @@ impl PromPlanner { Ok(table_ref) } - fn build_time_index_filter(&self, offset_duration: i64) -> Result> { + fn build_time_index_filter( + &self, + offset_duration: i64, + schema: &DFSchemaRef, + ) -> Result> { let start = self.ctx.start; let end = self.ctx.end; if end < start { return InvalidTimeRangeSnafu { start, end }.fail(); } - let lookback_delta = self.ctx.lookback_delta; - let range = self.ctx.range.unwrap_or_default(); - let interval = self.ctx.interval; let time_index_expr = self.create_time_index_column_expr()?; - let num_points = (end - start) / interval; - - // Prometheus semantics: - // - Instant selector lookback: (eval_ts - lookback_delta, eval_ts] - // - Range selector: (eval_ts - range, eval_ts] - // - // So samples positioned exactly at the lower boundary must be excluded. We align the scan - // lower bound with Prometheus by shifting it forward by 1ms (millisecond granularity), - // while still using a `>=` filter. - let selector_window = if range == 0 { lookback_delta } else { range }; - let lower_exclusive_adjustment = if selector_window > 0 { 1 } else { 0 }; - - // Scan a continuous time range - if (end - start) / interval > MAX_SCATTER_POINTS || interval <= INTERVAL_1H { - let single_time_range = time_index_expr - .clone() - .gt_eq(DfExpr::Literal( - ScalarValue::TimestampMillisecond( - Some( - self.ctx.start - offset_duration - selector_window - + lower_exclusive_adjustment, - ), - None, - ), - None, - )) - .and(time_index_expr.lt_eq(DfExpr::Literal( - ScalarValue::TimestampMillisecond(Some(self.ctx.end - offset_duration), None), - None, - ))); - return Ok(Some(single_time_range)); - } - - // Otherwise scan scatter ranges separately - let mut filters = Vec::with_capacity(num_points as usize + 1); - for timestamp in (start..=end).step_by(interval as usize) { - filters.push( + let time_index_name = self.ctx.time_index_column.as_ref().unwrap(); + let unit = schema + .index_of_column_by_name(None, time_index_name) + .and_then(|index| match schema.field(index).data_type() { + ArrowDataType::Timestamp(unit, _) => Some(*unit), + _ => None, + }) + .unwrap_or(ArrowTimeUnit::Millisecond); + let native_value = |milliseconds: i128| match unit { + ArrowTimeUnit::Second => milliseconds.div_euclid(1_000), + ArrowTimeUnit::Millisecond => milliseconds, + ArrowTimeUnit::Microsecond => milliseconds * 1_000, + ArrowTimeUnit::Nanosecond => milliseconds * 1_000_000, + }; + let scalar = |milliseconds: i128| -> Option { + let value = i64::try_from(native_value(milliseconds)).ok()?; + Some(match unit { + ArrowTimeUnit::Second => ScalarValue::TimestampSecond(Some(value), None), + ArrowTimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(value), None), + ArrowTimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(value), None), + ArrowTimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(value), None), + }) + }; + let window = self.ctx.range.unwrap_or(self.ctx.lookback_delta); + let filter = |lower_ms: i128, upper_ms: i128| { + let lower_value = native_value(lower_ms); + let upper_value = native_value(upper_ms); + if lower_value > i128::from(i64::MAX) || upper_value < i128::from(i64::MIN) { + return Some(lit(false)); + } + let lower_filter = (lower_value >= i128::from(i64::MIN)).then(|| { + let lower = DfExpr::Literal(scalar(lower_ms).unwrap(), None); + if window == 0 { + time_index_expr.clone().gt_eq(lower) + } else if unit == ArrowTimeUnit::Millisecond + && let Some(inclusive_lower) = lower_ms + .checked_add(1) + .and_then(|lower| i64::try_from(lower).ok()) + .and_then(|lower| scalar(i128::from(lower))) + { + time_index_expr + .clone() + .gt_eq(DfExpr::Literal(inclusive_lower, None)) + } else { + time_index_expr.clone().gt(lower) + } + }); + let upper_filter = (upper_value <= i128::from(i64::MAX)).then(|| { time_index_expr .clone() - .gt_eq(DfExpr::Literal( - ScalarValue::TimestampMillisecond( - Some( - timestamp - offset_duration - selector_window - + lower_exclusive_adjustment, - ), - None, - ), - None, - )) - .and(time_index_expr.clone().lt_eq(DfExpr::Literal( - ScalarValue::TimestampMillisecond(Some(timestamp - offset_duration), None), - None, - ))), - ) - } + .lt_eq(DfExpr::Literal(scalar(upper_ms).unwrap(), None)) + }); + // An underflowing lower bound must not discard a representable upper + // bound: without it, LastRow could retain a future row and discard the + // older eligible sample before the manipulator can check its time. + match (lower_filter, upper_filter) { + (Some(lower), Some(upper)) => Some(lower.and(upper)), + (Some(filter), None) | (None, Some(filter)) => Some(filter), + (None, None) => None, + } + }; + let bounds = |timestamp: i64| { + let upper = i128::from(timestamp) - i128::from(offset_duration); + (upper - i128::from(window), upper) + }; + let num_points = (end as i128 - start as i128) / self.ctx.interval as i128; + if num_points > MAX_SCATTER_POINTS as i128 || self.ctx.interval <= INTERVAL_1H { + let (lower, _) = bounds(start); + let (_, upper) = bounds(end); + return Ok(filter(lower, upper)); + } + let mut filters = Vec::new(); + for timestamp in (start..=end).step_by(self.ctx.interval as usize) { + let (lower, upper) = bounds(timestamp); + let Some(filter) = filter(lower, upper) else { + // A point whose native bounds cannot be represented may cover the whole native + // time domain, so its disjunct cannot be omitted. + return Ok(None); + }; + filters.push(filter); + } Ok(filters.into_iter().reduce(DfExpr::or)) } @@ -2884,14 +2971,16 @@ impl PromPlanner { self.ctx.tag_columns.clone() }; - let is_time_index_ms = scan_table + let time_index_data_type = scan_table .schema() .timestamp_column() .with_context(|| TimeIndexNotFoundSnafu { table: maybe_phy_table_ref.to_quoted_string(), })? .data_type - == ConcreteDataType::timestamp_millisecond_datatype(); + .clone(); + let is_time_index_second = + time_index_data_type == ConcreteDataType::timestamp_second_datatype(); let scan_projection = if table_id_filter.is_some() { let mut required_columns = HashSet::new(); @@ -2944,8 +3033,11 @@ impl PromPlanner { .context(DataFusionPlanningSnafu)?; } - if !is_time_index_ms { - // cast to ms if time_index not in Millisecond precision + if is_time_index_second { + // Promote seconds so millisecond offsets remain exact; retain finer precision. + // Later manipulators compare native sample ticks, while PromQL evaluation and + // emitted timestamps remain millisecond-based, so this projection must not + // silently truncate a finer-grained time index. let expr: Vec<_> = self .create_field_column_exprs()? .into_iter() @@ -2980,8 +3072,15 @@ impl PromPlanner { .context(DataFusionPlanningSnafu)? .build() .context(DataFusionPlanningSnafu)?; - } else if table_id_filter.is_some() { - // Drop the internal `__table_id` column after filtering. + } else if table_id_filter.is_some() + || time_index_data_type == ConcreteDataType::timestamp_microsecond_datatype() + || time_index_data_type == ConcreteDataType::timestamp_nanosecond_datatype() + { + // Drop the internal `__table_id` column after filtering and preserve PromQL's + // field/tag/timestamp column order for native microsecond/nanosecond timestamps. + // Keeping the original time column also lets the existing ordering hints + // use PerSeries scans without a cast, repartition, and sort. This benefits + // multi-evaluation selectors too; only a single evaluation can use LastRow. let project_exprs = self .create_field_column_exprs()? .into_iter() @@ -8832,7 +8931,7 @@ mod test { \n Projection: some_metric.timestamp, value AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ \n Projection: some_metric.timestamp, __promql_timestamp_value_ AS value, some_metric.tag_0 [timestamp:Timestamp(ms), value:Float64, tag_0:Utf8]\ \n PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ - \n Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(some_metric.timestamp AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ + \n Projection: some_metric.tag_0, some_metric.timestamp, some_metric.field_0, CAST(CAST(CAST(CAST(some_metric.timestamp AS Int64) AS Decimal128(19, 0)) * Decimal128(Some(1),1,0) + Decimal128(Some(0),19,0) AS Int64) AS Float64) / Float64(1000) AS __promql_timestamp_value_ [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N, __promql_timestamp_value_:Float64]\ \n PromSeriesDivide: tags=[\"tag_0\"] [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n Sort: some_metric.tag_0 ASC NULLS FIRST, some_metric.timestamp ASC NULLS FIRST [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ \n Filter: some_metric.tag_0 != Utf8(\"bar\") AND some_metric.timestamp >= TimestampMillisecond(-999, None) AND some_metric.timestamp <= TimestampMillisecond(100000000, None) [tag_0:Utf8, timestamp:Timestamp(ms), field_0:Float64;N]\ @@ -9051,7 +9150,17 @@ mod test { let manipulate = find_instant_manipulate(&plan).unwrap(); let exec = manipulate.to_execution_plan(Arc::new(DataSourceExec::new(Arc::new( - MemorySourceConfig::try_new(&[], Arc::new(ArrowSchema::empty()), None).unwrap(), + MemorySourceConfig::try_new( + &[], + Arc::new( + datafusion_expr::UserDefinedLogicalNodeCore::inputs(manipulate)[0] + .schema() + .as_arrow() + .clone(), + ), + None, + ) + .unwrap(), )))); assert!(format!("{exec:?}").contains("reuse_tsid_column: true")); } @@ -10814,6 +10923,7 @@ mod test { 1_000, 5_000, 1_000, + 0, "timestamp".to_string(), Vec::new(), Some(greptime_native_histogram().to_string()), @@ -11443,6 +11553,7 @@ mod test { 3000, 3000, 1000, + 0, 3000, "timestamp".to_string(), planner.ctx.field_columns.clone(), @@ -12135,6 +12246,101 @@ mod test { } } + #[tokio::test] + async fn native_scan_bounds_preserve_zero_lookback_and_overflow() { + let table_provider = build_test_table_provider( + &[(DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string())], + 1, + 1, + ) + .await; + let mut planner = PromPlanner { + table_provider, + ctx: PromPlannerContext::from_eval_stmt(&build_eval_stmt("some_metric")), + promql_annotations: None, + }; + planner.ctx.time_index_column = Some("timestamp".to_string()); + planner.ctx.start = 1_000; + planner.ctx.lookback_delta = 0; + let schema = Arc::new( + DFSchema::try_from(ArrowSchema::new(vec![Field::new( + "timestamp", + ArrowDataType::Timestamp(ArrowTimeUnit::Nanosecond, None), + false, + )])) + .unwrap(), + ); + for (end, interval, windows) in [ + (1_000, 1_000, 1), + (2_000, 1_000, 1), + (7_201_000, 7_200_000, 2), + ] { + planner.ctx.end = end; + planner.ctx.interval = interval; + let filter = planner + .build_time_index_filter(0, &schema) + .unwrap() + .unwrap() + .to_string(); + assert_eq!(filter.matches(">=").count(), windows, "{filter}"); + assert!( + filter.contains("TimestampNanosecond(1000000000, None)"), + "{filter}" + ); + } + planner.ctx.end = i64::MAX; + let filter = planner + .build_time_index_filter(0, &schema) + .unwrap() + .unwrap() + .to_string(); + assert!( + filter.contains("timestamp >= TimestampNanosecond(1000000000, None)"), + "{filter}" + ); + + // A lookback subtraction can underflow milliseconds while the upper bound remains + // representable. Keep that upper bound so LastRow cannot select a future sample. + let ms_schema = Arc::new( + DFSchema::try_from(ArrowSchema::new(vec![Field::new( + "timestamp", + ArrowDataType::Timestamp(ArrowTimeUnit::Millisecond, None), + false, + )])) + .unwrap(), + ); + planner.ctx.start = i64::MIN + 100; + planner.ctx.end = planner.ctx.start; + planner.ctx.lookback_delta = 200; + let filter = planner + .build_time_index_filter(0, &ms_schema) + .unwrap() + .unwrap() + .to_string(); + assert_eq!( + filter, + format!( + "timestamp <= TimestampMillisecond({}, None)", + i64::MIN + 100 + ) + ); + + // The lower bound can also overflow while converting milliseconds to native nanoseconds. + // Its representable upper bound still has to reach the scan. + planner.ctx.start = 0; + planner.ctx.end = 0; + planner.ctx.lookback_delta = 300_000; + let filter = planner + .build_time_index_filter(9_223_372_036_854, &schema) + .unwrap() + .unwrap() + .to_string(); + assert_eq!( + filter, + "timestamp <= TimestampNanosecond(-9223372036854000000, None)" + ); + } + #[tokio::test] async fn test_non_ms_precision() { let catalog_list = MemoryCatalogManager::with_default_setup(); @@ -12205,12 +12411,7 @@ mod test { .unwrap(); assert_eq!( plan.display_indent_schema().to_string(), - "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\ - \n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\ - \n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\ - \n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp >= TimestampMillisecond(-999, None) AND metrics.timestamp <= TimestampMillisecond(100000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\ - \n Projection: metrics.field, metrics.tag, CAST(metrics.timestamp AS Timestamp(ms)) AS timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\ - \n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]" + "PromInstantManipulate: range=[0..100000000], lookback=[1000], interval=[5000], time index=[timestamp] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp > TimestampNanosecond(-1000000000, None) AND metrics.timestamp <= TimestampNanosecond(100000000000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Projection: metrics.field, metrics.tag, metrics.timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]" ); let plan = PromPlanner::stmt_to_plan( DfTableSourceProvider::new( @@ -12235,15 +12436,7 @@ mod test { .unwrap(); assert_eq!( plan.display_indent_schema().to_string(), - "Filter: prom_avg_over_time(timestamp_range,field) IS NOT NULL [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\ - \n Projection: metrics.timestamp, prom_avg_over_time(timestamp_range, field) AS prom_avg_over_time(timestamp_range,field), metrics.tag [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\ - \n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[5000], time index=[timestamp], values=[\"field\"] [field:Dictionary(Int64, Float64);N, tag:Utf8, timestamp:Timestamp(ms), timestamp_range:Dictionary(Int64, Timestamp(ms))]\ - \n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\ - \n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\ - \n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\ - \n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp >= TimestampMillisecond(-4999, None) AND metrics.timestamp <= TimestampMillisecond(100000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\ - \n Projection: metrics.field, metrics.tag, CAST(metrics.timestamp AS Timestamp(ms)) AS timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ms)]\ - \n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]" + "Filter: prom_avg_over_time(timestamp_range,field) IS NOT NULL [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\n Projection: metrics.timestamp, prom_avg_over_time(timestamp_range, field) AS prom_avg_over_time(timestamp_range,field), metrics.tag [timestamp:Timestamp(ms), prom_avg_over_time(timestamp_range,field):Float64;N, tag:Utf8]\n PromRangeManipulate: req range=[0..100000000], interval=[5000], eval range=[5000], time index=[timestamp], values=[\"field\"] [field:Dictionary(Int64, Float64);N, tag:Utf8, timestamp:Timestamp(ms), timestamp_range:Dictionary(Int64, Timestamp(ms))]\n PromSeriesNormalize: offset=[0], time index=[timestamp], filter NaN: [true] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n PromSeriesDivide: tags=[\"tag\"] [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Sort: metrics.tag ASC NULLS FIRST, metrics.timestamp ASC NULLS FIRST [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Filter: metrics.tag = Utf8(\"1\") AND metrics.timestamp > TimestampNanosecond(-5000000000, None) AND metrics.timestamp <= TimestampNanosecond(100000000000000, None) [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n Projection: metrics.field, metrics.tag, metrics.timestamp [field:Float64;N, tag:Utf8, timestamp:Timestamp(ns)]\n TableScan: metrics [tag:Utf8, timestamp:Timestamp(ns), field:Float64;N]" ); } diff --git a/src/query/src/promql/planner/test/delta.rs b/src/query/src/promql/planner/test/delta.rs index 93e214d45b5..f46e211c758 100644 --- a/src/query/src/promql/planner/test/delta.rs +++ b/src/query/src/promql/planner/test/delta.rs @@ -453,6 +453,7 @@ async fn delta_mixed_ranges_drop_and_float_ranges_sum() { 2_000, 2_000, 1_000, + 0, 2_000, "ts".to_string(), planner.ctx.field_columns.clone(), @@ -757,3 +758,121 @@ async fn binary_joins_align_only_the_temporality_marker() { let (_, batches) = execute(set, &build_query_engine_state()).await; assert_eq!(1, batches.iter().map(RecordBatch::num_rows).sum::()); } + +#[tokio::test] +async fn delta_offsets_survive_optimized_plan_serialization() { + let eval_time = UNIX_EPOCH.checked_add(Duration::from_secs(120)).unwrap(); + for (name, query, expected) in [ + ( + "selector positive offset", + r#"delta_metric{series="cumulative"} offset 60s"#, + 10.0, + ), + ( + "selector negative offset", + r#"delta_metric{series="cumulative"} offset -60s"#, + 30.0, + ), + ( + "timestamp positive offset", + r#"timestamp(delta_metric{series="cumulative"} offset 60s)"#, + 120.0, + ), + ( + "timestamp negative offset", + r#"timestamp(delta_metric{series="cumulative"} offset -60s)"#, + 120.0, + ), + ( + "range positive offset", + r#"last_over_time(delta_metric{series="cumulative"}[60s] offset 60s)"#, + 10.0, + ), + ( + "range negative offset", + r#"last_over_time(delta_metric{series="cumulative"}[60s] offset -60s)"#, + 30.0, + ), + ( + "subquery positive offset", + r#"last_over_time((delta_metric{series="cumulative"} offset 60s)[60s:60s])"#, + 10.0, + ), + ( + "subquery negative offset", + r#"last_over_time((delta_metric{series="cumulative"} offset -60s)[60s:60s])"#, + 30.0, + ), + ] { + let eval_stmt = EvalStmt { + expr: parser::parse(query).unwrap(), + start: eval_time, + end: eval_time, + interval: Duration::from_secs(60), + lookback_delta: Duration::from_secs(300), + }; + let (provider, state, datafusion_table) = delta_temporality_table_provider(); + let raw = PromPlanner::stmt_to_plan(provider, &eval_stmt, &state) + .await + .unwrap(); + let context = QueryEngineContext::new(state.session_state(), QueryContext::arc()); + let optimized = state.optimize_by_extension_rules(raw, &context).unwrap(); + let optimized = state.optimize_logical_plan(optimized).unwrap(); + + let context = SessionContext::new_with_state(state.session_state()); + let catalog = Arc::new(MemoryCatalogProvider::new()); + let schema = Arc::new(MemorySchemaProvider::new()); + schema + .register_table("delta_metric".to_string(), datafusion_table) + .unwrap(); + catalog + .register_schema(DEFAULT_SCHEMA_NAME, schema) + .unwrap(); + context.register_catalog("datafusion", catalog); + let decoder = DefaultPlanDecoder::new(context.state(), &QueryContext::arc()).unwrap(); + let decoded = decoder + .decode( + DFLogicalSubstraitConvertor + .encode(&optimized, DefaultSerializer) + .unwrap(), + context.state().catalog_list().clone(), + false, + ) + .await + .unwrap(); + + let mut outputs = Vec::new(); + for plan in [optimized, decoded] { + let (_, batches) = execute(plan, &state).await; + let mut output = Vec::new(); + for batch in batches { + let value_field = batch + .schema() + .fields() + .iter() + .find(|field| field.data_type() == &ArrowDataType::Float64) + .unwrap() + .name() + .clone(); + let values = batch + .column_by_name(&value_field) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let timestamps = batch + .column_by_name(greptime_timestamp()) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + output.push((timestamps.value(row), values.value(row))); + } + } + assert_eq!(vec![(120_000, expected)], output, "{name}"); + outputs.push(output); + } + assert_eq!(outputs[0], outputs[1], "{name}"); + } +} diff --git a/src/query/src/query_engine/default_serializer.rs b/src/query/src/query_engine/default_serializer.rs index e009d80954f..f4d6b29fa2c 100644 --- a/src/query/src/query_engine/default_serializer.rs +++ b/src/query/src/query_engine/default_serializer.rs @@ -437,6 +437,7 @@ mod tests { 0, 1000, 1000, + 0, 1000, "timestamp".to_string(), vec!["float".to_string(), "histogram".to_string()], diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index f34623bc91c..a3d8432421b 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -284,6 +284,7 @@ impl QueryEngineState { catalog_manager: catalog_list, dyn_filter_registry_manager: Arc::new(DynFilterRegistryManager::default()), function_state: Arc::new(FunctionState { + plugins: plugins.clone(), table_mutation_handler, procedure_service_handler, flow_service_handler, diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml index 509699d8ba2..dbb956c0d5e 100644 --- a/src/servers/Cargo.toml +++ b/src/servers/Cargo.toml @@ -146,7 +146,7 @@ vrl.workspace = true zstd.workspace = true [target.'cfg(not(windows))'.dependencies] -tikv-jemalloc-ctl = { version = "0.6", features = ["use_std", "stats"] } +tikv-jemalloc-ctl = { version = "0.7", features = ["use_std", "stats"] } [dev-dependencies] auth = { workspace = true, features = ["testing"] } @@ -172,7 +172,10 @@ tokio-postgres.workspace = true tokio-postgres-rustls = "0.14" [target.'cfg(unix)'.dev-dependencies] -pprof = { version = "0.14", features = ["criterion", "flamegraph"] } +pprof = { version = "0.15", features = ["criterion", "flamegraph"] } + +[target.'cfg(not(windows))'.dev-dependencies] +tikv-jemallocator = "0.7" [build-dependencies] common-version.workspace = true @@ -190,6 +193,10 @@ required-features = ["testing"] name = "to_http_output" harness = false +[[bench]] +name = "prometheus_response" +harness = false + [[bench]] name = "loki_labels" harness = false diff --git a/src/servers/benches/prometheus_response.rs b/src/servers/benches/prometheus_response.rs new file mode 100644 index 00000000000..5d9325253a2 --- /dev/null +++ b/src/servers/benches/prometheus_response.rs @@ -0,0 +1,231 @@ +// 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 std::time::Duration; + +use axum::response::IntoResponse; +use common_query::Output; +use common_recordbatch::{RecordBatch, RecordBatches}; +use criterion::{Criterion, Throughput, criterion_group, criterion_main}; +use datatypes::data_type::ConcreteDataType; +use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; +use datatypes::vectors::{Float64Vector, StringVector, TimestampMillisecondVector, VectorRef}; +use promql_parser::parser::value::ValueType; +use servers::http::prometheus::{ + PromQueryResult, PromSampleValue, PrometheusJsonResponse, PrometheusResponse, +}; + +const SERIES: usize = 64; +const POINTS: usize = 2048; +/// Label count of the one-sample-per-series shape an instant query returns. +const LABELS: usize = 4; + +// Response building is dominated by allocation, so match the allocator the +// server binary uses instead of the platform default. +#[cfg(not(windows))] +#[global_allocator] +static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + +/// Builds a query result of `series_count` series with `points` samples each. +/// +/// `run_length` is the number of consecutive rows that belong to the same +/// series, and `permuted` shuffles timestamps so the response has to sort them. +fn input( + label_count: usize, + run_length: usize, + series_count: usize, + points: usize, + permuted: bool, +) -> (SchemaRef, Vec) { + let mut columns = vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), false), + ]; + columns.extend((0..label_count).map(|index| { + ColumnSchema::new( + format!("label_{index}"), + ConcreteDataType::string_datatype(), + false, + ) + })); + let schema = Arc::new(Schema::new(columns)); + let point = |row: usize| { + let point = row / (series_count * run_length) * run_length + row % run_length; + if permuted { + (point * 109 + 17) % points + } else { + point + } + }; + let mut vectors: Vec = vec![ + Arc::new(TimestampMillisecondVector::from_vec( + (0..series_count * points) + .map(|row| point(row) as i64 * 300_000) + .collect(), + )), + Arc::new(Float64Vector::from( + (0..series_count * points) + .map(|row| Some((point(row) % 1000) as f64 * 0.25)) + .collect::>(), + )), + ]; + for label in 0..label_count { + let values: Vec<_> = (0..series_count) + .map(|series| format!("label-{label}-series-{series:04}")) + .collect(); + vectors.push(Arc::new(StringVector::from( + (0..series_count * points) + .map(|row| Some(values[(row / run_length) % series_count].as_str())) + .collect::>(), + ))); + } + let batch = RecordBatch::new(schema.clone(), vectors).unwrap(); + let batches = (0..batch.num_rows()) + .step_by(1024) + .map(|offset| batch.slice(offset, 1024).unwrap()) + .collect(); + (schema, batches) +} + +fn bench_prometheus_response(c: &mut Criterion) { + let runtime = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + let mut group = c.benchmark_group("prometheus_response_complete"); + group.sample_size(20); + group.warm_up_time(Duration::from_millis(500)); + group.measurement_time(Duration::from_secs(2)); + group.throughput(Throughput::Elements((SERIES * POINTS) as u64)); + for labels in [1, 4] { + for run in [1, 32, POINTS] { + for permuted in [false, true] { + let (schema, batches) = input(labels, run, SERIES, POINTS, permuted); + let convert = || { + runtime.block_on(PrometheusJsonResponse::from_query_result( + Ok(Output::new_with_record_batches( + RecordBatches::try_new(schema.clone(), batches.clone()).unwrap(), + )), + Some("metric".to_string()), + ValueType::Matrix, + None, + )) + }; + let response = convert(); + assert_eq!(response.status, "success"); + let PrometheusResponse::PromData(data) = response.data else { + panic!("expected Prometheus data"); + }; + let PromQueryResult::Matrix(series) = data.result else { + panic!("expected matrix"); + }; + assert_eq!(series.len(), SERIES); + for (index, series) in series.iter().enumerate() { + assert_eq!(series.metric.len(), labels + 1); + assert_eq!(series.metric["__name__"], "metric"); + for label in 0..labels { + assert_eq!( + series.metric[&format!("label_{label}")], + format!("label-{label}-series-{index:04}") + ); + } + assert!(series.histograms.is_empty()); + assert_eq!(series.values.len(), POINTS); + for (point, (timestamp, value)) in series.values.iter().enumerate() { + assert_eq!(*timestamp, point as f64 * 300.0); + assert!( + matches!(value, PromSampleValue::Number(value) if *value == (point % 1000) as f64 * 0.25) + ); + } + } + let order = if permuted { "permuted" } else { "ordered" }; + group.bench_function(format!("labels{labels}_run{run}_{order}"), |b| { + b.iter(|| black_box(convert().into_response())); + }); + } + } + } + group.finish(); +} + +fn bench_prometheus_single_point(c: &mut Criterion) { + let runtime = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + let mut group = c.benchmark_group("prometheus_single_point_complete"); + group.sample_size(20); + group.warm_up_time(Duration::from_millis(500)); + group.measurement_time(Duration::from_secs(2)); + for series_count in [4096, 65536] { + let (schema, batches) = input(LABELS, 1, series_count, 1, false); + group.throughput(Throughput::Elements(series_count as u64)); + for result_type in [ValueType::Matrix, ValueType::Vector] { + let convert = || { + runtime.block_on(PrometheusJsonResponse::from_query_result( + Ok(Output::new_with_record_batches( + RecordBatches::try_new(schema.clone(), batches.clone()).unwrap(), + )), + Some("metric".to_string()), + result_type, + None, + )) + }; + let response = convert(); + assert_eq!(response.status, "success"); + let PrometheusResponse::PromData(data) = response.data else { + panic!("expected Prometheus data"); + }; + match data.result { + PromQueryResult::Matrix(series) => { + assert_eq!(series.len(), series_count); + assert!(series.iter().all(|series| { + series.metric.len() == LABELS + 1 + && series.values.len() == 1 + && series.values[0].0 == 0.0 + && matches!(series.values[0].1, PromSampleValue::Number(0.0)) + && series.histograms.is_empty() + })); + } + PromQueryResult::Vector(series) => { + assert_eq!(series.len(), series_count); + assert!(series.iter().all(|series| { + series.metric.len() == LABELS + 1 + && series.value.as_ref().is_some_and(|(timestamp, value)| { + *timestamp == 0.0 && value == "0.0" + }) + && series.histogram.is_none() + })); + } + _ => panic!("expected matrix or vector"), + } + group.bench_function( + format!("{result_type}_labels{LABELS}_series{series_count}"), + |b| b.iter(|| black_box(convert().into_response())), + ); + } + } + group.finish(); +} + +criterion_group!( + benches, + bench_prometheus_response, + bench_prometheus_single_point +); +criterion_main!(benches); diff --git a/src/servers/src/grpc/context_auth.rs b/src/servers/src/grpc/context_auth.rs index 39c4fc5c88d..2e0bda0939c 100644 --- a/src/servers/src/grpc/context_auth.rs +++ b/src/servers/src/grpc/context_auth.rs @@ -21,6 +21,7 @@ use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; use common_catalog::parse_catalog_and_schema_from_db_string; use common_error::ext::ErrorExt; use session::context::{Channel, QueryContextBuilder, QueryContextRef}; +use session::hints::INSERT_SKIP_WAL_HINT; use snafu::{OptionExt, ResultExt}; use tonic::Status; use tonic::metadata::MetadataMap; @@ -28,6 +29,7 @@ use tonic::metadata::MetadataMap; use crate::error::Error::UnsupportedAuthScheme; use crate::error::{AuthSnafu, InvalidParameterSnafu, NotFoundAuthHeaderSnafu, Result}; use crate::grpc::TonicResult; +use crate::hint_headers; use crate::http::AUTHORIZATION_HEADER; use crate::http::header::constants::GREPTIME_DB_HEADER_NAME; use crate::metrics::METRIC_AUTH_FAILURE; @@ -45,13 +47,26 @@ pub fn create_query_context_from_grpc_metadata( ) }; - Ok(Arc::new( - QueryContextBuilder::default() - .current_catalog(catalog) - .current_schema(schema) - .channel(Channel::Grpc) - .build(), - )) + let ctx = QueryContextBuilder::default() + .current_catalog(catalog) + .current_schema(schema) + .channel(Channel::Grpc) + .build(); + // OTEL Arrow uses ordinary inserts. Accept only its request-level WAL hint, + // leaving unrelated hints and reserved internal extensions unchanged. + if let Some((key, value)) = hint_headers::extract_hints(headers) + .into_iter() + .find(|(key, _)| key == INSERT_SKIP_WAL_HINT) + { + let skip_wal = value.parse::().map_err(|_| { + InvalidParameterSnafu { + reason: format!("Invalid {key} hint: expected true or false, got {value:?}"), + } + .build() + })?; + ctx.set_skip_wal(skip_wal); + } + Ok(Arc::new(ctx)) } /// Helper function to extract a header from the metadata map. @@ -161,3 +176,58 @@ pub async fn auth( .inc(); }) } + +#[cfg(test)] +mod tests { + use session::hints::{HINTS_KEY, REMOTE_QUERY_ID_EXTENSION_KEY, RESERVED_EXTENSION_KEYS}; + + use super::*; + + #[test] + fn test_arrow_insert_hint_does_not_accept_reserved_extensions() { + let mut headers = MetadataMap::new(); + assert_eq!( + create_query_context_from_grpc_metadata(&headers) + .unwrap() + .extension(INSERT_SKIP_WAL_HINT), + None + ); + for (value, expected) in [("true", true), ("false", false)] { + let mut hints = format!("insert_skip_wal={value},ttl=7d"); + for key in RESERVED_EXTENSION_KEYS { + hints.push_str(&format!(",{key}=external")); + } + headers.insert(HINTS_KEY, hints.parse().unwrap()); + let ctx = create_query_context_from_grpc_metadata(&headers).unwrap(); + assert_eq!(ctx.skip_wal(), expected); + assert_eq!(ctx.extension(INSERT_SKIP_WAL_HINT), None); + assert_eq!(ctx.extension("ttl"), None); + for key in RESERVED_EXTENSION_KEYS { + if key == REMOTE_QUERY_ID_EXTENSION_KEY { + // The builder generates this ID; external hints must not replace it. + assert!(ctx.remote_query_id().is_some()); + assert_ne!(ctx.extension(key), Some("external")); + } else { + assert_eq!(ctx.extension(key), None); + } + } + } + // Only the first matching hint is parsed and applied. + for (hints, expected) in [ + ("insert_skip_wal=true,insert_skip_wal=false", true), + ("insert_skip_wal=false,insert_skip_wal=true", false), + ("insert_skip_wal=true,insert_skip_wal=invalid", true), + ] { + headers.insert(HINTS_KEY, hints.parse().unwrap()); + let ctx = create_query_context_from_grpc_metadata(&headers).unwrap(); + assert_eq!(ctx.skip_wal(), expected); + } + for value in ["", "TRUE", "1", "invalid"] { + headers.insert( + HINTS_KEY, + format!("insert_skip_wal={value}").parse().unwrap(), + ); + assert!(create_query_context_from_grpc_metadata(&headers).is_err()); + } + } +} diff --git a/src/servers/src/grpc/greptime_handler.rs b/src/servers/src/grpc/greptime_handler.rs index 98786d30adc..dcd8c641a64 100644 --- a/src/servers/src/grpc/greptime_handler.rs +++ b/src/servers/src/grpc/greptime_handler.rs @@ -38,7 +38,7 @@ use common_telemetry::{debug, error, tracing, warn}; use common_time::timezone::parse_timezone; use futures_util::StreamExt; use session::context::{Channel, QueryContextBuilder, QueryContextRef}; -use session::hints::{READ_PREFERENCE_HINT, is_reserved_extension_key}; +use session::hints::{INSERT_SKIP_WAL_HINT, READ_PREFERENCE_HINT, is_reserved_extension_key}; use snafu::{OptionExt, ResultExt}; use tokio::sync::mpsc; use tokio::sync::mpsc::error::TrySendError; @@ -202,7 +202,7 @@ pub fn get_request_type(request: &GreptimeRequest) -> &'static str { pub(crate) fn create_query_context( channel: Channel, header: Option<&RequestHeader>, - mut extensions: Vec<(String, String)>, + extensions: Vec<(String, String)>, snapshot_seqs: HashMap, ) -> Result { let (catalog, schema) = header @@ -240,29 +240,36 @@ pub(crate) fn create_query_context( .channel(channel) .snapshot_seqs(Arc::new(RwLock::new(snapshot_seqs))); - if let Some(x) = extensions - .iter() - .position(|(k, _)| k == READ_PREFERENCE_HINT) - { - let (k, v) = extensions.swap_remove(x); - let Ok(read_preference) = ReadPreference::from_str(&v) else { - return UnknownHintSnafu { - hint: format!("{k}={v}"), - } - .fail(); - }; - ctx_builder = ctx_builder.read_preference(read_preference); - } - for (key, value) in extensions { - if is_reserved_extension_key(&key) { - debug!( - key = key.as_str(), - "Ignoring reserved external query context extension key" - ); - continue; + match key.as_str() { + READ_PREFERENCE_HINT => { + let Ok(read_preference) = ReadPreference::from_str(&value) else { + return UnknownHintSnafu { + hint: format!("{key}={value}"), + } + .fail(); + }; + ctx_builder = ctx_builder.read_preference(read_preference); + } + INSERT_SKIP_WAL_HINT => { + let skip_wal = value.parse::().map_err(|_| { + UnknownHintSnafu { + hint: format!("{key}={value}"), + } + .build() + })?; + ctx_builder = ctx_builder.skip_wal(skip_wal); + } + _ if is_reserved_extension_key(&key) => { + debug!( + key = key.as_str(), + "Ignoring reserved external query context extension key" + ); + } + _ => { + ctx_builder = ctx_builder.set_extension(key, value); + } } - ctx_builder = ctx_builder.set_extension(key, value); } Ok(ctx_builder.build().into()) } @@ -322,6 +329,81 @@ mod tests { use super::*; use crate::error::{ExecuteGrpcRequestSnafu, InvalidParameterSnafu}; + #[test] + fn test_create_query_context_typed_skip_wal() { + let ctx = create_query_context(Channel::Grpc, None, vec![], HashMap::new()).unwrap(); + assert!(!ctx.skip_wal()); + let legacy = create_query_context( + Channel::Grpc, + None, + vec![("skip_wal".to_string(), "true".to_string())], + HashMap::new(), + ) + .unwrap(); + assert!(!legacy.skip_wal()); + assert_eq!(legacy.extension("skip_wal"), Some("true")); + for (value, expected) in [("true", true), ("false", false)] { + let ctx = create_query_context( + Channel::Grpc, + None, + vec![(INSERT_SKIP_WAL_HINT.to_string(), value.to_string())], + HashMap::new(), + ) + .unwrap(); + assert_eq!(ctx.skip_wal(), expected); + assert_eq!(ctx.extension(INSERT_SKIP_WAL_HINT), None); + } + for value in ["", "TRUE", "1", "invalid"] { + assert!( + create_query_context( + Channel::Grpc, + None, + vec![(INSERT_SKIP_WAL_HINT.to_string(), value.to_string())], + HashMap::new() + ) + .is_err() + ); + } + let ctx = create_query_context( + Channel::Grpc, + None, + vec![ + (INSERT_SKIP_WAL_HINT.to_string(), "true".to_string()), + (INSERT_SKIP_WAL_HINT.to_string(), "false".to_string()), + ], + HashMap::new(), + ) + .unwrap(); + assert!(!ctx.skip_wal()); + assert_eq!(ctx.extension(INSERT_SKIP_WAL_HINT), None); + } + + #[test] + fn test_create_query_context_read_preference_duplicates() { + for (values, valid) in [ + (["leader", "LEADER"], true), + (["invalid", "leader"], false), + (["leader", "invalid"], false), + ] { + let result = create_query_context( + Channel::Grpc, + None, + values + .into_iter() + .map(|value| (READ_PREFERENCE_HINT.to_string(), value.to_string())) + .collect(), + HashMap::new(), + ); + if valid { + let ctx = result.unwrap(); + assert!(matches!(ctx.read_preference(), ReadPreference::Leader)); + assert_eq!(ctx.extension(READ_PREFERENCE_HINT), None); + } else { + assert!(result.is_err()); + } + } + } + #[test] fn test_create_query_context() { let header = RequestHeader { diff --git a/src/servers/src/grpc/prom_query_gateway.rs b/src/servers/src/grpc/prom_query_gateway.rs index 99940ad9a24..92bd5d722e5 100644 --- a/src/servers/src/grpc/prom_query_gateway.rs +++ b/src/servers/src/grpc/prom_query_gateway.rs @@ -135,11 +135,15 @@ impl PrometheusGatewayService { }; let (metric_name, mut result_type) = retrieve_metric_name_and_result_type(query.expr()); let query_id = ctx.remote_query_id().map(str::to_string); - let result = self.handler.do_query_parsed(query, ctx).await; - // range query only returns matrix - if is_range_query { + // A range query only returns a matrix, and matrix serialization sorts + // samples and series, so execution order never reaches the response. + let query = if is_range_query { result_type = ValueType::Matrix; + query.with_unordered_output() + } else { + query }; + let result = self.handler.do_query_parsed(query, ctx).await; PrometheusJsonResponse::from_query_result( result, diff --git a/src/servers/src/hint_headers.rs b/src/servers/src/hint_headers.rs index 5f64be584fa..3d13dd9bf5e 100644 --- a/src/servers/src/hint_headers.rs +++ b/src/servers/src/hint_headers.rs @@ -59,6 +59,22 @@ mod tests { use super::*; + #[test] + fn test_extract_skip_wal_hint() { + use session::hints::INSERT_SKIP_WAL_HINT; + + let mut headers = HeaderMap::new(); + headers.insert(HINTS_KEY, HeaderValue::from_static("insert_skip_wal=true")); + let mut metadata = MetadataMap::new(); + metadata.insert( + HINTS_KEY, + MetadataValue::from_static("insert_skip_wal=true"), + ); + let expected = vec![(INSERT_SKIP_WAL_HINT.to_string(), "true".to_string())]; + assert_eq!(extract_hints(&headers), expected); + assert_eq!(extract_hints(&metadata), expected); + } + #[test] fn test_extract_hints_with_full_header_map() { let mut headers = HeaderMap::new(); diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index 311ed48d912..40b7bb515b8 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -117,6 +117,7 @@ mod client_ip; use crate::prom_remote_write::validation::PromValidationMode; mod hints; mod read_preference; +mod skip_wal; #[cfg(any(test, feature = "testing"))] pub mod test_helpers; @@ -1059,7 +1060,8 @@ impl HttpServer { .layer(middleware::from_fn(client_ip::log_error_with_client_ip)) .layer(middleware::from_fn( read_preference::extract_read_preference, - )), + )) + .layer(middleware::from_fn(skip_wal::extract_skip_wal)), ); // Debug handlers are part of the complete router; the API listener hides diff --git a/src/servers/src/http/header.rs b/src/servers/src/http/header.rs index 1dc94bb76eb..625bee41ee7 100644 --- a/src/servers/src/http/header.rs +++ b/src/servers/src/http/header.rs @@ -44,6 +44,7 @@ pub mod constants { pub const GREPTIME_DB_HEADER_METRICS: &str = "x-greptime-metrics"; pub const GREPTIME_DB_HEADER_NAME: &str = "x-greptime-db-name"; pub const GREPTIME_DB_HEADER_READ_PREFERENCE: &str = "x-greptime-read-preference"; + pub const GREPTIME_INSERT_SKIP_WAL_HEADER_NAME: &str = "x-greptime-insert-skip-wal"; pub const GREPTIME_TIMEZONE_HEADER_NAME: &str = "x-greptime-timezone"; pub const GREPTIME_DB_HEADER_ERROR_CODE: &str = common_error::GREPTIME_DB_HEADER_ERROR_CODE; @@ -94,6 +95,10 @@ pub static GREPTIME_TIMEZONE_HEADER_NAME: HeaderName = pub static GREPTIME_DB_HEADER_READ_PREFERENCE: HeaderName = HeaderName::from_static(constants::GREPTIME_DB_HEADER_READ_PREFERENCE); +/// Request-level WAL policy, independent of the table-level skip_wal option. +pub static GREPTIME_INSERT_SKIP_WAL_HEADER_NAME: HeaderName = + HeaderName::from_static(constants::GREPTIME_INSERT_SKIP_WAL_HEADER_NAME); + pub static CONTENT_TYPE_PROTOBUF_STR: &str = "application/x-protobuf"; pub static CONTENT_TYPE_PROTOBUF: HeaderValue = HeaderValue::from_static(CONTENT_TYPE_PROTOBUF_STR); pub static CONTENT_ENCODING_SNAPPY: HeaderValue = HeaderValue::from_static("snappy"); diff --git a/src/servers/src/http/prometheus.rs b/src/servers/src/http/prometheus.rs index e1147c92b65..1f9988ede4c 100644 --- a/src/servers/src/http/prometheus.rs +++ b/src/servers/src/http/prometheus.rs @@ -71,7 +71,7 @@ use table::requests::{ SEMANTIC_METRIC_UNIT, SEMANTIC_VALUE_MIXED, }; -pub use super::result::prometheus_resp::PrometheusJsonResponse; +pub use super::result::prometheus_resp::{PromSampleValue, PrometheusJsonResponse}; use crate::error::{ CollectRecordbatchSnafu, ConvertScalarValueSnafu, DataFusionSnafu, Error, InvalidQuerySnafu, NotSupportedSnafu, Result, TableNotFoundSnafu, UnexpectedResultSnafu, @@ -100,7 +100,7 @@ pub struct PromSeriesVector { pub struct PromSeriesMatrix { pub metric: BTreeMap, #[serde(skip_serializing_if = "Vec::is_empty", default)] - pub values: Vec<(f64, String)>, + pub values: Vec<(f64, PromSampleValue)>, #[serde(skip_serializing_if = "Vec::is_empty", default)] pub histograms: Vec<(f64, PromNativeHistogram)>, } @@ -624,7 +624,11 @@ async fn do_range_query( ) -> PrometheusJsonResponse { let (metric_name, _) = retrieve_metric_name_and_result_type(prom_query.expr()); let query_id = query_ctx.remote_query_id().map(str::to_string); - let result = handler.do_query_parsed(prom_query, query_ctx).await; + // Matrix serialization sorts samples and series, so execution order never + // reaches the response. + let result = handler + .do_query_parsed(prom_query.with_unordered_output(), query_ctx) + .await; PrometheusJsonResponse::from_query_result( result, metric_name, @@ -2404,6 +2408,7 @@ mod tests { denied_table: Option<&'static str>, metric_names: Vec, queries: Mutex>, + ordered_outputs: Mutex>, } #[async_trait::async_trait] @@ -2417,6 +2422,10 @@ mod tests { query: ParsedPromQuery, _: QueryContextRef, ) -> Result { + self.ordered_outputs + .lock() + .unwrap() + .push(query.requires_output_ordering()); self.queries .lock() .unwrap() @@ -2531,6 +2540,59 @@ mod tests { ); } + #[tokio::test] + async fn range_query_does_not_require_execution_output_ordering() { + let handler = Arc::new(TestPrometheusHandler { + catalog_manager: MemoryCatalogManager::new(), + deny_operation: false, + denied_table: None, + metric_names: Vec::new(), + queries: Mutex::new(Vec::new()), + ordered_outputs: Mutex::new(Vec::new()), + }); + let state: PrometheusHandlerRef = handler.clone(); + instant_query( + State(state.clone()), + Query(InstantQuery { + query: Some("sort(vector(1))".to_string()), + time: Some("0".to_string()), + ..Default::default() + }), + Extension(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + )), + Form(InstantQuery::default()), + ) + .await; + + // Both a single-point and a multi-step range query take the same path. + for end in ["0", "1"] { + range_query( + State(state.clone()), + Query(RangeQuery { + query: Some("sort(vector(1))".to_string()), + start: Some("0".to_string()), + end: Some(end.to_string()), + step: Some("1s".to_string()), + ..Default::default() + }), + Extension(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + )), + Form(RangeQuery::default()), + ) + .await; + } + + // `sort()` stays observable for instant queries, but not for range queries. + assert_eq!( + *handler.ordered_outputs.lock().unwrap(), + vec![true, false, false] + ); + } + #[tokio::test] async fn test_promql_timer_records_parse_errors() { let handler: PrometheusHandlerRef = Arc::new(TestPrometheusHandler { @@ -2539,6 +2601,7 @@ mod tests { denied_table: None, metric_names: Vec::new(), queries: Mutex::new(Vec::new()), + ordered_outputs: Mutex::new(Vec::new()), }); let query_ctx = QueryContext::with("promql_timer_test", "parse_error"); let db = query_ctx.get_db_string(); @@ -2622,6 +2685,7 @@ mod tests { denied_table: Some("denied"), metric_names: Vec::new(), queries: Mutex::new(Vec::new()), + ordered_outputs: Mutex::new(Vec::new()), })), Path(FIELD_NAME_LABEL.to_string()), Extension(query_ctx), @@ -2665,6 +2729,7 @@ mod tests { denied_table: None, metric_names: vec!["cpu_user".to_string(), "cpu_system".to_string()], queries: Mutex::new(Vec::new()), + ordered_outputs: Mutex::new(Vec::new()), }); let state: PrometheusHandlerRef = handler.clone(); let query_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME); @@ -3571,6 +3636,7 @@ mod tests { denied_table: None, metric_names: Vec::new(), queries: Mutex::new(Vec::new()), + ordered_outputs: Mutex::new(Vec::new()), })), Query(MetadataQuery::default()), Extension(query_ctx.clone()), @@ -3584,6 +3650,7 @@ mod tests { denied_table: Some("denied"), metric_names: Vec::new(), queries: Mutex::new(Vec::new()), + ordered_outputs: Mutex::new(Vec::new()), }); let response = metadata_query( State(handler.clone()), diff --git a/src/servers/src/http/result/prometheus_resp.rs b/src/servers/src/http/result/prometheus_resp.rs index bc1d0849037..eca38338719 100644 --- a/src/servers/src/http/result/prometheus_resp.rs +++ b/src/servers/src/http/result/prometheus_resp.rs @@ -15,8 +15,12 @@ //! prom supply the prometheus HTTP API Server compliance use std::cmp::Ordering; use std::collections::{BTreeMap, HashMap}; +use std::hash::BuildHasher; +use std::ops::Range; -use arrow::array::{Array, AsArray, StructArray}; +use arrow::array::{Array, ArrayRef, AsArray, StructArray}; +use arrow::buffer::BooleanBuffer; +use arrow::compute::kernels::cmp::distinct; use arrow::datatypes::{Float64Type, TimestampMillisecondType}; use arrow_schema::DataType; use axum::Json; @@ -36,10 +40,13 @@ use common_recordbatch::RecordBatches; use datatypes::arrow_array::string_array_value_at_index; use datatypes::prelude::ConcreteDataType; use indexmap::IndexMap; +use indexmap::map::RawEntryApiV1; +use indexmap::map::raw_entry_v1::RawEntryMut; +use itertools::Either; use promql_parser::label::METRIC_NAME; use promql_parser::parser::value::ValueType; use ryu::Buffer; -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Serialize, Serializer}; use serde_json::Value; use snafu::{OptionExt, ResultExt}; @@ -55,10 +62,45 @@ use crate::http::prometheus::{ #[derive(Default)] struct PromSeriesSamples { - values: Vec<(f64, String)>, + values: Vec<(f64, PromSampleValue)>, histograms: Vec<(f64, PromNativeHistogram)>, } +/// A sample value of the Prometheus HTTP API JSON format. +/// +/// Samples read out of a query result are kept as `f64` and formatted while the +/// response is serialized, which avoids one `String` per sample. Samples parsed +/// from a JSON body keep their original spelling, so a response that is +/// deserialized and serialized again is unchanged. +#[derive(Debug, Clone, Deserialize, PartialEq)] +#[serde(untagged)] +pub enum PromSampleValue { + #[serde(skip_deserializing)] + Number(f64), + Text(String), +} + +impl Serialize for PromSampleValue { + fn serialize(&self, serializer: S) -> std::result::Result { + match self { + Self::Number(value) if value.is_finite() => { + serializer.serialize_str(Buffer::new().format_finite(*value)) + } + Self::Number(value) => serializer.collect_str(value), + Self::Text(value) => serializer.serialize_str(value), + } + } +} + +impl PromSampleValue { + fn into_string(self) -> String { + match self { + Self::Number(value) => format_prometheus_sample_value(value), + Self::Text(value) => value, + } + } +} + fn prometheus_native_histogram(histogram: &NativeHistogram) -> Result { Ok(PromNativeHistogram { count: format_prometheus_float(histogram.count), @@ -321,14 +363,8 @@ impl PrometheusJsonResponse { // Tag order matters, e.g., after sorc and sort_desc, the output order must be kept. let mut buffer = IndexMap::, PromSeriesSamples>::new(); - // Query output is clustered by series (the range plan sorts by series - // key + timestamp), so consecutive rows usually belong to the same - // series. Remember the index of the previous row's entry in `buffer`, - // and reuse it directly when its tags are unchanged. This avoids - // building and hashing the label vector on every row; the worst case - // adds one `Vec` comparison per series transition before falling back - // to the map lookup. - let mut last_entry_index = None; + // Only a series that is new to `buffer` needs its own key vector. + let mut tags = Vec::with_capacity(num_label_columns + 1); let schema = batches.schema(); for batch in batches.iter() { @@ -364,72 +400,95 @@ impl PrometheusJsonResponse { }) .transpose()?; + // Read the labels once per run of rows that share them instead of + // once per row. `label_runs` finds every boundary, so the probe only + // decides whether looking for runs is worth its cost. + let label_runs = if prefer_label_runs(&tag_columns, batch.num_rows()) { + Either::Left(label_runs(&tag_columns, batch.num_rows())?.into_iter()) + } else { + Either::Right((0..batch.num_rows()).map(|row| row..row + 1)) + }; + // assemble rows - for row_index in 0..batch.num_rows() { - let value = field_column.and_then(|field_column| { - if !field_column.is_valid(row_index) { - return None; + for run in label_runs { + let mut run_entry_index = None; + for row_index in run { + let value = field_column.and_then(|field_column| { + if !field_column.is_valid(row_index) { + return None; + } + let value = field_column.value(row_index); + (!is_prometheus_stale_nan(value)) + .then_some((timestamp_column.value(row_index), value)) + }); + let histogram = native_histogram_column + .and_then(|column| { + read_histogram(column, row_index) + .context(DataFusionSnafu) + .transpose() + }) + .transpose()? + .filter(|histogram| !is_prometheus_stale_nan(histogram.sum)) + .map(|histogram| { + prometheus_native_histogram(&histogram) + .map(|histogram| (timestamp_column.value(row_index), histogram)) + }) + .transpose()?; + + if value.is_none() && histogram.is_none() { + continue; } - let value = field_column.value(row_index); - (!is_prometheus_stale_nan(value)) - .then_some((timestamp_column.value(row_index), value)) - }); - let histogram = native_histogram_column - .and_then(|column| { - read_histogram(column, row_index) - .context(DataFusionSnafu) - .transpose() - }) - .transpose()? - .filter(|histogram| !is_prometheus_stale_nan(histogram.sum)) - .map(|histogram| { - prometheus_native_histogram(&histogram) - .map(|histogram| (timestamp_column.value(row_index), histogram)) - }) - .transpose()?; - if value.is_none() && histogram.is_none() { - continue; - } + let entry_index = match run_entry_index { + Some(index) => index, + None => { + // retrieve tags + tags.clear(); + if let Some(metric_name) = &metric_name { + tags.push((METRIC_NAME, metric_name.as_str())); + } + for (tag_column, tag_name) in tag_columns.iter().zip(tag_names.iter()) { + if let Some(tag_value) = + string_array_value_at_index(tag_column, row_index) + { + tags.push((tag_name, tag_value)); + } + } - // retrieve tags - let mut tags = Vec::with_capacity(num_label_columns + 1); - if let Some(metric_name) = &metric_name { - tags.push((METRIC_NAME, metric_name.as_str())); - } - for (tag_column, tag_name) in tag_columns.iter().zip(tag_names.iter()) { - if let Some(tag_value) = string_array_value_at_index(tag_column, row_index) { - tags.push((tag_name, tag_value)); + let hash = buffer.hasher().hash_one(&tags); + let entry = buffer + .raw_entry_mut_v1() + .from_key_hashed_nocheck(hash, &tags); + let index = entry.index(); + if let RawEntryMut::Vacant(entry) = entry { + // Hand the key buffer over with the hash that was + // just computed for the lookup. + let key = std::mem::replace( + &mut tags, + Vec::with_capacity(num_label_columns + 1), + ); + entry.insert_hashed_nocheck( + hash, + key, + PromSeriesSamples::default(), + ); + } + run_entry_index = Some(index); + index + } + }; + let samples = &mut buffer[entry_index]; + if let Some((timestamp_millis, histogram)) = histogram { + samples + .histograms + .push((timestamp_millis as f64 / 1000.0, histogram)); + } else if let Some((timestamp_millis, value)) = value { + samples.values.push(( + timestamp_millis as f64 / 1000.0, + PromSampleValue::Number(value), + )); } } - - let reuse = last_entry_index.filter(|index| { - buffer - .get_index(*index) - .is_some_and(|(key, _)| key == &tags) - }); - let samples = if let Some(index) = reuse { - buffer - .get_index_mut(index) - .map(|(_, samples)| samples) - .with_context(|| UnexpectedResultSnafu { - reason: "reused series entry must exist", - })? - } else { - let entry = buffer.entry(tags); - last_entry_index = Some(entry.index()); - entry.or_default() - }; - if let Some((timestamp_millis, histogram)) = histogram { - samples - .histograms - .push((timestamp_millis as f64 / 1000.0, histogram)); - } else if let Some((timestamp_millis, value)) = value { - samples.values.push(( - timestamp_millis as f64 / 1000.0, - format_prometheus_sample_value(value), - )); - } } } @@ -451,7 +510,10 @@ impl PrometheusJsonResponse { PromQueryResult::Vector(ref mut v) => { let histogram = samples.histograms.pop(); let value = if histogram.is_none() { - samples.values.pop() + samples + .values + .pop() + .map(|(timestamp, value)| (timestamp, value.into_string())) } else { None }; @@ -481,7 +543,10 @@ impl PrometheusJsonResponse { }); } PromQueryResult::Scalar(ref mut v) => { - *v = samples.values.pop(); + *v = samples + .values + .pop() + .map(|(timestamp, value)| (timestamp, value.into_string())); } PromQueryResult::String(ref mut _v) => { // TODO(ruihang): Not supported yet @@ -505,6 +570,71 @@ impl PrometheusJsonResponse { } } +/// Ranges of consecutive rows whose label values are all equal. +/// +/// `arrow::compute::partition` computes the same ranges, but its contract takes +/// lexicographically sorted columns, which query output is not: range queries +/// run without the plan's output sort, and `sort`/`topk` order by value. +/// `distinct` is element-wise, so it holds for any row order, and its null +/// handling is the one a series key needs: a null label and an empty one are +/// distinct, and two nulls are not. +fn label_runs(columns: &[&ArrayRef], rows: usize) -> Result>> { + let mut boundaries: Option = None; + if rows >= 2 { + for column in columns { + let changed = distinct(&column.slice(0, rows - 1), &column.slice(1, rows - 1)) + .context(ArrowSnafu)?; + boundaries = Some(match boundaries { + Some(accumulated) => &accumulated | changed.values(), + None => changed.values().clone(), + }); + } + } + + let mut runs = Vec::new(); + let mut start = 0; + for boundary in boundaries.iter().flat_map(BooleanBuffer::set_indices) { + runs.push(start..boundary + 1); + start = boundary + 1; + } + runs.push(start..rows); + Ok(runs) +} + +/// Decides whether to group the rows of a batch into runs that share their +/// label values. +/// +/// Sampling adjacent row pairs keeps the decision independent of the batch +/// size. Probe positions come from a xorshift sequence instead of a fixed +/// stride, which would alias with periodic series layouts. A wrong guess only +/// costs time: `partition` still validates every boundary, and the row-by-row +/// path builds the labels of every row. +fn prefer_label_runs(columns: &[&ArrayRef], rows: usize) -> bool { + if columns.is_empty() { + return true; + } + if rows < 2 { + return false; + } + let mut position = 0x9e37_79b9_u32; + let mut changes = 0; + for _ in 0..8 { + position ^= position << 13; + position ^= position >> 17; + position ^= position << 5; + let row = position as usize % (rows - 1); + if columns.iter().any(|column| { + string_array_value_at_index(column, row) != string_array_value_at_index(column, row + 1) + }) { + changes += 1; + if changes > 2 { + return false; + } + } + } + true +} + fn merge_annotations(target: &mut Option>, source: Option>) { let Some(source) = source else { return; @@ -665,6 +795,101 @@ mod tests { ); } + #[test] + fn sample_value_serialization_matches_eager_formatting() { + let mut values = vec![ + 0.0, + -0.0, + f64::MAX, + f64::MIN, + f64::MIN_POSITIVE, + f64::from_bits(1), + 1e-7, + 1e21, + f64::NAN, + f64::INFINITY, + f64::NEG_INFINITY, + ]; + let mut bits = 0x1234_5678_9876_5432_u64; + for _ in 0..1000 { + bits ^= bits << 13; + bits ^= bits >> 7; + bits ^= bits << 17; + values.push(f64::from_bits(bits)); + } + for value in values { + assert_eq!( + serde_json::to_string(&PromSampleValue::Number(value)).unwrap(), + serde_json::to_string(&format_prometheus_sample_value(value)).unwrap() + ); + } + for value in ["1.00", "+Inf", "-0", "NaN", "not-a-number", "", "\"\\\n"] { + let json = serde_json::to_string(value).unwrap(); + let parsed: PromSampleValue = serde_json::from_str(&json).unwrap(); + assert_eq!(serde_json::to_string(&parsed).unwrap(), json); + } + for value in ["1", "null", "true", "[]", "{}"] { + assert!(serde_json::from_str::(value).is_err()); + } + } + + #[tokio::test] + async fn matrix_response_body_matches_eagerly_formatted_json() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), false), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), false), + ])); + let batches = RecordBatches::try_new( + schema.clone(), + vec![ + RecordBatch::new( + schema, + vec![ + Arc::new(TimestampMillisecondVector::from_values([ + 1000, 2000, 3000, 4000, 5000, + ])) as _, + Arc::new(StringVector::from(vec![Some("a"); 5])) as _, + Arc::new(Float64Vector::from_values([ + -0.0, + f64::NAN, + 1e-7, + f64::INFINITY, + f64::NEG_INFINITY, + ])) as _, + ], + ) + .unwrap(), + ], + ) + .unwrap(); + let actual = PrometheusJsonResponse::from_query_result( + Ok(Output::new_with_record_batches(batches)), + None, + ValueType::Matrix, + None, + ) + .await; + // Deserializing the expectation yields `PromSampleValue::Text`, so this + // compares the deferred numeric encoding against eagerly built strings. + let expected: PrometheusJsonResponse = serde_json::from_value(serde_json::json!({ + "status": "success", + "data": {"resultType": "matrix", "result": [{ + "metric": {"host": "a"}, + "values": [[1.0, "-0.0"], [2.0, "NaN"], [3.0, "1e-7"], [4.0, "inf"], [5.0, "-inf"]] + }]} + })) + .unwrap(); + assert_eq!( + serde_json::to_string(&actual).unwrap(), + serde_json::to_string(&expected).unwrap() + ); + } + #[test] fn matrix_response_preserves_ordinary_nan_and_filters_stale_markers() { let schema = Arc::new(Schema::new(vec![ @@ -704,8 +929,8 @@ mod tests { assert_eq!(series.len(), 1); assert_eq!( - series[0].values, - vec![(1.0, "1.0".to_string()), (2.0, "NaN".to_string())] + serde_json::to_value(&series[0].values).unwrap(), + serde_json::json!([[1.0, "1.0"], [2.0, "NaN"]]) ); } @@ -766,7 +991,10 @@ mod tests { ((index + 1) as f64, expected_value) }) .collect::>(); - assert_eq!(series[0].values, expected); + assert_eq!( + serde_json::to_value(&series[0].values).unwrap(), + serde_json::to_value(expected).unwrap() + ); } #[test] @@ -809,21 +1037,16 @@ mod tests { assert_eq!(series.len(), 1); assert_eq!( - series[0].values, - vec![ - (1.0, "inf".to_string()), - (2.0, "-inf".to_string()), - (3.0, "NaN".to_string()), - ] + serde_json::to_value(&series[0].values).unwrap(), + serde_json::json!([[1.0, "inf"], [2.0, "-inf"], [3.0, "NaN"]]) ); } #[test] - fn record_batches_to_data_reuses_entries_for_clustered_series() { - // Rows are clustered by series (a, b, a, a, b, c): consecutive rows of - // the same series exercise the entry-reuse fast path, while series - // transitions fall back to the map lookup. The result must keep the - // first-occurrence order and accumulate values per series as before. + fn record_batches_to_data_groups_clustered_series() { + // Rows are clustered by series (a, b, a, a, b, c) and `a` is revisited + // after `b`. The result must keep the first-occurrence order and + // accumulate values per series. let schema = Arc::new(Schema::new(vec![ ColumnSchema::new( "timestamp", @@ -889,6 +1112,286 @@ mod tests { assert_eq!(series[2].value, Some((6.0, "6.0".to_string()))); } + #[test] + fn label_strategy_switches_preserve_all_rows_when_probes_miss_changes() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), false), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), false), + ])); + let mut batches = Vec::new(); + let mut expected = BTreeMap::>::new(); + for layout in 0..3 { + let labels = (0..1024) + .map(|row| { + let alternating = match layout { + 0 => true, + 1 => row < 64, + _ => row >= 64, + }; + if alternating && row % 2 != 0 { + "b" + } else { + "a" + } + }) + .collect::>(); + for (row, label) in labels.iter().enumerate() { + let value = (layout * 1024 + row) as f64; + expected + .entry((*label).to_string()) + .or_default() + .push((value, PromSampleValue::Number(value))); + } + let batch = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_values( + (0..1024).map(|row| (layout * 1024 + row) as i64 * 1000), + )) as _, + Arc::new(StringVector::from( + labels.into_iter().map(Some).collect::>(), + )) as _, + Arc::new(Float64Vector::from( + (0..1024) + .map(|row| Some((layout * 1024 + row) as f64)) + .collect::>(), + )) as _, + ], + ) + .unwrap(); + // The middle batch alternates only over a prefix the probes miss, + // so it takes the run path even though most rows are not runs. + assert_eq!( + prefer_label_runs(&[batch.column(1)], batch.num_rows()), + layout == 1 + ); + batches.push(batch); + } + let response = PrometheusJsonResponse::record_batches_to_data( + RecordBatches::try_new(schema, batches).unwrap(), + None, + ValueType::Matrix, + ) + .unwrap(); + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Matrix(series), + .. + }) = response + else { + panic!("expected matrix response"); + }; + assert_eq!(series.len(), expected.len()); + for series in series { + assert_eq!(series.metric.len(), 1); + assert!(series.histograms.is_empty()); + assert_eq!( + series.values, + expected.remove(&series.metric["host"]).unwrap() + ); + } + assert!(expected.is_empty()); + } + + #[test] + fn label_runs_keep_null_and_empty_labels_apart_across_batches() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("rack", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), false), + ])); + // Two batches of two runs each, with only `rack` changing: a null label + // and an empty one must not share a run, and the run opening the second + // batch continues the series that ended the first one. + let mut batches = Vec::new(); + let mut expected = vec![Vec::new(), Vec::new()]; + for (batch_index, leading_null) in [true, false].into_iter().enumerate() { + let mut racks = Vec::new(); + let mut values = Vec::new(); + for row in 0..1024 { + let value = (batch_index * 1024 + row) as f64; + let null_rack = (row < 512) == leading_null; + racks.push((!null_rack).then_some("")); + values.push(Some(value)); + expected[usize::from(!null_rack)].push((value, PromSampleValue::Number(value))); + } + batches.push( + RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_values( + values.iter().map(|value| value.unwrap() as i64 * 1000), + )) as _, + Arc::new(StringVector::from(vec![Some("a"); 1024])) as _, + Arc::new(StringVector::from(racks)) as _, + Arc::new(Float64Vector::from(values)) as _, + ], + ) + .unwrap(), + ); + } + for batch in &batches { + assert!(prefer_label_runs( + &[batch.column(1), batch.column(2)], + batch.num_rows() + )); + } + for series in &mut expected { + series.sort_by(|left, right| left.0.total_cmp(&right.0)); + } + + let response = PrometheusJsonResponse::record_batches_to_data( + RecordBatches::try_new(schema, batches).unwrap(), + None, + ValueType::Matrix, + ) + .unwrap(); + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Matrix(series), + .. + }) = response + else { + panic!("expected matrix response"); + }; + assert_eq!(series.len(), 2); + assert_eq!( + series[0].metric, + BTreeMap::from([("host".into(), "a".into())]) + ); + assert_eq!(series[0].values, expected[0]); + assert_eq!( + series[1].metric, + BTreeMap::from([("host".into(), "a".into()), ("rack".into(), "".into())]) + ); + assert_eq!(series[1].values, expected[1]); + } + + #[test] + fn matrix_response_is_independent_of_input_row_order() { + // Range queries run without the plan's output sort, so this function sees + // series interleaved across batches with timestamps out of order. The + // serialized matrix must be the same either way. + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new( + "timestamp", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("rack", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true), + ColumnSchema::new("histogram", native_histogram_value_type().clone(), true), + ])); + let histogram = |sum: f64| NativeHistogram { + sum, + ..sample_histogram() + }; + // timestamp, host, rack, float value, histogram value + type Row = ( + i64, + Option<&'static str>, + Option<&'static str>, + Option, + Option, + ); + let rows: Vec = vec![ + (1_000, Some("a"), Some("r"), Some(1.0), None), + (3_000, Some("a"), Some("r"), Some(3.0), None), + (2_000, Some("a"), Some("r"), Some(2.0), None), + (5_000, Some("a"), None, Some(5.0), None), + (4_000, Some("a"), None, Some(4.0), None), + (7_000, Some(""), None, Some(7.0), None), + (8_000, None, None, Some(8.0), None), + (6_000, None, None, Some(6.0), None), + (2_000, Some("h"), None, None, Some(histogram(20.0))), + (1_000, Some("h"), None, None, Some(histogram(10.0))), + ]; + let matrix = |order: &[usize], splits: &[usize]| { + let batch = RecordBatch::new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondVector::from_vec( + order.iter().map(|&row| rows[row].0).collect(), + )) as _, + Arc::new(StringVector::from( + order.iter().map(|&row| rows[row].1).collect::>(), + )) as _, + Arc::new(StringVector::from( + order.iter().map(|&row| rows[row].2).collect::>(), + )) as _, + Arc::new(Float64Vector::from( + order.iter().map(|&row| rows[row].3).collect::>(), + )) as _, + histogram_vector( + &order + .iter() + .map(|&row| rows[row].4.clone()) + .collect::>(), + ), + ], + ) + .unwrap(); + let mut batches = Vec::new(); + let mut start = 0; + for &end in splits.iter().chain(std::iter::once(&order.len())) { + batches.push(batch.slice(start, end - start).unwrap()); + start = end; + } + let response = PrometheusJsonResponse::record_batches_to_data( + RecordBatches::try_new(schema.clone(), batches).unwrap(), + Some("metric".to_string()), + ValueType::Matrix, + ) + .unwrap(); + let PrometheusResponse::PromData(PromData { + result: PromQueryResult::Matrix(series), + .. + }) = response + else { + panic!("expected matrix response"); + }; + series + }; + + let clustered = matrix(&[0, 2, 1, 4, 3, 5, 7, 6, 9, 8], &[5]); + let interleaved = matrix(&[8, 5, 1, 6, 0, 3, 9, 2, 7, 4], &[3, 6]); + assert_eq!( + serde_json::to_value(&interleaved).unwrap(), + serde_json::to_value(&clustered).unwrap() + ); + + // Pin the canonical arrangement itself, not only its stability. + assert_eq!( + serde_json::to_value(&clustered[..4]).unwrap(), + serde_json::json!([ + {"metric": {"__name__": "metric"}, "values": [[6.0, "6.0"], [8.0, "8.0"]]}, + {"metric": {"__name__": "metric", "host": ""}, "values": [[7.0, "7.0"]]}, + {"metric": {"__name__": "metric", "host": "a"}, + "values": [[4.0, "4.0"], [5.0, "5.0"]]}, + {"metric": {"__name__": "metric", "host": "a", "rack": "r"}, + "values": [[1.0, "1.0"], [2.0, "2.0"], [3.0, "3.0"]]}, + ]) + ); + assert_eq!(clustered[4].metric["host"], "h"); + assert_eq!( + clustered[4] + .histograms + .iter() + .map(|(timestamp, histogram)| (*timestamp, histogram.sum.as_str())) + .collect::>(), + vec![(1.0, "10"), (2.0, "20")] + ); + } + #[test] fn record_batches_to_data_preserves_mixed_float_and_histogram_rows() { let schema = Arc::new(Schema::new(vec![ diff --git a/src/servers/src/http/skip_wal.rs b/src/servers/src/http/skip_wal.rs new file mode 100644 index 00000000000..ea91552789c --- /dev/null +++ b/src/servers/src/http/skip_wal.rs @@ -0,0 +1,56 @@ +// 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 axum::body::Body; +use axum::http::Request; +use axum::middleware::Next; +use axum::response::{IntoResponse, Response}; +use session::context::QueryContext; + +use crate::error::InvalidParameterSnafu; +use crate::http::header::GREPTIME_INSERT_SKIP_WAL_HEADER_NAME; +use crate::http::result::error_result::ErrorResponse; + +/// Extract the request-level WAL policy from the dedicated HTTP header. +pub async fn extract_skip_wal(mut request: Request, next: Next) -> Response { + let skip_wal = match request.headers().get(&GREPTIME_INSERT_SKIP_WAL_HEADER_NAME) { + None => false, + Some(value) => match value + .to_str() + .ok() + .and_then(|value| value.parse::().ok()) + { + Some(skip_wal) => skip_wal, + None => { + return ( + http::StatusCode::BAD_REQUEST, + ErrorResponse::from_error( + InvalidParameterSnafu { + reason: format!( + "{} must be true or false", + GREPTIME_INSERT_SKIP_WAL_HEADER_NAME + ), + } + .build(), + ), + ) + .into_response(); + } + }, + }; + if let Some(query_ctx) = request.extensions_mut().get_mut::() { + query_ctx.set_skip_wal(skip_wal); + } + next.run(request).await +} diff --git a/src/servers/src/pending_rows_batcher.rs b/src/servers/src/pending_rows_batcher.rs index 03eb90dce28..f8c4a869d5e 100644 --- a/src/servers/src/pending_rows_batcher.rs +++ b/src/servers/src/pending_rows_batcher.rs @@ -223,6 +223,7 @@ struct BatchKey { catalog: String, schema: String, physical_table: String, + skip_wal: bool, } /// An aligned logical record batch and its timestamp column index. @@ -327,8 +328,7 @@ enum WorkerCommand { Ack { ack_tx: oneshot::Sender<()> }, } -// Batch key is derived from QueryContext; it assumes catalog/schema/physical_table fully -// define the write target and must remain consistent across the batch. +// Requests can share a batch only when their write target and WAL policy match. fn batch_key_from_ctx(ctx: &QueryContextRef) -> BatchKey { let physical_table = ctx .extension(PHYSICAL_TABLE_KEY) @@ -338,6 +338,7 @@ fn batch_key_from_ctx(ctx: &QueryContextRef) -> BatchKey { catalog: ctx.current_catalog().to_string(), schema: ctx.current_schema(), physical_table, + skip_wal: ctx.skip_wal(), } } @@ -1617,7 +1618,7 @@ pub async fn flush_batch_physical( )?; let resolved_batches = resolve_region_targets(planned_batches, partition_manager).await?; - let region_writes = encode_region_write_requests(resolved_batches)?; + let region_writes = encode_region_write_requests(resolved_batches, ctx.skip_wal())?; flush_region_writes_concurrently(node_manager, region_writes).await } @@ -1794,6 +1795,7 @@ async fn resolve_region_targets( fn encode_region_write_requests( resolved_batches: Vec, + skip_wal: bool, ) -> Result> { let mut region_writes = Vec::with_capacity(resolved_batches.len()); for resolved in resolved_batches { @@ -1811,6 +1813,7 @@ fn encode_region_write_requests( ..Default::default() }), body: Some(region_request::Body::BulkInsert(BulkInsertRequest { + skip_wal, region_id: region_id.as_u64(), partition_expr_version: None, // Set aligned_schema_version to None so that datanode will check the batch schema again to see if any @@ -1929,12 +1932,13 @@ mod tests { PendingRowsBatcher, PendingWorker, PhysicalFlushCatalogProvider, PhysicalFlushNodeRequester, PhysicalFlushPartitionProvider, PhysicalTableMetadata, PlannedRegionBatch, RecordBatchWithTsIdx, ResolvedRegionBatch, TableBatch, WorkerCommand, - columns_taxonomy, drain_batch, encode_region_write_requests, extract_timestamps, - flush_batch, flush_batch_physical, flush_region_writes_concurrently, greptime_timestamp, - notify_flow_dirty_windows_after_flush, plan_region_batches, remove_worker_if_same_channel, - should_close_worker_on_idle_timeout, should_dispatch_concurrently, - start_flow_notification_worker, start_worker, strip_partition_columns_from_batch, - transform_logical_batches_to_physical, try_enqueue_flow_notification, + batch_key_from_ctx, columns_taxonomy, drain_batch, encode_region_write_requests, + extract_timestamps, flush_batch, flush_batch_physical, flush_region_writes_concurrently, + greptime_timestamp, notify_flow_dirty_windows_after_flush, plan_region_batches, + remove_worker_if_same_channel, should_close_worker_on_idle_timeout, + should_dispatch_concurrently, start_flow_notification_worker, start_worker, + strip_partition_columns_from_batch, transform_logical_batches_to_physical, + try_enqueue_flow_notification, }; use crate::error; use crate::metrics::FLOW_NOTIFICATION_DROPPED; @@ -2901,6 +2905,31 @@ mod tests { } } + #[test] + fn test_batch_key_groups_by_skip_wal() { + let wal_ctx = session::context::QueryContext::arc(); + let skip_wal_ctx = session::context::QueryContext::arc(); + skip_wal_ctx.set_skip_wal(true); + let another_skip_wal_ctx = session::context::QueryContext::arc(); + another_skip_wal_ctx.set_skip_wal(true); + + let mut batches = HashMap::new(); + *batches.entry(batch_key_from_ctx(&wal_ctx)).or_insert(0) += 1; + *batches + .entry(batch_key_from_ctx(&skip_wal_ctx)) + .or_insert(0) += 1; + *batches + .entry(batch_key_from_ctx(&another_skip_wal_ctx)) + .or_insert(0) += 1; + *batches + .entry(batch_key_from_ctx(&session::context::QueryContext::arc())) + .or_insert(0) += 1; + + assert_eq!(batches.len(), 2); + assert_eq!(batches[&batch_key_from_ctx(&wal_ctx)], 2); + assert_eq!(batches[&batch_key_from_ctx(&skip_wal_ctx)], 2); + } + #[test] fn test_remove_worker_if_same_channel_removes_matching_entry() { let workers = DashMap::new(); @@ -2908,6 +2937,7 @@ mod tests { catalog: "greptime".to_string(), schema: "public".to_string(), physical_table: "phy".to_string(), + skip_wal: false, }; let (tx, _rx) = mpsc::channel::(1); @@ -2924,6 +2954,7 @@ mod tests { catalog: "greptime".to_string(), schema: "public".to_string(), physical_table: "phy".to_string(), + skip_wal: false, }; let (stale_tx, _stale_rx) = mpsc::channel::(1); @@ -3015,6 +3046,7 @@ mod tests { catalog: "greptime".to_string(), schema: "public".to_string(), physical_table: "phy".to_string(), + skip_wal: false, }; let workers = Arc::new(DashMap::new()); let (worker_tx, worker_rx) = mpsc::channel(1); @@ -3745,6 +3777,11 @@ mod tests { #[test] fn test_encode_region_write_requests_builds_bulk_insert_requests() { + check_encode_region_write_requests(false); + check_encode_region_write_requests(true); + } + + fn check_encode_region_write_requests(skip_wal: bool) { let planned_batch = PlannedRegionBatch { region_id: RegionId::new(1024, 1), batch: RecordBatch::try_new( @@ -3772,7 +3809,7 @@ mod tests { addr: "node-1".to_string(), }, }; - let writes = encode_region_write_requests(vec![resolved_batch]).unwrap(); + let writes = encode_region_write_requests(vec![resolved_batch], skip_wal).unwrap(); assert_eq!(1, writes.len()); assert_eq!(1, writes[0].datanode.id); @@ -3780,5 +3817,6 @@ mod tests { panic!("expected bulk insert request"); }; assert_eq!(RegionId::new(1024, 1).as_u64(), request.region_id); + assert_eq!(skip_wal, request.skip_wal); } } diff --git a/src/servers/src/prometheus_handler.rs b/src/servers/src/prometheus_handler.rs index d56668d0160..4f4b7c5d19c 100644 --- a/src/servers/src/prometheus_handler.rs +++ b/src/servers/src/prometheus_handler.rs @@ -40,6 +40,7 @@ pub type PrometheusHandlerRef = Arc; pub struct ParsedPromQuery { query: PromQuery, statement: QueryStatement, + requires_output_ordering: bool, } impl ParsedPromQuery { @@ -51,7 +52,11 @@ impl ParsedPromQuery { query: query.clone(), } })?; - Ok(Self { query, statement }) + Ok(Self { + query, + statement, + requires_output_ordering: true, + }) } /// Returns the original query parameters. @@ -64,6 +69,18 @@ impl ParsedPromQuery { &self.statement } + /// Returns whether the caller observes the query output in execution order. + /// When it does not, the executor may drop the plan's output sort. + pub fn requires_output_ordering(&self) -> bool { + self.requires_output_ordering + } + + /// Marks the output order as irrelevant to the caller. + pub(crate) fn with_unordered_output(mut self) -> Self { + self.requires_output_ordering = false; + self + } + /// Returns the parsed PromQL expression. pub fn expr(&self) -> &PromqlExpr { let QueryStatement::Promql(eval_stmt, _) = &self.statement else { diff --git a/src/servers/tests/http/http_handler_test.rs b/src/servers/tests/http/http_handler_test.rs index 7456d28b362..42f53c370c9 100644 --- a/src/servers/tests/http/http_handler_test.rs +++ b/src/servers/tests/http/http_handler_test.rs @@ -980,3 +980,102 @@ async fn get_body(response: Response) -> Bytes { .await .unwrap() } + +#[tokio::test] +async fn test_http_skip_wal_header() { + struct SkipWalSqlHandler { + inner: ServerSqlQueryHandlerRef, + observed: AtomicUsize, + } + + #[async_trait] + impl SqlQueryHandler for SkipWalSqlHandler { + async fn do_query(&self, query: &str, ctx: QueryContextRef) -> Vec> { + self.observed + .store(usize::from(ctx.skip_wal()), Ordering::Relaxed); + self.inner.do_query(query, ctx).await + } + + async fn do_analyze_stream_query( + &self, + query: &str, + ctx: QueryContextRef, + ) -> Result { + self.inner.do_analyze_stream_query(query, ctx).await + } + + async fn do_exec_plan( + &self, + plan: LogicalPlan, + stmt: Option, + ctx: QueryContextRef, + ) -> Result { + self.inner.do_exec_plan(plan, stmt, ctx).await + } + + async fn do_promql_query( + &self, + query: &PromQuery, + ctx: QueryContextRef, + ) -> Vec> { + self.inner.do_promql_query(query, ctx).await + } + + async fn do_describe( + &self, + stmt: Statement, + ctx: QueryContextRef, + ) -> Result> { + self.inner.do_describe(stmt, ctx).await + } + + async fn is_valid_schema(&self, catalog: &str, schema: &str) -> Result { + self.inner.is_valid_schema(catalog, schema).await + } + } + + let handler = Arc::new(SkipWalSqlHandler { + inner: create_testing_sql_query_handler(MemTable::default_numbers_table()), + observed: AtomicUsize::new(2), + }); + let server = HttpServerBuilder::new(HttpOptions::default()) + .with_sql_handler(handler.clone()) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + for (header, expected) in [ + (Some(("x-greptime-insert-skip-wal", "true")), Some(true)), + (Some(("x-greptime-insert-skip-wal", "false")), Some(false)), + (None, Some(false)), + (Some(("x-greptime-insert-skip-wal", "yes")), None), + (Some(("x-greptime-insert-skip-wal", "1")), None), + (Some(("x-greptime-insert-skip-wal", "")), None), + (Some(("x-greptime-skip-wal", "true")), Some(false)), + (Some(("x-greptime-hints", "skip_wal=true")), Some(false)), + ( + Some(("x-greptime-hints", "insert_skip_wal=true")), + Some(false), + ), + ] { + // Sentinel also proves invalid values are rejected before the SQL handler. + handler.observed.store(2, Ordering::Relaxed); + let mut request = client.get("/v1/sql?sql=SELECT%201"); + if let Some((key, value)) = header { + request = request.header(key, value); + } + let response = request.send().await; + match expected { + Some(expected) => { + assert_eq!(response.status(), StatusCode::OK, "{header:?}"); + assert_eq!( + handler.observed.load(Ordering::Relaxed), + usize::from(expected), + "{header:?}" + ); + } + None => { + assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{header:?}"); + assert_eq!(handler.observed.load(Ordering::Relaxed), 2); + } + } + } +} diff --git a/src/session/src/context.rs b/src/session/src/context.rs index d383aef477e..2e8b49dae37 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -153,6 +153,15 @@ impl QueryContextBuilder { self } + pub fn skip_wal(mut self, skip_wal: bool) -> Self { + self.mutable_session_data + .get_or_insert_default() + .write() + .unwrap() + .skip_wal = skip_wal; + self + } + pub fn read_preference(mut self, read_preference: ReadPreference) -> Self { self.mutable_session_data .get_or_insert_default() @@ -348,6 +357,15 @@ impl QueryContext { self.mutable_session_data.write().unwrap().timezone = timezone; } + /// Returns whether ordinary inserts in this request should skip WAL. + pub fn skip_wal(&self) -> bool { + self.mutable_session_data.read().unwrap().skip_wal + } + + pub fn set_skip_wal(&self, skip_wal: bool) { + self.mutable_session_data.write().unwrap().skip_wal = skip_wal; + } + pub fn read_preference(&self) -> ReadPreference { self.mutable_session_data.read().unwrap().read_preference } @@ -739,6 +757,37 @@ mod test { assert_eq!(fork.current_schema(), "private"); } + #[test] + fn test_skip_wal_default_builder_and_fork() { + let default_context = QueryContext::with(DEFAULT_CATALOG_NAME, "public"); + assert!(!default_context.skip_wal()); + assert!(!QueryContextBuilder::default().build().skip_wal()); + let context = QueryContextBuilder::default().skip_wal(true).build(); + assert!(context.skip_wal()); + let fork = context.fork(); + assert!(fork.skip_wal()); + fork.set_skip_wal(false); + assert!(context.skip_wal()); + assert!(!fork.skip_wal()); + context.set_skip_wal(false); + fork.set_skip_wal(true); + assert!(!context.skip_wal()); + assert!(fork.skip_wal()); + } + + #[test] + fn test_skip_wal_is_not_serialized_in_query_context() { + let context = QueryContextBuilder::default().skip_wal(true).build(); + let api_context: api::v1::QueryContext = context.into(); + assert!( + !api_context + .extensions + .contains_key(crate::hints::INSERT_SKIP_WAL_HINT) + ); + let restored: QueryContext = api_context.into(); + assert!(!restored.skip_wal()); + } + #[test] fn test_api_query_context_roundtrip_with_sequences() { let api_ctx = api::v1::QueryContext { diff --git a/src/session/src/hints.rs b/src/session/src/hints.rs index e656b34c82a..971426f4cc0 100644 --- a/src/session/src/hints.rs +++ b/src/session/src/hints.rs @@ -23,6 +23,9 @@ pub const SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY: &str = "query.support_flight_metrics_before_batch"; pub const LIVE_ANALYZE_METRICS_EXTENSION_KEY: &str = "query.live_analyze_metrics"; +/// Skip WAL for this insert only; never persisted as a table option. +pub const INSERT_SKIP_WAL_HINT: &str = "insert_skip_wal"; + pub const READ_PREFERENCE_HINT: &str = "read_preference"; pub const RESERVED_EXTENSION_KEYS: [&str; 4] = [ REMOTE_QUERY_ID_EXTENSION_KEY, diff --git a/src/session/src/lib.rs b/src/session/src/lib.rs index c1414fc5a3b..40d2e3e67fe 100644 --- a/src/session/src/lib.rs +++ b/src/session/src/lib.rs @@ -60,6 +60,8 @@ pub(crate) struct MutableInner { timezone: Timezone, query_timeout: Option, read_preference: ReadPreference, + /// Request-level WAL policy for ordinary inserts. + skip_wal: bool, #[debug(skip)] pub(crate) cursors: HashMap>, /// Warning messages for MySQL SHOW WARNINGS support @@ -74,6 +76,7 @@ impl Default for MutableInner { timezone: get_timezone(None).clone(), query_timeout: None, read_preference: ReadPreference::Leader, + skip_wal: false, cursors: HashMap::with_capacity(0), warnings: VecDeque::new(), } diff --git a/src/sql/src/parsers/alter_parser.rs b/src/sql/src/parsers/alter_parser.rs index 6c71702f5bf..d5a89f54076 100644 --- a/src/sql/src/parsers/alter_parser.rs +++ b/src/sql/src/parsers/alter_parser.rs @@ -29,7 +29,7 @@ use sqlparser::tokenizer::{Token, TokenWithSpan}; use crate::ast::ObjectNamePartExt; use crate::error::{self, InvalidColumnOptionSnafu, Result, SetFulltextOptionSnafu}; use crate::parser::ParserContext; -use crate::parsers::create_parser::INVERTED; +use crate::parsers::create_parser::{INVERTED, parse_json2_type_and_options}; use crate::parsers::utils::{ parse_with_options, validate_column_fulltext_create_option, validate_column_skipping_index_create_option, @@ -462,10 +462,19 @@ impl ParserContext<'_> { .context(error::SyntaxSnafu)?; self.parse_alter_table_drop_default(column_name) } else { - let data_type = self.parser.parse_data_type().context(error::SyntaxSnafu)?; + let (data_type, json2_options) = + if let Some(json2) = parse_json2_type_and_options(&mut self.parser)? { + json2 + } else { + ( + self.parser.parse_data_type().context(error::SyntaxSnafu)?, + None, + ) + }; Ok(AlterTableOperation::ModifyColumnType { column_name, target_type: data_type, + json2_options, }) } } @@ -1537,9 +1546,11 @@ ALTER TABLE metrics REPARTITION AlterTableOperation::ModifyColumnType { column_name, target_type, + json2_options, } => { assert_eq!("a", column_name.value); assert_eq!(DataType::String(None), *target_type); + assert!(json2_options.is_none()); } _ => unreachable!(), } @@ -1548,6 +1559,66 @@ ALTER TABLE metrics REPARTITION } } + #[test] + fn test_parse_alter_json2() { + let sql = r#"ALTER TABLE application_logs +MODIFY COLUMN attrs JSON2 ( + max_auto_expanded_paths = 2000, + trace_id STRING, + user.id STRING NOT NULL, + user.name STRING DEFAULT 'anonymous', + request_id STRING INVERTED INDEX +)"#; + let mut statements = + ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default()) + .unwrap(); + + let Statement::AlterTable(alter_table) = statements.remove(0) else { + unreachable!() + }; + let AlterTableOperation::ModifyColumnType { + column_name, + target_type, + json2_options: Some(options), + } = alter_table.alter_operation() + else { + unreachable!() + }; + + assert_eq!("attrs", column_name.value); + assert_eq!("JSON2", target_type.to_string()); + assert_eq!(Some(2000), options.max_auto_expanded_paths); + assert_eq!(4, options.type_hints.len()); + assert_eq!(vec!["user", "id"], options.type_hints[1].path); + assert!(!options.type_hints[1].nullable); + assert!(options.type_hints[2].default.is_some()); + assert!(options.type_hints[3].inverted_index); + + let formatted = alter_table.to_string(); + let reparsed = ParserContext::create_with_dialect( + &formatted, + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap(); + assert_eq!(Statement::AlterTable(alter_table), reparsed[0]); + + let mut empty = ParserContext::create_with_dialect( + "ALTER TABLE application_logs MODIFY COLUMN attrs JSON2 ()", + &GreptimeDbDialect {}, + ParseOptions::default(), + ) + .unwrap(); + let Statement::AlterTable(empty) = empty.remove(0) else { + unreachable!() + }; + let AlterTableOperation::ModifyColumnType { json2_options, .. } = empty.alter_operation() + else { + unreachable!() + }; + assert!(json2_options.is_none()); + } + #[test] fn test_parse_alter_change_column_alias_type() { let sql_1 = "ALTER TABLE my_metric_1 MODIFY COLUMN a MediumText"; @@ -1571,9 +1642,11 @@ ALTER TABLE metrics REPARTITION AlterTableOperation::ModifyColumnType { column_name, target_type, + json2_options, } => { assert_eq!("a", column_name.value); assert_eq!(DataType::MediumText, *target_type); + assert!(json2_options.is_none()); } _ => unreachable!(), } @@ -1602,9 +1675,11 @@ ALTER TABLE metrics REPARTITION AlterTableOperation::ModifyColumnType { column_name, target_type, + json2_options, } => { assert_eq!("a", column_name.value); assert!(matches!(target_type, DataType::Timestamp(Some(6), _))); + assert!(json2_options.is_none()); } _ => unreachable!(), } diff --git a/src/sql/src/parsers/create_parser.rs b/src/sql/src/parsers/create_parser.rs index 81ff23d3bc9..82dcd949ae4 100644 --- a/src/sql/src/parsers/create_parser.rs +++ b/src/sql/src/parsers/create_parser.rs @@ -24,6 +24,7 @@ use datafusion_common::ScalarValue; use datatypes::arrow::datatypes::{DataType as ArrowDataType, IntervalUnit}; use datatypes::data_type::ConcreteDataType; use itertools::Itertools; +pub(crate) use json::parse_json2_type_and_options; pub use json::parse_json2_type_hint_path; use snafu::{OptionExt, ResultExt, ensure}; use sqlparser::ast::{ diff --git a/src/sql/src/parsers/create_parser/json.rs b/src/sql/src/parsers/create_parser/json.rs index 2f8a80c79a7..fadc37c6d2f 100644 --- a/src/sql/src/parsers/create_parser/json.rs +++ b/src/sql/src/parsers/create_parser/json.rs @@ -49,7 +49,7 @@ pub fn parse_json2_type_hint_path(path: &str) -> Result> { Ok(path) } -pub(super) fn parse_json2_type_and_options( +pub(crate) fn parse_json2_type_and_options( parser: &mut Parser<'_>, ) -> Result)>> { let token = parser.peek_token(); diff --git a/src/sql/src/statements/alter.rs b/src/sql/src/statements/alter.rs index dd1f8e19d26..2f98d8737ae 100644 --- a/src/sql/src/statements/alter.rs +++ b/src/sql/src/statements/alter.rs @@ -26,7 +26,7 @@ use sqlparser::ast::{ColumnDef, DataType, Expr, Ident, ObjectName, TableConstrai use sqlparser_derive::{Visit, VisitMut}; use crate::statements::OptionMap; -use crate::statements::create::Partitions; +use crate::statements::create::{Json2Options, Partitions}; #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)] pub struct AlterTable { @@ -86,6 +86,7 @@ pub enum AlterTableOperation { ModifyColumnType { column_name: Ident, target_type: DataType, + json2_options: Option, }, /// `SET =
` SetTableOptions { @@ -256,8 +257,13 @@ impl Display for AlterTableOperation { AlterTableOperation::ModifyColumnType { column_name, target_type, + json2_options, } => { - write!(f, r#"MODIFY COLUMN {column_name} {target_type}"#) + write!(f, r#"MODIFY COLUMN {column_name} {target_type}"#)?; + if let Some(options) = json2_options { + write!(f, "{options}")?; + } + Ok(()) } AlterTableOperation::SetTableOptions { options } => { let kvs = options diff --git a/src/sql/src/statements/create.rs b/src/sql/src/statements/create.rs index aa39d42c291..a4736d0a7aa 100644 --- a/src/sql/src/statements/create.rs +++ b/src/sql/src/statements/create.rs @@ -141,7 +141,7 @@ pub struct ColumnExtensions { } /// JSON2-specific options represented in the SQL AST. -#[derive(Debug, PartialEq, Eq, Clone, Visit, VisitMut, Serialize)] +#[derive(Debug, PartialEq, Eq, Clone, Visit, VisitMut, Default, Serialize)] pub struct Json2Options { /// Maximum number of unhinted JSON2 paths expanded into Arrow fields. pub(crate) max_auto_expanded_paths: Option, diff --git a/src/store-api/src/mito_engine_options.rs b/src/store-api/src/mito_engine_options.rs index c308e915654..bc9b31ada95 100644 --- a/src/store-api/src/mito_engine_options.rs +++ b/src/store-api/src/mito_engine_options.rs @@ -88,6 +88,31 @@ pub const MAX_ROW_GROUP_ROW_COUNT: &str = "max_row_group_row_count"; pub const MAX_ROW_GROUP_ROW_COUNT_LIMIT: usize = 10 * 1024 * 1024; /// Option key for preserving per-row sequence numbers through flush and compaction. pub const PRESERVE_ROW_SEQUENCE: &str = "preserve_row_sequence"; +/// Option key for experimental Parquet float field encoding. +pub const EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING: &str = "experimental_sst_float_field_encoding"; + +/// Encoding policy for direct floating-point field columns in Parquet SSTs. +#[derive( + Debug, + Default, + Clone, + Copy, + PartialEq, + Eq, + Hash, + serde::Serialize, + serde::Deserialize, + strum::EnumString, +)] +#[serde(rename_all = "snake_case")] +#[strum(serialize_all = "snake_case")] +pub enum FloatFieldEncoding { + /// The existing Parquet writer behavior. + #[default] + Default, + /// Parquet byte-stream-split encoding. + ByteStreamSplit, +} // Note: Adding new options here should also check if this option should be removed in [metric_engine::engine::create::region_options_for_metadata_region]. /// Conflicting values supplied through the legacy and canonical TWCS trigger options. @@ -155,6 +180,7 @@ pub fn is_mito_engine_option_key(key: &str) -> bool { SST_FORMAT_KEY, MAX_ROW_GROUP_ROW_COUNT, PRESERVE_ROW_SEQUENCE, + EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, ] .contains(&key) } @@ -217,9 +243,29 @@ mod tests { assert!(is_mito_engine_option_key("append_mode")); assert!(is_mito_engine_option_key("max_row_group_row_count")); assert!(is_mito_engine_option_key("preserve_row_sequence")); + assert!(is_mito_engine_option_key( + EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING + )); assert!(!is_mito_engine_option_key("foo")); } + #[test] + fn test_float_field_encoding_values() { + assert_eq!( + "default".parse::(), + Ok(FloatFieldEncoding::Default) + ); + assert_eq!( + "byte_stream_split".parse::(), + Ok(FloatFieldEncoding::ByteStreamSplit) + ); + assert_eq!( + serde_json::from_str::("\"byte_stream_split\"").unwrap(), + FloatFieldEncoding::ByteStreamSplit + ); + assert!(serde_json::from_str::("\"unknown\"").is_err()); + } + #[test] fn test_normalize_twcs_trigger_aliases_to_legacy_key() { let expected = HashMap::from([(TWCS_TRIGGER_FILE_NUM.to_string(), "4".to_string())]); diff --git a/src/store-api/src/region_request.rs b/src/store-api/src/region_request.rs index 137087bc579..6a40f3dbc6b 100644 --- a/src/store-api/src/region_request.rs +++ b/src/store-api/src/region_request.rs @@ -221,6 +221,7 @@ fn make_region_puts(inserts: InsertRequests) -> Result Result Result> { let region_id = request.region_id.into(); + let skip_wal = request.skip_wal; let partition_expr_version = request.partition_expr_version.map(|v| v.value); let aligned_schema_version = request.aligned_schema_version.map(|v| v.schema_version); let Some(Body::ArrowIpc(request)) = request.body else { @@ -492,6 +494,7 @@ fn make_region_bulk_inserts(request: BulkInsertRequest) -> Result, + /// Skip WAL for this insert without changing region options. + /// Metadata writes must not inherit this option from user inserts. + pub skip_wal: bool, /// Partition expression version for the region. pub partition_expr_version: Option, } @@ -1751,6 +1757,8 @@ pub struct RegionCatchupRequest { #[derive(Debug, Clone)] pub struct RegionBulkInsertsRequest { + /// Whether this request should skip WAL. + pub skip_wal: bool, pub region_id: RegionId, pub payload: DfRecordBatch, pub raw_data: ArrowIpc, @@ -1864,6 +1872,36 @@ mod tests { use super::*; use crate::metadata::RegionMetadataBuilder; + #[test] + fn test_make_region_puts_preserves_skip_wal() { + let region_id = RegionId::new(42, 3); + let rows = Rows::default(); + let requests = make_region_puts(InsertRequests { + requests: [false, true, false] + .into_iter() + .map(|skip_wal| api::v1::region::InsertRequest { + region_id: region_id.as_u64(), + rows: Some(rows.clone()), + partition_expr_version: Some(api::v1::PartitionExprVersion { value: 7 }), + skip_wal, + }) + .collect(), + }) + .unwrap(); + + assert_eq!(3, requests.len()); + for ((id, request), skip_wal) in requests.into_iter().zip([false, true, false]) { + assert_eq!(region_id, id); + let RegionRequest::Put(request) = request else { + panic!("expected a put request"); + }; + assert_eq!(rows, request.rows); + assert_eq!(skip_wal, request.skip_wal); + assert_eq!(Some(7), request.partition_expr_version); + assert!(request.hint.is_none()); + } + } + #[test] fn test_make_region_compact_with_time_range() { let requests = make_region_compact(CompactRequest { diff --git a/src/store-api/src/storage/file.rs b/src/store-api/src/storage/file.rs index 0dcfcdc8cdd..b0b76146825 100644 --- a/src/store-api/src/storage/file.rs +++ b/src/store-api/src/storage/file.rs @@ -51,6 +51,11 @@ impl FileId { pub fn as_bytes(&self) -> &[u8] { self.0.as_bytes() } + + /// Constructs a file id from its packed 16-byte UUID representation. + pub fn from_bytes(bytes: [u8; 16]) -> FileId { + FileId(Uuid::from_bytes(bytes)) + } } impl From for Uuid { diff --git a/src/table/src/metadata.rs b/src/table/src/metadata.rs index 69803c7855e..b664c5cfd55 100644 --- a/src/table/src/metadata.rs +++ b/src/table/src/metadata.rs @@ -40,9 +40,10 @@ use store_api::storage::{ColumnDescriptor, ColumnDescriptorBuilder, ColumnId}; use crate::error::{self, Result}; use crate::requests::{ AddColumnRequest, AlterKind, AnnotationContext, AnnotationFamily, AnnotationValidationError, - ModifyColumnTypeRequest, REPARTITION_COLUMN_HINT_KEY, SetDefaultRequest, SetIndexOption, - TableOptions, UnsetIndexOption, has_stable_string_form, parse_entity_columns, - parse_entity_option_key, validate_and_normalize_annotation, + ModifyColumnTypeRequest, REPARTITION_COLUMN_HINT_KEY, REPARTITION_PARTITION_NUM_HINT_KEY, + SetDefaultRequest, SetIndexOption, TableOptions, UnsetIndexOption, has_stable_string_form, + parse_entity_columns, parse_entity_option_key, validate_and_normalize_annotation, + validate_annotation_keys, }; use crate::table_reference::TableReference; @@ -461,13 +462,13 @@ impl TableMeta { family: AnnotationFamily, options: &[(String, String)], ) -> Result { - ensure!( - !family.requires_single_key() || options.len() == 1, + validate_annotation_keys(options.iter().map(|(key, _)| key.as_str())).map_err(|err| { error::InvalidAlterRequestSnafu { table: table_name, - err: family.mixed_batch_error(), + err: err.to_string(), } - ); + .build() + })?; let cx = AnnotationContext { schema: &self.schema, partition_key_indices: &self.partition_key_indices, @@ -516,13 +517,13 @@ impl TableMeta { family: AnnotationFamily, keys: &[String], ) -> Result { - ensure!( - !family.requires_single_key() || keys.len() == 1, + validate_annotation_keys(keys.iter().map(String::as_str)).map_err(|err| { error::InvalidAlterRequestSnafu { table: table_name, - err: family.mixed_batch_error(), + err: err.to_string(), } - ); + .build() + })?; let mut new_options = self.options.clone(); for key in keys { ensure!( @@ -1491,6 +1492,7 @@ impl TableInfo { pub fn to_region_options(&self) -> HashMap { let mut options = HashMap::from(&self.meta.options); options.remove(REPARTITION_COLUMN_HINT_KEY); + options.remove(REPARTITION_PARTITION_NUM_HINT_KEY); options } @@ -1943,6 +1945,160 @@ mod tests { ); } + #[test] + fn test_set_unset_repartition_hints_together() { + let meta = TableMetaBuilder::empty() + .schema(Arc::new(new_test_schema())) + .primary_key_indices(vec![0]) + .engine("engine") + .next_column_id(3) + .build() + .unwrap(); + let options = vec![ + (REPARTITION_COLUMN_HINT_KEY.to_string(), "col1".to_string()), + ( + REPARTITION_PARTITION_NUM_HINT_KEY.to_string(), + "10".to_string(), + ), + ]; + for reverse in [false, true] { + let mut options = options.clone(); + if reverse { + options.reverse(); + } + let updated = meta + .builder_with_alter_kind( + "t", + &AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: options.clone(), + }, + ) + .unwrap() + .build() + .unwrap(); + assert_eq!( + updated.options.extra_options, + options.iter().cloned().collect() + ); + for (key, value) in [ + (REPARTITION_COLUMN_HINT_KEY, "missing"), + (REPARTITION_PARTITION_NUM_HINT_KEY, "0"), + ] { + let invalid = options + .iter() + .map(|(k, v)| { + ( + k.clone(), + if k == key { + value.to_string() + } else { + v.clone() + }, + ) + }) + .collect(); + assert!( + updated + .builder_with_alter_kind( + "t", + &AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: invalid, + } + ) + .is_err() + ); + assert_eq!( + updated.options.extra_options, + options.iter().cloned().collect() + ); + } + let cleared = updated + .builder_with_alter_kind( + "t", + &AlterKind::UnsetAnnotations { + family: AnnotationFamily::RepartitionHint, + keys: options.into_iter().map(|(key, _)| key).collect(), + }, + ) + .unwrap() + .build() + .unwrap(); + assert!(cleared.options.extra_options.is_empty()); + } + } + + #[test] + fn test_repartition_partition_num_hint() { + for partition_key_indices in [vec![], vec![0]] { + let mut meta = TableMetaBuilder::empty() + .schema(Arc::new(new_test_schema())) + .primary_key_indices(vec![0]) + .partition_key_indices(partition_key_indices) + .engine("engine") + .next_column_id(3) + .build() + .unwrap(); + for value in ["", " ", "0", "-1", "1.5", "abc", "4294967296"] { + let err = meta + .builder_with_alter_kind( + "t", + &AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: vec![( + REPARTITION_PARTITION_NUM_HINT_KEY.to_string(), + value.to_string(), + )], + }, + ) + .err() + .unwrap(); + assert!( + err.to_string().contains("expects a positive integer"), + "{err}" + ); + } + for (value, expected) in [(" 8 ", "8"), ("1", "1"), ("4294967295", "4294967295")] { + meta = meta + .builder_with_alter_kind( + "t", + &AlterKind::SetAnnotations { + family: AnnotationFamily::RepartitionHint, + options: vec![( + REPARTITION_PARTITION_NUM_HINT_KEY.to_string(), + value.to_string(), + )], + }, + ) + .unwrap() + .build() + .unwrap(); + assert_eq!( + meta.options.extra_options, + HashMap::from([( + REPARTITION_PARTITION_NUM_HINT_KEY.to_string(), + expected.to_string() + )]) + ); + } + for _ in 0..2 { + meta = meta + .builder_with_alter_kind( + "t", + &AlterKind::UnsetAnnotations { + family: AnnotationFamily::RepartitionHint, + keys: vec![REPARTITION_PARTITION_NUM_HINT_KEY.to_string()], + }, + ) + .unwrap() + .build() + .unwrap(); + assert!(meta.options.extra_options.is_empty()); + } + } + } + #[test] fn test_set_repartition_column_hint() { let meta = TableMetaBuilder::empty() @@ -2141,11 +2297,15 @@ mod tests { } #[test] - fn test_repartition_column_hint_is_not_region_option() { + fn test_repartition_hints_are_not_region_options() { let mut table_options = TableOptions::default(); table_options .extra_options .insert(REPARTITION_COLUMN_HINT_KEY.to_string(), "col1".to_string()); + table_options.extra_options.insert( + REPARTITION_PARTITION_NUM_HINT_KEY.to_string(), + "8".to_string(), + ); let table_info = TableInfoBuilder::default() .table_id(1) .table_version(0) @@ -2165,6 +2325,11 @@ mod tests { .build() .unwrap(); + assert!( + !table_info + .to_region_options() + .contains_key(REPARTITION_PARTITION_NUM_HINT_KEY) + ); assert!( !table_info .to_region_options() @@ -2972,7 +3137,7 @@ mod tests { } #[test] - fn test_repartition_hint_batch_must_be_single_key() { + fn test_repartition_hint_batch_rejects_duplicate_keys() { let meta = TableMetaBuilder::empty() .schema(Arc::new(new_test_schema())) .primary_key_indices(vec![0]) @@ -2995,7 +3160,7 @@ mod tests { .err() .unwrap(); assert!( - err.to_string().contains("must be altered separately"), + err.to_string().contains("duplicate repartition hint keys"), "{err}" ); @@ -3011,7 +3176,7 @@ mod tests { .err() .unwrap(); assert!( - err.to_string().contains("must be altered separately"), + err.to_string().contains("duplicate repartition hint keys"), "{err}" ); } diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index 3c9d0679aed..a21a8f6fed8 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -36,10 +36,10 @@ use store_api::metric_engine_consts::{ LOGICAL_TABLE_METADATA_KEY, PHYSICAL_TABLE_METADATA_KEY, is_metric_engine_option_key, }; use store_api::mito_engine_options::{ - APPEND_MODE_KEY, COMPACTION_TYPE, MEMTABLE_BULK_ENCODE_BYTES_THRESHOLD, - MEMTABLE_BULK_ENCODE_ROW_THRESHOLD, MEMTABLE_BULK_MAX_MERGE_GROUPS, - MEMTABLE_BULK_MERGE_THRESHOLD, MEMTABLE_TYPE, MERGE_MODE_KEY, SST_FORMAT_KEY, - TWCS_ACTIVE_WINDOW_L1_MERGE_TRIGGER, TWCS_ACTIVE_WINDOW_TRIGGER_FILE_NUM, + APPEND_MODE_KEY, COMPACTION_TYPE, EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, FloatFieldEncoding, + MEMTABLE_BULK_ENCODE_BYTES_THRESHOLD, MEMTABLE_BULK_ENCODE_ROW_THRESHOLD, + MEMTABLE_BULK_MAX_MERGE_GROUPS, MEMTABLE_BULK_MERGE_THRESHOLD, MEMTABLE_TYPE, MERGE_MODE_KEY, + SST_FORMAT_KEY, TWCS_ACTIVE_WINDOW_L1_MERGE_TRIGGER, TWCS_ACTIVE_WINDOW_TRIGGER_FILE_NUM, TWCS_FALLBACK_TO_LOCAL, TWCS_INACTIVE_WINDOW_L1_MERGE_TRIGGER, TWCS_INACTIVE_WINDOW_TRIGGER_FILE_NUM, TWCS_MAX_OUTPUT_FILE_SIZE, TWCS_TIME_WINDOW, TWCS_TRIGGER_FILE_NUM, is_mito_engine_option_key, normalize_twcs_trigger_options, @@ -77,7 +77,7 @@ pub fn is_trace_v1_table(table_info: &crate::metadata::TableInfo) -> bool { pub const OTLP_METRIC_COMPAT_KEY: &str = "otlp_metric_compat"; pub const OTLP_METRIC_COMPAT_PROM: &str = "prom"; -pub const VALID_TABLE_OPTION_KEYS: [&str; 14] = [ +pub const VALID_TABLE_OPTION_KEYS: [&str; 15] = [ // common keys: WRITE_BUFFER_SIZE_KEY, TTL_KEY, @@ -96,6 +96,7 @@ pub const VALID_TABLE_OPTION_KEYS: [&str; 14] = [ TABLE_DATA_MODEL, OTLP_METRIC_COMPAT_KEY, REPARTITION_COLUMN_HINT_KEY, + REPARTITION_PARTITION_NUM_HINT_KEY, ]; pub const DDL_TIMEOUT: &str = "timeout"; @@ -209,6 +210,9 @@ pub const SKIP_WAL_KEY: &str = store_api::mito_engine_options::SKIP_WAL_KEY; pub const TRACE_TABLE_PARTITIONS_HINT_KEY: &str = "trace_table_partitions"; pub const REPARTITION_COLUMN_HINT_KEY: &str = "repartition.column.hint"; +/// Table-level partition count hint consumed by the auto-repartition planner. +pub const REPARTITION_PARTITION_NUM_HINT_KEY: &str = "repartition.partition.num.hint"; + impl TableOptions { pub fn try_from_iter>( iter: U, @@ -252,6 +256,16 @@ impl TableOptions { options.ttl = Some(ttl_value); } + if let Some(encoding) = kvs.get(EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING) { + encoding.parse::().map_err(|_| { + ParseTableOptionSnafu { + key: EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, + value: encoding, + } + .build() + })?; + } + if let Some(skip_wal) = kvs.get(SKIP_WAL_KEY) { options.skip_wal = skip_wal.parse().map_err(|_| { ParseTableOptionSnafu { @@ -350,24 +364,26 @@ pub struct ModifyColumnTypeRequest { pub enum AnnotationFamily { /// `greptime.semantic.*` options (see the [`semantic`] module). Semantic, - /// `repartition.column.hint`, consumed by the auto-repartition planner. + /// Column and partition count hints consumed by the auto-repartition planner. RepartitionHint, } impl AnnotationFamily { - /// The key namespace: a prefix for [`Self::Semantic`], the exact key for - /// [`Self::RepartitionHint`]. + /// The key namespace used in diagnostics; accepted keys are classified by [`Self::of_key`]. pub fn namespace(self) -> &'static str { match self { Self::Semantic => SEMANTIC_PREFIX, - Self::RepartitionHint => REPARTITION_COLUMN_HINT_KEY, + Self::RepartitionHint => "repartition.", } } pub fn of_key(key: &str) -> Option { if key.starts_with(SEMANTIC_PREFIX) { Some(Self::Semantic) - } else if key == REPARTITION_COLUMN_HINT_KEY { + } else if matches!( + key, + REPARTITION_COLUMN_HINT_KEY | REPARTITION_PARTITION_NUM_HINT_KEY + ) { Some(Self::RepartitionHint) } else { None @@ -384,10 +400,8 @@ impl AnnotationFamily { } } - /// Whether this family's SET/UNSET batch must contain exactly one key. - /// The repartition hint is a single marker; a batch with several hint - /// entries (duplicates included) has no meaningful order. - pub fn requires_single_key(self) -> bool { + /// Whether duplicate keys are rejected in this family's SET/UNSET batch. + pub fn requires_unique_keys(self) -> bool { matches!(self, Self::RepartitionHint) } @@ -398,12 +412,60 @@ impl AnnotationFamily { "`{SEMANTIC_PREFIX}*` options must be altered separately from other table options" ), Self::RepartitionHint => { - format!("{REPARTITION_COLUMN_HINT_KEY} must be altered separately") + "repartition hints must be altered separately from other table options".to_string() } } } } +/// Why an annotation SET/UNSET key batch was rejected. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AnnotationKeyError { + MixedFamilies { family: AnnotationFamily }, + DuplicateKey, +} + +impl fmt::Display for AnnotationKeyError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::MixedFamilies { family } => f.write_str(&family.mixed_batch_error()), + Self::DuplicateKey => f.write_str("duplicate repartition hint keys"), + } + } +} + +/// Validates a SET/UNSET batch and returns its annotation family. +/// +/// Empty batches and batches containing only non-annotation keys return `None`. +/// Annotation keys cannot share a batch with another family or region options. +/// Duplicate keys are rejected only for families that require unique keys. +pub fn validate_annotation_keys<'a>( + keys: impl IntoIterator, +) -> std::result::Result, AnnotationKeyError> { + let mut keys = keys.into_iter(); + let Some(first) = keys.next() else { + return Ok(None); + }; + let family = AnnotationFamily::of_key(first); + let reject_duplicates = family.is_some_and(|family| family.requires_unique_keys()); + let mut seen = HashSet::new(); + if reject_duplicates { + seen.insert(first); + } + for key in keys { + let this = AnnotationFamily::of_key(key); + if this != family + && let Some(family) = family.or(this) + { + return Err(AnnotationKeyError::MixedFamilies { family }); + } + if reject_duplicates && !seen.insert(key) { + return Err(AnnotationKeyError::DuplicateKey); + } + } + Ok(family) +} + /// Table shape an annotation option is validated against. pub struct AnnotationContext<'a> { pub schema: &'a Schema, @@ -431,6 +493,9 @@ pub enum AnnotationValidationError { column: String, ty: ConcreteDataType, }, + InvalidPartitionNumHint { + value: String, + }, NotSingleColumn, PartitionMetadataConflict, TimeIndexConflict, @@ -449,6 +514,10 @@ impl fmt::Display for AnnotationValidationError { "entity column `{column}` (option `{key}`) has type `{ty}`, \ which cannot render as a string" ), + Self::InvalidPartitionNumHint { value } => write!( + f, + "{REPARTITION_PARTITION_NUM_HINT_KEY} expects a positive integer within u32 range, got `{value}`" + ), Self::NotSingleColumn => write!( f, "{REPARTITION_COLUMN_HINT_KEY} expects exactly one column name" @@ -466,7 +535,7 @@ impl fmt::Display for AnnotationValidationError { } /// Validates one annotation option and returns the value to store — the -/// repartition hint is trimmed to the bare column name, semantic values pass +/// repartition hints are trimmed, semantic values pass /// through unchanged. pub(crate) fn validate_and_normalize_annotation( family: AnnotationFamily, @@ -505,6 +574,15 @@ pub(crate) fn validate_and_normalize_annotation( } Ok(value.to_string()) } + AnnotationFamily::RepartitionHint if key == REPARTITION_PARTITION_NUM_HINT_KEY => { + let value = value.trim(); + if !matches!(value.parse::(), Ok(1..)) { + return Err(AnnotationValidationError::InvalidPartitionNumHint { + value: value.to_string(), + }); + } + Ok(value.to_string()) + } AnnotationFamily::RepartitionHint => { let column_name = value.trim(); if column_name.is_empty() || column_name.contains(',') { @@ -646,6 +724,8 @@ pub struct InsertRequest { pub schema_name: String, pub table_name: String, pub columns_values: HashMap, + /// Whether this insert should skip WAL. + pub skip_wal: bool, } /// Delete (by primary key) request @@ -763,6 +843,49 @@ mod tests { use super::*; + #[test] + fn test_validate_annotation_keys() { + let column = REPARTITION_COLUMN_HINT_KEY; + let count = REPARTITION_PARTITION_NUM_HINT_KEY; + for (keys, expected) in [ + (vec![], None), + (vec![TTL_KEY, TTL_KEY], None), + (vec!["repartition.unknown.hint"], None), + (vec![column], Some(AnnotationFamily::RepartitionHint)), + (vec![count], Some(AnnotationFamily::RepartitionHint)), + (vec![column, count], Some(AnnotationFamily::RepartitionHint)), + (vec![count, column], Some(AnnotationFamily::RepartitionHint)), + ( + vec!["greptime.semantic.source", "greptime.semantic.source"], + Some(AnnotationFamily::Semantic), + ), + ] { + assert_eq!(validate_annotation_keys(keys), Ok(expected)); + } + for keys in [ + vec![column, column], + vec![count, count], + vec![column, count, column], + ] { + assert_eq!( + validate_annotation_keys(keys), + Err(AnnotationKeyError::DuplicateKey) + ); + } + for keys in [ + vec![column, TTL_KEY], + vec![TTL_KEY, count], + vec![column, "greptime.semantic.source"], + ] { + assert_eq!( + validate_annotation_keys(keys), + Err(AnnotationKeyError::MixedFamilies { + family: AnnotationFamily::RepartitionHint, + }) + ); + } + } + #[test] fn test_validate_table_option() { assert!(validate_table_option(FILE_TABLE_LOCATION_KEY)); @@ -772,7 +895,10 @@ mod tests { assert!(validate_table_option(WRITE_BUFFER_SIZE_KEY)); assert!(validate_table_option(STORAGE_KEY)); assert!(validate_table_option(MEMTABLE_BULK_MERGE_THRESHOLD)); + assert!(validate_table_option(EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING)); assert!(validate_table_option(REPARTITION_COLUMN_HINT_KEY)); + assert!(validate_table_option(REPARTITION_PARTITION_NUM_HINT_KEY)); + assert_eq!(AnnotationFamily::of_key("repartition.unknown.hint"), None); assert!(!validate_table_option("foo")); // Only whitelisted semantic keys are accepted. @@ -808,6 +934,25 @@ mod tests { assert!(!validate_database_option("foo")); } + #[test] + fn test_parse_float_field_encoding_option() { + let options = TableOptions::try_from_iter([( + EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, + "byte_stream_split", + )]) + .unwrap(); + assert_eq!( + options + .extra_options + .get(EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING), + Some(&"byte_stream_split".to_string()) + ); + assert!( + TableOptions::try_from_iter([(EXPERIMENTAL_SST_FLOAT_FIELD_ENCODING, "invalid",)]) + .is_err() + ); + } + #[test] fn test_database_trigger_value_boundaries() { let maximum = usize::MAX.to_string(); diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 5e2a63015b7..a2d6c70e394 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -106,6 +106,7 @@ hex.workspace = true http.workspace = true itertools.workspace = true jsonb.workspace = true +mito2 = { workspace = true, features = ["test"] } mysql_async = { version = "0.37", default-features = false, features = [ "time", "default-rustls-ring", diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index f37d5ee638a..2486dddee05 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -49,6 +49,7 @@ mod test { use futures_util::{Stream, StreamExt}; use hyper_util::rt::TokioIo; use itertools::Itertools; + use rstest::rstest; use servers::grpc::builder::GrpcServerBuilder; use servers::grpc::flight::{ FlightCraft, FlightCraftWrapper, FlightRecordBatchSource, FlightRecordBatchStream, @@ -59,6 +60,7 @@ mod test { use servers::query_handler::grpc::GrpcQueryHandler; use servers::server::Server; use session::context::QueryContextRef; + use session::hints::INSERT_SKIP_WAL_HINT; use tokio::net::TcpListener; use tokio_stream::wrappers::TcpListenerStream; use tonic::transport::Server as TonicServer; @@ -67,7 +69,7 @@ mod test { use crate::cluster::GreptimeDbClusterBuilder; use crate::grpc::query_and_expect; - use crate::test_util::{StorageType, setup_grpc_server}; + use crate::test_util::{MockInstanceImpl, StorageType, assert_wal_delta, setup_grpc_server}; use crate::tests::test_util::MockInstance; struct SlowFlightCraft; @@ -399,6 +401,68 @@ mod test { assert!(start.elapsed() >= Duration::from_secs(1)); assert!(matches!(output.data, OutputData::Stream(_))); } + + #[rstest] + #[case::standalone_single_region(false, false)] + #[case::standalone_partitioned(false, true)] + #[case::distributed_single_region(true, false)] + #[case::distributed_partitioned(true, true)] + #[tokio::test(flavor = "multi_thread")] + async fn test_flight_bulk_skip_wal(#[case] distributed: bool, #[case] partitioned: bool) { + let mut db = + MockInstanceImpl::new(&format!("flight_bulk_skip_wal_{partitioned}"), distributed) + .await; + let runtime = common_runtime::global_runtime().clone(); + let handler = GreptimeRequestHandler::new( + db.frontend(), + None, + Some(runtime.clone()), + FlightCompression::default(), + ); + let mut server = GrpcServerBuilder::new(GrpcServerConfig::default(), runtime) + .flight_handler(Arc::new(handler)) + .build(); + server.start("127.0.0.1:0".parse().unwrap()).await.unwrap(); + let client = Database::new( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + Client::with_urls(vec![server.bind_addr().unwrap().to_string()]), + ); + let partition = if partitioned { + " PARTITION ON COLUMNS (a) (a < 0, a >= 0)" + } else { + "" + }; + client.sql(&format!( + "CREATE TABLE foo (ts TIMESTAMP TIME INDEX, a INT NOT NULL, \"B\" STRING, PRIMARY KEY (a)){partition}" + )).await.unwrap(); + let mut before = db.flush_and_snapshot_wal().await; + assert_eq!(before.len(), if partitioned { 2 } else { 1 }); + // The same stream spans both partitions and includes multiple data messages. + // Repeated rows still exercise real writes; each round flushes before the next. + for hint in [None, Some("true"), Some("false"), None] { + let hints = hint.map(|value| [(INSERT_SKIP_WAL_HINT, value)]); + test_put_record_batches_with_hints( + &client, + create_record_batches(-4), + hints.as_ref().map_or(&[], |hints| hints.as_slice()), + ) + .await; + let after = db.flush_and_snapshot_wal().await; + assert_wal_delta(&before, &after, hint == Some("true")); + // Both regions must receive rows, not merely exist in the catalog. + for (id, &(flushed_sequence, _)) in &after { + assert!( + flushed_sequence > before[id].0, + "region {id:?} received no rows" + ); + } + before = after; + } + server.shutdown().await.unwrap(); + db.shutdown().await; + } + #[tokio::test(flavor = "multi_thread")] async fn test_standalone_flight_do_put() { common_telemetry::init_default_ut_logging(); @@ -735,6 +799,14 @@ mod test { } async fn test_put_record_batches(client: &Database, record_batches: Vec) { + test_put_record_batches_with_hints(client, record_batches, &[]).await; + } + + async fn test_put_record_batches_with_hints( + client: &Database, + record_batches: Vec, + hints: &[(&str, &str)], + ) { let requests_count = record_batches.len(); let schema = record_batches[0].schema.arrow_schema().clone(); @@ -767,7 +839,7 @@ mod test { ) .boxed(); - let response_stream = client.do_put(stream).await.unwrap(); + let response_stream = client.do_put_with_hints(stream, hints).await.unwrap(); let responses = response_stream.collect::>().await; let responses_count = responses.len(); diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index 0704b184e89..18e1d144024 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -62,6 +62,8 @@ use standalone::{StandaloneDatanodeManager, StandaloneRepartitionProcedureFactor use crate::test_util::{self, StorageType, TestGuard, create_tmp_dir_and_datanode_opts}; pub struct GreptimeDbStandalone { + /// Storage engine for assertions across protocol and storage boundaries. + pub mito_engine: mito2::engine::MitoEngine, pub frontend: Arc, pub opts: StandaloneOptions, pub guard: TestGuard, @@ -342,6 +344,7 @@ impl GreptimeDbStandaloneBuilder { }; GreptimeDbStandalone { + mito_engine: datanode.region_server().mito_engine().unwrap(), frontend: Arc::new(frontend), opts, guard, diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 9d12848ce61..e0a749b813f 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -12,6 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::BTreeMap; +#[cfg(test)] +use std::collections::HashMap; use std::env; use std::fmt::Display; use std::net::SocketAddr; @@ -26,7 +29,10 @@ use axum::Router; use catalog::kvbackend::KvBackendCatalogManager; use client::{Client, Database}; use common_base::Plugins; +use common_catalog::consts::MIN_USER_TABLE_ID; use common_config::Configurable; +#[cfg(test)] +use common_meta::DatanodeId; use common_meta::key::TableMetadataManager; use common_meta::key::catalog_name::CatalogNameKey; use common_meta::key::schema_name::SchemaNameKey; @@ -38,8 +44,13 @@ use common_test_util::ports; use common_test_util::temp_dir::{TempDir, create_temp_dir}; use common_wal::config::DatanodeWalConfig; use datanode::config::{DatanodeOptions, StorageConfig}; +#[cfg(test)] +use datanode::datanode::Datanode; use frontend::instance::Instance; use frontend::service_config::{MysqlOptions, PostgresOptions}; +#[cfg(test)] +use meta_srv::metasrv::Metasrv; +use mito2::engine::MitoEngine; use mito2::gc::GcConfig; use object_store::config::{ AzblobConfig, FileConfig, GcsConfig, ObjectStoreConfig, OssConfig, S3Config, @@ -62,9 +73,178 @@ use servers::request_memory_limiter::ServerMemoryLimiter; use servers::server::Server; use servers::tls::ReloadableTlsServerConfig; use session::context::QueryContext; +use store_api::metric_engine_consts::METRIC_METADATA_REGION_GROUP; +use store_api::region_engine::RegionEngine; +use store_api::region_request::{RegionFlushRequest, RegionRequest}; +use store_api::storage::RegionId; +use crate::cluster::{GreptimeDbCluster, GreptimeDbClusterBuilder}; use crate::standalone::{GreptimeDbStandalone, GreptimeDbStandaloneBuilder}; +/// Maps `(node_id, region_id)` to `(flushed_sequence, flushed_entry_id)`. +pub type WalSnapshot = BTreeMap<(u64, RegionId), (u64, u64)>; + +/// Flush user data regions before observing their persisted WAL watermarks. +/// Metadata regions write WAL independently and are excluded from this check. +async fn flush_and_snapshot_region_wal(engine: &MitoEngine) -> WalSnapshot { + let mut snapshot = BTreeMap::new(); + for region in engine.regions() { + let id = region.region_id(); + if id.table_id() < MIN_USER_TABLE_ID || id.region_group() == METRIC_METADATA_REGION_GROUP { + continue; + } + engine + .handle_request(id, RegionRequest::Flush(RegionFlushRequest::default())) + .await + .unwrap(); + let statistic = engine.region_statistic(id).unwrap(); + snapshot.insert( + (0, id), + ( + region.flushed_sequence(), + statistic.manifest.data_flushed_entry_id(), + ), + ); + } + snapshot +} + +/// Require a real write, then check WAL policy independently of row deduplication. +pub fn assert_wal_delta(before: &WalSnapshot, after: &WalSnapshot, skip_wal: bool) { + assert_eq!( + before.keys().collect::>(), + after.keys().collect::>(), + "warm up table creation before taking the snapshot" + ); + let mut written = 0; + for (id, &(flushed_sequence, flushed_entry_id)) in after { + let (previous_flushed_sequence, previous_flushed_entry_id) = before[id]; + if flushed_sequence > previous_flushed_sequence { + written += 1; + if skip_wal { + assert_eq!( + flushed_entry_id, previous_flushed_entry_id, + "region {id:?} wrote WAL despite the request policy" + ); + } else { + assert!( + flushed_entry_id > previous_flushed_entry_id, + "region {id:?} did not write WAL" + ); + } + } + } + assert!( + written > 0, + "request must reach a data region, not merely return success" + ); +} + +/// A test instance backed by embedded storage or a distributed cluster. +pub enum MockInstanceImpl { + Standalone(GreptimeDbStandalone), + Distributed(GreptimeDbCluster), +} + +impl MockInstanceImpl { + /// Returns the metasrv of a distributed instance. + /// + /// # Panics + /// Panics if this is a standalone instance. + #[cfg(test)] + pub(crate) fn metasrv(&self) -> &Arc { + match self { + Self::Standalone(_) => unreachable!(), + Self::Distributed(instance) => &instance.metasrv, + } + } + + /// Returns the datanodes of a distributed instance. + /// + /// # Panics + /// Panics if this is a standalone instance. + #[cfg(test)] + pub(crate) fn datanodes(&self) -> &HashMap { + match self { + Self::Standalone(_) => unreachable!(), + Self::Distributed(instance) => &instance.datanode_instances, + } + } + + /// Creates a standalone instance or a three-datanode cluster using local storage. + pub async fn new(name: &str, distributed: bool) -> Self { + let name = format!( + "{name}_{}", + if distributed { + "distributed" + } else { + "standalone" + } + ); + if distributed { + // The repository cluster harness uses real tonic/protobuf services + // over duplex transports between FE and datanodes, not direct calls. + Self::Distributed( + GreptimeDbClusterBuilder::new(&name) + .await + .with_datanodes(3) + .build(false) + .await, + ) + } else { + Self::Standalone(GreptimeDbStandaloneBuilder::new(&name).build().await) + } + } + + /// Returns the frontend instance. + pub fn frontend(&self) -> Arc { + match self { + Self::Standalone(instance) => instance.fe_instance().clone(), + Self::Distributed(cluster) => cluster.fe_instance().clone(), + } + } + + /// Flushes user data regions and returns their write and WAL watermarks. + pub async fn flush_and_snapshot_wal(&self) -> WalSnapshot { + match self { + Self::Standalone(instance) => { + flush_and_snapshot_region_wal(&instance.mito_engine).await + } + Self::Distributed(cluster) => { + let mut result = BTreeMap::new(); + for (node_id, datanode) in &cluster.datanode_instances { + let engine = datanode.region_server().mito_engine().unwrap(); + result.extend( + flush_and_snapshot_region_wal(&engine) + .await + .into_iter() + .map(|((_, region_id), watermarks)| { + ((*node_id, region_id), watermarks) + }), + ); + } + result + } + } + } + + /// Shuts down cluster services and removes the test storage. + pub async fn shutdown(&mut self) { + match self { + Self::Standalone(instance) => instance.guard.remove_all().await, + Self::Distributed(cluster) => { + cluster.metasrv.shutdown().await.unwrap(); + for datanode in cluster.datanode_instances.values_mut() { + datanode.shutdown().await.unwrap(); + } + for guard in &mut cluster.guards { + guard.remove_all().await; + } + } + } + } +} + pub const PEER_PLACEHOLDER_ADDR: &str = "127.0.0.1:3001"; #[derive(Debug, Clone, Copy, Eq, PartialEq)] @@ -663,11 +843,27 @@ async fn setup_test_prom_app_with_frontend_inner( let sql = "INSERT INTO mito(host, val, ts) VALUES (1, 1.1, 0)"; run_sql(sql, &instance).await; + let http_server = build_test_prom_server( + instance.fe_instance().clone(), + enable_batcher, + experimental_enable_prometheus_native_histogram, + ) + .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap()) + .build(); + let app = http_server.build(http_server.make_app()).unwrap(); + (app, instance.guard) +} + +/// Builds Prometheus HTTP routes for either a standalone or distributed frontend. +pub fn build_test_prom_server( + frontend_ref: Arc, + enable_batcher: bool, + experimental_enable_prometheus_native_histogram: bool, +) -> HttpServerBuilder { let http_opts = HttpOptions { addr: format!("127.0.0.1:{}", ports::get_port()), ..Default::default() }; - let frontend_ref = instance.fe_instance().clone(); // Mirror the production wiring at `frontend::server`: build the batcher from the // instance's managers. A short flush interval keeps the test responsive. let pending_rows_batcher = if enable_batcher { @@ -688,9 +884,10 @@ async fn setup_test_prom_app_with_frontend_inner( } else { None }; - let http_server = HttpServerBuilder::new(http_opts) + assert_eq!(pending_rows_batcher.is_some(), enable_batcher); + HttpServerBuilder::new(http_opts) .with_sql_handler(frontend_ref.clone()) - .with_logs_handler(instance.fe_instance().clone()) + .with_logs_handler(frontend_ref.clone()) .with_prom_handler( frontend_ref.clone(), Some(frontend_ref.clone()), @@ -700,10 +897,6 @@ async fn setup_test_prom_app_with_frontend_inner( pending_rows_batcher, ) .with_prometheus_handler(frontend_ref) - .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap()) - .build(); - let app = http_server.build(http_server.make_app()).unwrap(); - (app, instance.guard) } pub async fn setup_grpc_server( diff --git a/tests-integration/src/tests/gc.rs b/tests-integration/src/tests/gc.rs index badaef5f7ac..20242223422 100644 --- a/tests-integration/src/tests/gc.rs +++ b/tests-integration/src/tests/gc.rs @@ -208,6 +208,7 @@ async fn test_gc_basic(store_type: &StorageType) { let procedure = BatchGcProcedure::new( metasrv.mailbox().clone(), metasrv.table_metadata_manager().clone(), + metasrv.runtime_switch_manager().clone(), metasrv.options().grpc.server_addr.clone(), regions.clone(), false, // full_file_listing diff --git a/tests-integration/src/tests/gc/repart.rs b/tests-integration/src/tests/gc/repart.rs index 7e7a3edd5a5..93969dbb3b3 100644 --- a/tests-integration/src/tests/gc/repart.rs +++ b/tests-integration/src/tests/gc/repart.rs @@ -90,6 +90,7 @@ CREATE TABLE test_cleanup_repartition ( let mut procedure = BatchGcProcedure::new_update_repartition_for_test( metasrv.mailbox().clone(), metasrv.table_metadata_manager().clone(), + metasrv.runtime_switch_manager().clone(), metasrv.options().grpc.server_addr.clone(), regions, manifest, @@ -183,6 +184,7 @@ async fn test_cleanup_region_repartition_preserve_uninvolved_entries() { let mut procedure = BatchGcProcedure::new_update_repartition_for_test( metasrv.mailbox().clone(), metasrv.table_metadata_manager().clone(), + metasrv.runtime_switch_manager().clone(), metasrv.options().grpc.server_addr.clone(), regions, manifest, @@ -267,6 +269,7 @@ async fn test_cleanup_region_repartition_remove_when_tmp_refs_empty() { let mut procedure = BatchGcProcedure::new_update_repartition_for_test( metasrv.mailbox().clone(), metasrv.table_metadata_manager().clone(), + metasrv.runtime_switch_manager().clone(), metasrv.options().grpc.server_addr.clone(), regions, manifest, diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index a716284c758..3cd625ba204 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -1174,6 +1174,55 @@ async fn test_execute_query_external_table_csv(instance: Arc) check_output_stream(output, expect).await; } +#[tokio::test(flavor = "multi_thread")] +async fn test_copy_from_inherits_skip_wal() { + use crate::test_util::{MockInstanceImpl, assert_wal_delta}; + + // Local COPY access is intentionally restricted to the standalone harness. + let mut env = MockInstanceImpl::new("copy_from_skip_wal", false).await; + let instance = env.frontend(); + let directory = create_local_file_test_dir("copy_from_skip_wal"); + let csv_path = directory.path().join("copy_skip_wal.csv"); + std::fs::copy( + find_testing_resource("/tests/data/csv/headerless.csv"), + &csv_path, + ) + .unwrap(); + execute_sql( + &instance, + "CREATE TABLE copy_skip_wal(host_id INT, host_name STRING, reading_value DOUBLE, ts TIMESTAMP TIME INDEX);", + ) + .await; + + let commands = [ + format!( + "COPY copy_skip_wal FROM '{}' WITH (FORMAT='csv', HEADERS='false');", + prepare_path(&csv_path.display().to_string()) + ), + format!( + "COPY DATABASE public FROM '{}/' WITH (FORMAT='csv', HEADERS='false');", + prepare_path(&directory.path().display().to_string()) + ), + ]; + for command in commands { + for skip_wal in [Some(true), Some(false), None] { + let ctx = QueryContext::arc(); + if let Some(skip_wal) = skip_wal { + ctx.set_skip_wal(skip_wal); + } + let before = env.flush_and_snapshot_wal().await; + let output = execute_sql_with(&instance, &command, ctx).await; + assert!(matches!(output.data, OutputData::AffectedRows(2))); + assert_wal_delta( + &before, + &env.flush_and_snapshot_wal().await, + skip_wal == Some(true), + ); + } + } + env.shutdown().await; +} + #[apply(standalone_instance_case)] async fn test_execute_copy_from_headerless_csv(instance: Arc) { let instance = instance.frontend(); diff --git a/tests-integration/src/tests/test_util.rs b/tests-integration/src/tests/test_util.rs index ab860367509..5ddcd0748a9 100644 --- a/tests-integration/src/tests/test_util.rs +++ b/tests-integration/src/tests/test_util.rs @@ -42,6 +42,7 @@ use session::context::{QueryContext, QueryContextRef}; use crate::cluster::{GreptimeDbCluster, GreptimeDbClusterBuilder}; use crate::standalone::{GreptimeDbStandalone, GreptimeDbStandaloneBuilder}; +pub(crate) use crate::test_util::MockInstanceImpl; use crate::test_util::StorageType; use crate::tests::{MockDistributedInstance, create_distributed_instance}; @@ -85,33 +86,9 @@ pub(crate) enum MockInstanceBuilder { Distributed(GreptimeDbClusterBuilder), } -pub(crate) enum MockInstanceImpl { - Standalone(GreptimeDbStandalone), - Distributed(GreptimeDbCluster), -} - -impl MockInstanceImpl { - pub(crate) fn metasrv(&self) -> &Arc { - match self { - MockInstanceImpl::Standalone(_) => unreachable!(), - MockInstanceImpl::Distributed(instance) => &instance.metasrv, - } - } - - pub(crate) fn datanodes(&self) -> &HashMap { - match self { - MockInstanceImpl::Standalone(_) => unreachable!(), - MockInstanceImpl::Distributed(instance) => &instance.datanode_instances, - } - } -} - impl MockInstance for MockInstanceImpl { fn frontend(&self) -> Arc { - match self { - MockInstanceImpl::Standalone(instance) => instance.frontend(), - MockInstanceImpl::Distributed(instance) => instance.fe_instance().clone(), - } + MockInstanceImpl::frontend(self) } fn is_distributed_mode(&self) -> bool { diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index 242161654bd..9607d611627 100644 --- a/tests-integration/tests/grpc.rs +++ b/tests-integration/tests/grpc.rs @@ -15,13 +15,16 @@ use std::sync::Arc; use api::v1::alter_table_expr::Kind; +use api::v1::greptime_database_client::GreptimeDatabaseClient; +use api::v1::greptime_request::Request as RequestBody; +use api::v1::greptime_response::Response as ResponseBody; use api::v1::promql_request::Promql; use api::v1::value::ValueData; use api::v1::{ - AddColumn, AddColumns, AlterTableExpr, Basic, Column, ColumnDataType, ColumnDef, - CreateTableExpr, InsertRequest, InsertRequests, PromInstantQuery, PromRangeQuery, - PromqlRequest, RequestHeader, Row, RowInsertRequest, RowInsertRequests, SemanticType, Value, - column, + AddColumn, AddColumns, AlterTableExpr, Basic, Column, ColumnDataType, ColumnDef, ColumnSchema, + CreateTableExpr, GreptimeRequest, InsertRequest, InsertRequests, PromInstantQuery, + PromRangeQuery, PromqlRequest, RequestHeader, Row, RowInsertRequest, RowInsertRequests, Rows, + SemanticType, Value, column, }; use auth::user_provider_from_option; use base64::prelude::{BASE64_STANDARD, Engine as _}; @@ -52,22 +55,27 @@ use otel_arrow_rust::proto::opentelemetry::arrow::v1::{ }; use otel_arrow_rust::proto::opentelemetry::metrics::v1::AggregationTemporality; use otel_arrow_rust::schema::consts as arrow_consts; +use rstest_reuse::apply; use servers::grpc::GrpcServerConfig; use servers::grpc::builder::GrpcServerBuilder; use servers::http::prometheus::{ - PromData, PromQueryResult, PromSeriesMatrix, PromSeriesVector, PrometheusJsonResponse, - PrometheusResponse, + PromData, PromQueryResult, PromSampleValue, PromSeriesMatrix, PromSeriesVector, + PrometheusJsonResponse, PrometheusResponse, }; use servers::request_memory_limiter::ServerMemoryLimiter; use servers::server::Server; use servers::tls::{TlsMode, TlsOption}; +use session::hints::{HINTS_KEY, INSERT_SKIP_WAL_HINT}; use tests_integration::test_util::{ - StorageType, setup_grpc_server, setup_grpc_server_with, + MockInstanceImpl, StorageType, assert_wal_delta, setup_grpc_server, + setup_grpc_server_for_frontend_instance, setup_grpc_server_with, setup_grpc_server_with_auto_create_table_disabled, setup_grpc_server_with_user_provider, }; use tonic::Request; use tonic::metadata::MetadataValue; +use crate::both_deployment_cases; + #[macro_export] macro_rules! grpc_test { ($service:ident, $($(#[$meta:meta])* $test:ident),*,) => { @@ -776,6 +784,44 @@ fn gauge_arrow_batch(batch_id: i64, reserved_attr: bool) -> BatchArrowRecords { } } +#[apply(both_deployment_cases)] +async fn test_skip_wal_otel_arrow_metrics(distributed: bool) { + // OTEL Arrow metrics decode into ordinary metric inserts, unlike Flight DoPut bulk inserts. + let mut env = MockInstanceImpl::new("skip_wal_otel_arrow", distributed).await; + let server = setup_grpc_server_for_frontend_instance(env.frontend(), None).await; + let mut client = + ArrowMetricsServiceClient::connect(format!("http://{}", server.bind_addr().unwrap())) + .await + .unwrap(); + // Warm up auto-created logical/physical tables and metadata before comparing data-region WAL. + for (batch_id, hint) in [None, Some("true"), Some("false"), None] + .into_iter() + .enumerate() + { + let before = env.flush_and_snapshot_wal().await; + let batch = gauge_arrow_batch(batch_id as i64, false); + let request = with_skip_wal_hint(futures::stream::iter([batch]), hint); + let mut response = client.arrow_metrics(request).await.unwrap().into_inner(); + let status = response.message().await.unwrap().unwrap(); + assert_eq!( + status.status_code, + ArrowStatusCode::Ok as i32, + "{}", + status.status_message + ); + assert!(response.message().await.unwrap().is_none()); + if batch_id != 0 { + assert_wal_delta( + &before, + &env.flush_and_snapshot_wal().await, + hint == Some("true"), + ); + } + } + server.shutdown().await.unwrap(); + env.shutdown().await; +} + pub async fn test_otel_arrow_delta_histogram(store_type: StorageType) { let (_instance, server) = setup_grpc_server(store_type, "test_otel_arrow_delta_histogram").await; @@ -1152,6 +1198,139 @@ fn expect_data() -> (Column, Column, Column, Column) { ) } +fn with_skip_wal_hint(body: T, hint: Option<&str>) -> Request { + let mut request = Request::new(body); + if let Some(hint) = hint { + request.metadata_mut().insert( + HINTS_KEY, + format!("{INSERT_SKIP_WAL_HINT}={hint}").parse().unwrap(), + ); + } + request +} + +fn skip_wal_insert_body(columnar: bool) -> RequestBody { + if columnar { + RequestBody::Inserts(InsertRequests { + inserts: vec![InsertRequest { + table_name: "skip_wal_grpc".to_string(), + row_count: 1, + columns: vec![Column { + column_name: "ts".to_string(), + semantic_type: SemanticType::Timestamp as i32, + datatype: ColumnDataType::TimestampMillisecond as i32, + values: Some(column::Values { + timestamp_millisecond_values: vec![1000], + ..Default::default() + }), + ..Default::default() + }], + }], + }) + } else { + RequestBody::RowInserts(RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: "skip_wal_grpc".to_string(), + rows: Some(Rows { + schema: vec![ColumnSchema { + column_name: "ts".to_string(), + semantic_type: SemanticType::Timestamp as i32, + datatype: ColumnDataType::TimestampMillisecond as i32, + ..Default::default() + }], + rows: vec![Row { + values: vec![Value { + value_data: Some(ValueData::TimestampMillisecondValue(1000)), + }], + }], + }), + }], + }) + } +} + +#[apply(both_deployment_cases)] +async fn test_skip_wal_grpc_unary_stream_and_flight_sql(distributed: bool) { + let mut env = MockInstanceImpl::new("skip_wal_grpc_protocols", distributed).await; + let server = setup_grpc_server_for_frontend_instance(env.frontend(), None).await; + let addr = server.bind_addr().unwrap().to_string(); + let mut grpc = GreptimeDatabaseClient::connect(format!("http://{addr}")) + .await + .unwrap(); + let database = Database::new_with_dbname("greptime-public", Client::with_urls(vec![addr])); + database + .sql("CREATE TABLE skip_wal_grpc (ts TIMESTAMP TIME INDEX)") + .await + .unwrap(); + + for streaming in [false, true] { + for columnar in [false, true] { + for hint in [Some("true"), Some("false"), None] { + let before = env.flush_and_snapshot_wal().await; + let request = GreptimeRequest { + header: Some(RequestHeader { + catalog: "greptime".to_string(), + schema: "public".to_string(), + ..Default::default() + }), + request: Some(skip_wal_insert_body(columnar)), + }; + let response = if streaming { + grpc.handle_requests(with_skip_wal_hint( + futures::stream::iter([request.clone(), request]), + hint, + )) + .await + .unwrap() + } else { + grpc.handle(with_skip_wal_hint(request, hint)) + .await + .unwrap() + } + .into_inner(); + let expected_rows = if streaming { 2 } else { 1 }; + assert!( + matches!(response.response, Some(ResponseBody::AffectedRows(rows)) if rows.value == expected_rows) + ); + assert_wal_delta( + &before, + &env.flush_and_snapshot_wal().await, + hint == Some("true"), + ); + } + } + } + for hint in [Some("true"), Some("false"), None] { + let before = env.flush_and_snapshot_wal().await; + let hints = hint + .map(|value| vec![(INSERT_SKIP_WAL_HINT, value)]) + .unwrap_or_default(); + database + .sql_with_hint("INSERT INTO skip_wal_grpc VALUES (1000)", &hints) + .await + .unwrap(); + assert_wal_delta( + &before, + &env.flush_and_snapshot_wal().await, + hint == Some("true"), + ); + } + // Strict validation must reject the request before any data write. + let before = env.flush_and_snapshot_wal().await; + assert!( + database + .sql_with_hint( + "INSERT INTO skip_wal_grpc VALUES (1000)", + &[(INSERT_SKIP_WAL_HINT, "yes")] + ) + .await + .is_err() + ); + assert_eq!(before, env.flush_and_snapshot_wal().await); + server.shutdown().await.unwrap(); + env.shutdown().await; +} + pub async fn test_insert_and_select(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (_db, fe_grpc_server) = setup_grpc_server(store_type, "test_insert_and_select").await; @@ -1538,7 +1717,7 @@ pub async fn test_prom_gateway_query(store_type: StorageType) { panic!("unexpected result type") }; - mat.sort_unstable_by_key(|v| v.values[0].1.clone()); + mat.sort_unstable_by_key(|v| serde_json::to_string(&v.values[0].1).unwrap()); assert_eq!( mat, @@ -1550,7 +1729,10 @@ pub async fn test_prom_gateway_query(store_type: StorageType) { ] .into_iter() .collect(), - values: vec![(5.0, "1.0".to_string()), (10.0, "1.0".to_string())], + values: vec![ + (5.0, PromSampleValue::Text("1.0".to_string())), + (10.0, PromSampleValue::Text("1.0".to_string())), + ], ..Default::default() }, PromSeriesMatrix { @@ -1560,7 +1742,10 @@ pub async fn test_prom_gateway_query(store_type: StorageType) { ] .into_iter() .collect(), - values: vec![(5.0, "2.0".to_string()), (10.0, "2.0".to_string())], + values: vec![ + (5.0, PromSampleValue::Text("2.0".to_string())), + (10.0, PromSampleValue::Text("2.0".to_string())), + ], ..Default::default() }, ] diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index c696bd40ac4..f53e44fcc21 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -58,8 +58,8 @@ use opentelemetry_proto::tonic::collector::trace::v1::{ }; use pipeline::GREPTIME_INTERNAL_TRACE_PIPELINE_V1_NAME; use prost::Message; +use rstest_reuse::apply; use serde_json::{Value, json}; -use servers::http::GreptimeQueryOutput; use servers::http::handler::HealthResponse; use servers::http::header::constants::{ GREPTIME_LOG_EXTRACT_KEYS_HEADER_NAME, GREPTIME_LOG_TABLE_NAME_HEADER_NAME, @@ -72,20 +72,23 @@ use servers::http::result::error_result::ErrorResponse; use servers::http::result::greptime_result_v1::GreptimedbV1Response; use servers::http::result::influxdb_result_v1::{InfluxdbOutput, InfluxdbV1Response}; use servers::http::test_helpers::{TestClient, TestResponse}; +use servers::http::{GreptimeQueryOutput, HttpOptions, HttpServerBuilder}; use servers::prom_remote_write::v2::test_util as remote_write_v2; +use servers::prom_remote_write::validation::PromValidationMode; use servers::prom_store::{self, mock_timeseries_new_label}; use servers::request_memory_limiter::ServerMemoryLimiter; use standalone::options::StandaloneOptions; use table::table_name::TableName; use tests_integration::test_util::{ - StorageType, setup_test_http_app, setup_test_http_app_with_frontend, - setup_test_http_app_with_frontend_and_slow_query_threshold, + MockInstanceImpl, StorageType, assert_wal_delta, build_test_prom_server, setup_test_http_app, + setup_test_http_app_with_frontend, setup_test_http_app_with_frontend_and_slow_query_threshold, setup_test_http_app_with_frontend_and_user_provider, setup_test_prom_app_with_frontend, setup_test_prom_app_with_frontend_batched, setup_test_prom_app_with_frontend_native_histogram, }; use urlencoding::encode; use yaml_rust::YamlLoader; +use crate::both_deployment_cases; use crate::event_recorder_test_util::assert_procedure_actor_by_table; #[macro_export] @@ -444,6 +447,33 @@ pub async fn test_cors() { guard.remove_all().await; } +#[apply(both_deployment_cases)] +async fn test_sql_skip_wal(distributed: bool) { + let mut cases = vec![HttpWalCase::new( + "SQL", + "/v1/sql", + "application/x-www-form-urlencoded", + "sql=INSERT+INTO+wal_sql+VALUES+(1%2C-1)%2C(2%2C1)", + )]; + let mut sql_get = HttpWalCase::new( + "SQL GET", + "/v1/sql?sql=INSERT%20INTO%20wal_sql%20VALUES%20(1%2C-1)%2C(2%2C1)", + "application/x-www-form-urlencoded", + Vec::new(), + ); + sql_get.get = true; + cases.push(sql_get); + for case in &mut cases { + case.create_table = Some(if distributed { + "CREATE TABLE IF NOT EXISTS wal_sql (ts TIMESTAMP TIME INDEX, v INT) PARTITION ON COLUMNS (v) (v < 0, v >= 0)" + } else { + "CREATE TABLE IF NOT EXISTS wal_sql (ts TIMESTAMP TIME INDEX, v INT)" + }); + case.expected_written_nodes = distributed.then_some(2); + } + check_http_skip_wal("sql", &cases, distributed).await; +} + pub async fn test_sql_api(store_type: StorageType) { let (app, mut guard) = setup_test_http_app_with_frontend(store_type, "sql_api").await; let client = TestClient::new(app).await; @@ -1628,6 +1658,25 @@ pub async fn test_splunk_health_is_public(store_type: StorageType) { assert_eq!(StatusCode::OK, res.status()); } +#[apply(both_deployment_cases)] +async fn test_splunk_skip_wal(distributed: bool) { + let cases = vec![ + HttpWalCase::new( + "Splunk event", + "/v1/splunk/services/collector/event", + "application/json", + r#"{"event":"wal test","time":1700000000,"index":"wal_splunk"}"#, + ), + HttpWalCase::new( + "Splunk raw", + "/v1/splunk/services/collector/raw?index=wal_splunk_raw&time=1700000000", + "text/plain", + "wal test", + ), + ]; + check_http_skip_wal("splunk", &cases, distributed).await; +} + pub async fn test_splunk_logs(store_type: StorageType) { common_telemetry::init_default_ut_logging(); @@ -2332,7 +2381,7 @@ manifest_checkpoint_distance = 10 experimental_manifest_keep_removed_file_count = 256 experimental_manifest_keep_removed_file_ttl = "1h" compress_manifest = false -experimental_series_index_root = "" +experimental_enable_series_index = false experimental_series_index_maintenance_interval = "5m" experimental_compaction_memory_limit = "unlimited" experimental_compaction_on_exhausted = "wait" @@ -2649,6 +2698,43 @@ pub async fn test_dashboard_api(store_type: StorageType) { #[cfg(not(feature = "dashboard"))] pub async fn test_dashboard_api(_: StorageType) {} +#[apply(both_deployment_cases)] +async fn test_prometheus_skip_wal(distributed: bool) { + let mut cases = Vec::new(); + let prom = api::prom_store::remote::WriteRequest { + timeseries: prom_store::mock_timeseries(), + ..Default::default() + }; + let mut prom_case = HttpWalCase::new( + "Prometheus v1", + "/v1/prometheus/write", + "application/x-protobuf", + prom_store::snappy_compress(&prom.encode_to_vec()).unwrap(), + ); + prom_case.headers.push(("content-encoding", "snappy")); + cases.push(prom_case); + let prom_v2 = servers::prom_remote_write::v2::test_util::request_with_labels_and_samples( + vec![ + (prom_store::METRIC_NAME_LABEL, "wal_prom_v2"), + ("host", "a"), + ], + vec![api::greptime_proto::io::prometheus::write::v2::Sample { + value: 1.0, + timestamp: 1700000000000, + start_timestamp: 0, + }], + ); + let mut prom_v2_case = HttpWalCase::new( + "Prometheus v2", + "/v1/prometheus/write", + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + prom_store::snappy_compress(&prom_v2.encode_to_vec()).unwrap(), + ); + prom_v2_case.headers.push(("content-encoding", "snappy")); + cases.push(prom_v2_case); + check_http_skip_wal("prometheus", &cases, distributed).await; +} + pub async fn test_prometheus_remote_write(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = @@ -3224,6 +3310,94 @@ pub async fn test_prometheus_remote_write_v2_native_histogram(store_type: Storag guard.remove_all().await; } +#[apply(both_deployment_cases)] +async fn test_prometheus_remote_write_batched_skip_wal(distributed: bool) { + check_prometheus_remote_write_batched_skip_wal(distributed, false).await; +} + +#[apply(both_deployment_cases)] +async fn test_prometheus_remote_write_v2_batched_skip_wal(distributed: bool) { + check_prometheus_remote_write_batched_skip_wal(distributed, true).await; +} + +async fn check_prometheus_remote_write_batched_skip_wal(distributed: bool, v2: bool) { + common_telemetry::init_default_ut_logging(); + let mut instance = + MockInstanceImpl::new(&format!("prom_bulk_skip_wal_v2_{v2}"), distributed).await; + let server = build_test_prom_server(instance.frontend(), true, false).build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + + write_prometheus_skip_wal_sample(&client, v2, 1000, None).await; + wait_for_data(&client, "SELECT count(*) FROM wal_prom_bulk", "[[1]]").await; + let mut before = instance.flush_and_snapshot_wal().await; + // Reuse the same batcher across policy transitions to detect leaked policy. + for (index, hint) in [None, Some("true"), Some("false"), None] + .into_iter() + .enumerate() + { + write_prometheus_skip_wal_sample(&client, v2, (index as i64 + 2) * 1000, hint).await; + wait_for_data( + &client, + "SELECT count(*) FROM wal_prom_bulk", + &format!("[[{}]]", index + 2), + ) + .await; + let after = instance.flush_and_snapshot_wal().await; + assert_wal_delta(&before, &after, hint == Some("true")); + before = after; + } + instance.shutdown().await; +} + +async fn write_prometheus_skip_wal_sample( + client: &TestClient, + v2: bool, + timestamp: i64, + hint: Option<&str>, +) { + let payload = if v2 { + remote_write_v2::request_with_labels_and_samples( + vec![(prom_store::METRIC_NAME_LABEL, "wal_prom_bulk")], + vec![RemoteWriteV2Sample { + value: 1.0, + timestamp, + start_timestamp: 0, + }], + ) + .encode_to_vec() + } else { + WriteRequest { + timeseries: vec![TimeSeries { + labels: vec![Label { + name: prom_store::METRIC_NAME_LABEL.to_string(), + value: "wal_prom_bulk".to_string(), + }], + samples: vec![Sample { + value: 1.0, + timestamp, + }], + ..Default::default() + }], + ..Default::default() + } + .encode_to_vec() + }; + let mut request = client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .body(prom_store::snappy_compress(&payload).unwrap()); + if v2 { + request = request.header( + "Content-Type", + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + ); + } + if let Some(hint) = hint { + request = request.header("x-greptime-insert-skip-wal", hint); + } + assert_eq!(request.send().await.status(), StatusCode::NO_CONTENT); +} + /// Covers the batched (pending-rows-batcher) Prometheus remote write path, which /// bypasses `PromStoreProtocolHandler::write`. Verifies the metric table is created /// asynchronously and still carries the Prometheus semantic identity stamped on the @@ -4135,6 +4309,25 @@ transform: guard.remove_all().await; } +#[apply(both_deployment_cases)] +async fn test_json_logs_skip_wal(distributed: bool) { + let cases = vec![ + HttpWalCase::new( + "JSON logs", + "/v1/ingest?table=wal_logs&pipeline_name=greptime_identity", + "application/json", + r#"[{"message":"wal test"}]"#, + ), + HttpWalCase::new( + "Events logs", + "/v1/events/logs?table=wal_events&pipeline_name=greptime_identity", + "application/json", + r#"[{"message":"wal test"}]"#, + ), + ]; + check_http_skip_wal("json_logs", &cases, distributed).await; +} + pub async fn test_identity_pipeline(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = @@ -5343,6 +5536,25 @@ transform: guard.remove_all().await; } +#[apply(both_deployment_cases)] +async fn test_influxdb_skip_wal(distributed: bool) { + let cases = vec![ + HttpWalCase::new( + "InfluxDB v1", + "/v1/influxdb/write?db=public", + "text/plain", + "wal_influx,host=a value=1 1700000000000000000", + ), + HttpWalCase::new( + "InfluxDB v2", + "/v1/influxdb/api/v2/write?bucket=public", + "text/plain", + "wal_influx_v2,host=a value=1 1700000000000000000", + ), + ]; + check_http_skip_wal("influxdb", &cases, distributed).await; +} + pub async fn test_influxdb_write_with_hints(storage_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = @@ -6523,6 +6735,21 @@ pub async fn test_pipeline_auto_transform_with_select(store_type: StorageType) { guard.remove_all().await; } +#[apply(both_deployment_cases)] +async fn test_otlp_metrics_skip_wal(distributed: bool) { + let mut cases = Vec::new(); + let metrics: ExportMetricsServiceRequest = serde_json::from_value(json!({ + "resourceMetrics":[{"scopeMetrics":[{"metrics":[{"name":"wal_otlp", "gauge":{"dataPoints":[{"timeUnixNano":"1700000000000000000","asDouble":1.0}]}}]}]}] + })).unwrap(); + cases.push(HttpWalCase::new( + "OTLP metrics", + "/v1/otlp/v1/metrics", + "application/x-protobuf", + metrics.encode_to_vec(), + )); + check_http_skip_wal("otlp_metrics", &cases, distributed).await; +} + pub async fn test_otlp_metrics_new(store_type: StorageType) { // init common_telemetry::init_default_ut_logging(); @@ -7318,6 +7545,42 @@ pub async fn test_otlp_metrics_resource_info_conflicts(store_type: StorageType) guard.remove_all().await; } +#[apply(both_deployment_cases)] +async fn test_otlp_traces_skip_wal(distributed: bool) { + let mut cases = Vec::new(); + let traces = make_trace_v1_request( + "wal-test", + vec![json!({ + "traceId":"0102030405060708090a0b0c0d0e0f10", + "spanId":"0102030405060708", "name":"wal test", + "startTimeUnixNano":"1700000000000000000", + "endTimeUnixNano":"1700000001000000000" + })], + ); + let mut traces_case = HttpWalCase::new( + "OTLP traces", + "/v1/otlp/v1/traces", + "application/x-protobuf", + traces.encode_to_vec(), + ); + traces_case + .headers + .push(("x-greptime-pipeline-name", "greptime_trace_v0")); + cases.push(traces_case); + let mut traces_v1 = HttpWalCase::new( + "OTLP traces v1", + "/v1/otlp/v1/traces", + "application/x-protobuf", + traces.encode_to_vec(), + ); + traces_v1.headers = vec![ + ("x-greptime-pipeline-name", "greptime_trace_v1"), + ("x-greptime-trace-table-name", "wal_traces_v1"), + ]; + cases.push(traces_v1); + check_http_skip_wal("otlp_traces", &cases, distributed).await; +} + pub async fn test_otlp_traces_v0(store_type: StorageType) { // init common_telemetry::init_default_ut_logging(); @@ -8388,6 +8651,22 @@ pub async fn test_otlp_traces_v1(store_type: StorageType) { guard.remove_all().await; } +#[apply(both_deployment_cases)] +async fn test_otlp_logs_skip_wal(distributed: bool) { + let mut cases = Vec::new(); + let logs = make_log_request(vec![json!({ + "timeUnixNano": "1700000000000000000", + "body": {"stringValue": "wal test"} + })]); + cases.push(HttpWalCase::new( + "OTLP logs", + "/v1/otlp/v1/logs", + "application/x-protobuf", + logs.encode_to_vec(), + )); + check_http_skip_wal("otlp_logs", &cases, distributed).await; +} + pub async fn test_otlp_logs(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = setup_test_http_app_with_frontend(store_type, "test_otlp_logs").await; @@ -8678,6 +8957,37 @@ pub async fn test_otlp_logs(store_type: StorageType) { guard.remove_all().await; } +#[apply(both_deployment_cases)] +async fn test_loki_skip_wal(distributed: bool) { + let mut cases = vec![HttpWalCase::new( + "Loki JSON", + "/v1/loki/api/v1/push", + "application/json", + r#"{"streams":[{"stream":{"host":"a"},"values":[["1700000000000000000","wal test"]]}]}"#, + )]; + let loki = loki_proto::logproto::PushRequest { + streams: vec![loki_proto::logproto::StreamAdapter { + labels: r#"{host="a"}"#.to_string(), + entries: vec![loki_proto::logproto::EntryAdapter { + timestamp: Some(loki_proto::prost_types::Timestamp { + seconds: 1700000000, + nanos: 0, + }), + line: "wal test".to_string(), + ..Default::default() + }], + ..Default::default() + }], + }; + cases.push(HttpWalCase::new( + "Loki protobuf", + "/v1/loki/api/v1/push", + "application/x-protobuf", + prom_store::snappy_compress(&loki.encode_to_vec()).unwrap(), + )); + check_http_skip_wal("loki", &cases, distributed).await; +} + pub async fn test_loki_pb_logs(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = setup_test_http_app_with_frontend(store_type, "test_loki_pb_logs").await; @@ -9086,6 +9396,17 @@ processors: guard.remove_all().await; } +#[apply(both_deployment_cases)] +async fn test_elasticsearch_skip_wal(distributed: bool) { + let cases = vec![HttpWalCase::new( + "Elasticsearch", + "/v1/elasticsearch/_bulk", + "application/json", + "{\"create\":{\"_index\":\"wal_elastic\"}}\n{\"message\":\"wal test\"}\n", + )]; + check_http_skip_wal("elasticsearch", &cases, distributed).await; +} + pub async fn test_elasticsearch_logs(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = @@ -10789,6 +11110,17 @@ pub async fn test_jaeger_query_api_for_trace_v1(store_type: StorageType) { guard.remove_all().await; } +#[apply(both_deployment_cases)] +async fn test_opentsdb_skip_wal(distributed: bool) { + let cases = vec![HttpWalCase::new( + "OpenTSDB", + "/v1/opentsdb/api/put", + "application/json", + r#"{"metric":"wal_tsdb","timestamp":1700000000,"value":1,"tags":{"host":"a"}}"#, + )]; + check_http_skip_wal("opentsdb", &cases, distributed).await; +} + pub async fn test_influxdb_write(store_type: StorageType) { common_telemetry::init_default_ut_logging(); let (app, mut guard) = @@ -10897,6 +11229,123 @@ async fn execute_sql(client: &TestClient, sql: &str) -> TestResponse { .await } +struct HttpWalCase { + name: &'static str, + get: bool, + create_table: Option<&'static str>, + expected_written_nodes: Option, + path: &'static str, + content_type: &'static str, + body: Vec, + headers: Vec<(&'static str, &'static str)>, +} + +impl HttpWalCase { + fn new( + name: &'static str, + path: &'static str, + content_type: &'static str, + body: impl Into>, + ) -> Self { + Self { + name, + get: false, + create_table: None, + expected_written_nodes: None, + path, + content_type, + body: body.into(), + headers: Vec::new(), + } + } +} + +async fn check_http_skip_wal(name: &str, cases: &[HttpWalCase], distributed: bool) { + common_telemetry::init_default_ut_logging(); + let mut instance = MockInstanceImpl::new(&format!("http_skip_wal_{name}"), distributed).await; + let fe = instance.frontend(); + let server = HttpServerBuilder::new(HttpOptions::default()) + .with_sql_handler(fe.clone()) + .with_influxdb_handler(fe.clone()) + .with_opentsdb_handler(fe.clone()) + .with_log_ingest_handler(fe.clone(), None, None) + .with_otlp_handler(fe.clone(), true, false) + // The pending batcher uses BulkInsert, deliberately outside this PR. + .with_prom_handler( + fe.clone(), + Some(fe), + true, + PromValidationMode::Strict, + false, + None, + ) + .build(); + let client = TestClient::new(server.build(server.make_app()).unwrap()).await; + for case in cases { + if let Some(create_table) = case.create_table { + let response = execute_sql(&client, create_table).await; + assert!(response.status().is_success(), "{}", response.text().await); + } + // Warm up schema-on-write, then change only the request hint. Reusing + // the same payload also verifies policy is not persisted on the table. + for (round, hint) in [None, Some("true"), Some("false"), None] + .into_iter() + .enumerate() + { + common_telemetry::info!("Protocol WAL case: {}, hint: {:?}", case.name, hint); + let before = instance.flush_and_snapshot_wal().await; + let mut headers = vec![( + HeaderName::from_static("content-type"), + HeaderValue::from_static(case.content_type), + )]; + headers.extend(case.headers.iter().map(|(key, value)| { + ( + HeaderName::from_static(key), + HeaderValue::from_static(value), + ) + })); + if let Some(value) = hint { + headers.push(( + HeaderName::from_static("x-greptime-insert-skip-wal"), + HeaderValue::from_static(value), + )); + } + let response = if case.get { + let mut request = client.get(case.path).body(case.body.clone()); + for (key, value) in headers { + request = request.header(key, value); + } + request.send().await + } else { + send_req(&client, headers, case.path, case.body.clone(), false).await + }; + let status = response.status(); + let body = response.text().await; + assert!(status.is_success(), "{}: {status}: {body}", case.name); + let after = instance.flush_and_snapshot_wal().await; + if before.keys().eq(after.keys()) { + assert_wal_delta(&before, &after, hint == Some("true")); + if let Some(expected_written_nodes) = case.expected_written_nodes { + let written_nodes = after + .iter() + .filter(|(key, (flushed_sequence, _))| *flushed_sequence > before[*key].0) + .map(|((node_id, _), _)| *node_id) + .collect::>(); + assert_eq!( + written_nodes.len(), + expected_written_nodes, + "{}: unexpected number of written datanodes", + case.name + ); + } + } else { + assert!(round == 0, "only warmup may create regions: {}", case.name); + } + } + } + instance.shutdown().await; +} + async fn send_req( client: &TestClient, headers: Vec<(HeaderName, HeaderValue)>, diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index 56108d4d3ff..28a81155724 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -41,6 +41,13 @@ mod reconciliation_event; mod view_ddl_event; mod wal_prune_event; +#[rstest_reuse::template] +#[rstest::rstest] +#[case::standalone(false)] +#[case::distributed(true)] +#[tokio::test(flavor = "multi_thread")] +fn both_deployment_cases(#[case] distributed: bool) {} + grpc_tests!(File, S3, S3WithCache, Oss, Azblob, Gcs); http_tests!(File, S3, S3WithCache, Oss, Azblob, Gcs); diff --git a/tests-integration/tests/reconciliation_event.rs b/tests-integration/tests/reconciliation_event.rs index ebd12f2af96..19eaea6ff2c 100644 --- a/tests-integration/tests/reconciliation_event.rs +++ b/tests-integration/tests/reconciliation_event.rs @@ -27,6 +27,216 @@ const EVENTS_TABLE: &str = "greptime_private.events"; const CATALOG: &str = "greptime"; const DATABASE: &str = "reconciliation_event_database"; +#[tokio::test(flavor = "multi_thread")] +async fn test_catalog_and_database_reconciliation_events() { + common_telemetry::init_default_ut_logging(); + + let cluster = GreptimeDbClusterBuilder::new("catalog_database_reconciliation_events") + .await + .with_datanodes(1) + .build(true) + .await; + let frontend = cluster.fe_instance().clone(); + run_sql(frontend.as_ref(), &format!("CREATE DATABASE {DATABASE}")).await; + run_sql( + frontend.as_ref(), + &format!("CREATE TABLE {DATABASE}.metrics (ts TIMESTAMP TIME INDEX)"), + ) + .await; + + let catalog_procedure_id = cluster + .metasrv + .reconciliation_manager() + .reconcile_catalog(CATALOG.to_string(), ResolveStrategy::UseLatest, 1) + .await + .unwrap(); + let mut procedure_watcher = cluster + .metasrv + .procedure_manager() + .procedure_watcher(catalog_procedure_id) + .unwrap(); + watcher::wait(&mut procedure_watcher).await.unwrap(); + let catalog_procedure_id = catalog_procedure_id.to_string(); + + assert_eventually_eq( + &frontend, + &format!( + "SELECT count(*) AS event_count, \ + json_get_int(payload, 'version') AS version, \ + json_get_string(payload, 'resolve_strategy') AS resolve_strategy, \ + json_get_bool(payload, 'fail_fast') AS fail_fast, \ + json_get_int(payload, 'parallelism') AS parallelism \ + FROM {EVENTS_TABLE} \ + WHERE type = 'reconcile_catalog' AND procedure_id = '{catalog_procedure_id}' \ + AND json_get_string(procedure_trigger, 'type') = 'Submitted' \ + AND catalog_name = '{CATALOG}' AND schema_name IS NULL AND table_name IS NULL \ + AND table_id IS NULL AND physical_table_id IS NULL \ + GROUP BY json_get_int(payload, 'version'), \ + json_get_string(payload, 'resolve_strategy'), \ + json_get_bool(payload, 'fail_fast'), \ + json_get_int(payload, 'parallelism')" + ), + "\ ++-------------+---------+------------------+-----------+-------------+ +| event_count | version | resolve_strategy | fail_fast | parallelism | ++-------------+---------+------------------+-----------+-------------+ +| 1 | 1 | use_latest | false | 1 | ++-------------+---------+------------------+-----------+-------------+", + ) + .await; + + let catalog_result_payload = find_eventually_string( + &frontend, + &format!( + "SELECT json_to_string(payload) AS payload FROM {EVENTS_TABLE} \ + WHERE type = 'reconcile_catalog' AND procedure_id = '{catalog_procedure_id}' \ + AND json_get_string(procedure_trigger, 'type') = 'Succeeded' \ + AND catalog_name = '{CATALOG}' AND schema_name IS NULL AND table_name IS NULL \ + AND table_id IS NULL AND physical_table_id IS NULL LIMIT 1" + ), + "payload", + ) + .await; + let catalog_result: serde_json::Value = serde_json::from_str(&catalog_result_payload).unwrap(); + assert_eq!(catalog_result["version"], serde_json::json!(1)); + assert_eq!(catalog_result["complete"], serde_json::json!(true)); + let processed_database_count = catalog_result["processed_database_count"].as_u64().unwrap(); + let succeeded_database_count = catalog_result["succeeded_database_count"].as_u64().unwrap(); + let failed_database_count = catalog_result["failed_database_count"].as_u64().unwrap(); + assert_eq!( + processed_database_count, + succeeded_database_count + failed_database_count + ); + assert!(processed_database_count > 0); + assert_eq!(failed_database_count, 0); + + let database_procedure_id = find_eventually_string( + &frontend, + &format!( + "SELECT procedure_id FROM {EVENTS_TABLE} \ + WHERE type = 'reconcile_database' AND catalog_name = '{CATALOG}' \ + AND schema_name = '{DATABASE}' \ + AND json_get_string(procedure_trigger, 'type') = 'Submitted' LIMIT 1" + ), + "procedure_id", + ) + .await; + + assert_eventually_eq( + &frontend, + &format!( + "SELECT parent.procedure_id AS parent_procedure_id, \ + child.procedure_id AS child_procedure_id, \ + parent.type AS parent_event_type, \ + child.type AS child_event_type, \ + parent.catalog_name AS parent_catalog_name, \ + parent.schema_name AS parent_schema_name, \ + child.catalog_name AS child_catalog_name, \ + child.schema_name AS child_schema_name \ + FROM {EVENTS_TABLE} AS parent \ + JOIN {EVENTS_TABLE} AS child \ + ON json_get_string(parent.procedure_trigger, 'procedure_id') = child.procedure_id \ + WHERE parent.procedure_id = '{catalog_procedure_id}' \ + AND json_get_string(parent.procedure_trigger, 'type') = 'ChildSubmitted' \ + AND child.procedure_id = '{database_procedure_id}' \ + AND json_get_string(child.procedure_trigger, 'type') = 'Submitted'" + ), + &format!( + "\ ++--------------------------------------+--------------------------------------+-------------------+--------------------+---------------------+--------------------+--------------------+-------------------------------+ +| parent_procedure_id | child_procedure_id | parent_event_type | child_event_type | parent_catalog_name | parent_schema_name | child_catalog_name | child_schema_name | ++--------------------------------------+--------------------------------------+-------------------+--------------------+---------------------+--------------------+--------------------+-------------------------------+ +| {catalog_procedure_id} | {database_procedure_id} | reconcile_catalog | reconcile_database | greptime | | greptime | reconciliation_event_database | ++--------------------------------------+--------------------------------------+-------------------+--------------------+---------------------+--------------------+--------------------+-------------------------------+" + ), + ) + .await; + + assert_eventually_eq( + &frontend, + &format!( + "SELECT count(*) AS event_count, \ + json_get_string(procedure_trigger, 'outcome') AS outcome \ + FROM {EVENTS_TABLE} \ + WHERE type = 'reconcile_catalog' AND procedure_id = '{catalog_procedure_id}' \ + AND json_get_string(procedure_trigger, 'type') = 'ChildSubmitted' \ + AND json_get_string(procedure_trigger, 'procedure_id') = '{database_procedure_id}' \ + AND catalog_name = '{CATALOG}' AND schema_name IS NULL \ + AND json_is_null(payload) \ + GROUP BY json_get_string(procedure_trigger, 'outcome')" + ), + "\ ++-------------+----------+ +| event_count | outcome | ++-------------+----------+ +| 1 | Accepted | ++-------------+----------+", + ) + .await; + + assert_eventually_eq( + &frontend, + &format!( + "SELECT count(*) AS event_count, \ + json_get_int(payload, 'version') AS version, \ + json_get_string(payload, 'resolve_strategy') AS resolve_strategy, \ + json_get_bool(payload, 'fail_fast') AS fail_fast, \ + json_get_int(payload, 'parallelism') AS parallelism, \ + json_get_bool(payload, 'is_subprocedure') AS is_subprocedure \ + FROM {EVENTS_TABLE} \ + WHERE type = 'reconcile_database' AND procedure_id = '{database_procedure_id}' \ + AND json_get_string(procedure_trigger, 'type') = 'Submitted' \ + AND catalog_name = '{CATALOG}' AND schema_name = '{DATABASE}' \ + AND table_name IS NULL AND table_id IS NULL AND physical_table_id IS NULL \ + GROUP BY json_get_int(payload, 'version'), \ + json_get_string(payload, 'resolve_strategy'), \ + json_get_bool(payload, 'fail_fast'), \ + json_get_int(payload, 'parallelism'), \ + json_get_bool(payload, 'is_subprocedure')" + ), + "\ ++-------------+---------+------------------+-----------+-------------+-----------------+ +| event_count | version | resolve_strategy | fail_fast | parallelism | is_subprocedure | ++-------------+---------+------------------+-----------+-------------+-----------------+ +| 1 | 1 | use_latest | false | 1 | true | ++-------------+---------+------------------+-----------+-------------+-----------------+", + ) + .await; + + assert_eventually_eq( + &frontend, + &format!( + "SELECT count(*) AS event_count, \ + json_get_int(payload, 'version') AS version, \ + json_get_bool(payload, 'complete') AS complete, \ + json_get_int(payload, 'processed_table_count') AS processed_count, \ + json_get_int(payload, 'succeeded_table_count') AS succeeded_count, \ + json_get_int(payload, 'failed_table_count') AS failed_count, \ + json_get_int(payload, 'succeeded_subprocedure_count') AS succeeded_child_count, \ + json_get_int(payload, 'failed_subprocedure_count') AS failed_child_count \ + FROM {EVENTS_TABLE} \ + WHERE type = 'reconcile_database' AND procedure_id = '{database_procedure_id}' \ + AND json_get_string(procedure_trigger, 'type') = 'Succeeded' \ + AND catalog_name = '{CATALOG}' AND schema_name = '{DATABASE}' \ + AND table_name IS NULL AND table_id IS NULL AND physical_table_id IS NULL \ + GROUP BY json_get_int(payload, 'version'), \ + json_get_bool(payload, 'complete'), \ + json_get_int(payload, 'processed_table_count'), \ + json_get_int(payload, 'succeeded_table_count'), \ + json_get_int(payload, 'failed_table_count'), \ + json_get_int(payload, 'succeeded_subprocedure_count'), \ + json_get_int(payload, 'failed_subprocedure_count')" + ), + "\ ++-------------+---------+----------+-----------------+-----------------+--------------+-----------------------+--------------------+ +| event_count | version | complete | processed_count | succeeded_count | failed_count | succeeded_child_count | failed_child_count | ++-------------+---------+----------+-----------------+-----------------+--------------+-----------------------+--------------------+ +| 1 | 1 | true | 1 | 1 | 0 | 1 | 0 | ++-------------+---------+----------+-----------------+-----------------+--------------+-----------------------+--------------------+", + ) + .await; +} + #[tokio::test(flavor = "multi_thread")] async fn test_table_reconciliation_events() { common_telemetry::init_default_ut_logging(); diff --git a/tests-integration/tests/repartition.rs b/tests-integration/tests/repartition.rs index cf64f40aa00..d25efb92010 100644 --- a/tests-integration/tests/repartition.rs +++ b/tests-integration/tests/repartition.rs @@ -812,6 +812,7 @@ async fn trigger_table_gc(metasrv: &Arc, table_name: &str) { let procedure = BatchGcProcedure::new( metasrv.mailbox().clone(), metasrv.table_metadata_manager().clone(), + metasrv.runtime_switch_manager().clone(), metasrv.options().grpc.server_addr.clone(), region_ids.clone(), false, // full_file_listing diff --git a/tests/cases/standalone/common/aggregate/max_abs_regression.result b/tests/cases/standalone/common/aggregate/max_abs_regression.result new file mode 100644 index 00000000000..a0e0b394667 --- /dev/null +++ b/tests/cases/standalone/common/aggregate/max_abs_regression.result @@ -0,0 +1,131 @@ +-- Regression coverage for aggregate expressions with multiple MAX aggregates. +CREATE DATABASE max_abs_regression; + +Affected Rows: 1 + +USE max_abs_regression; + +Affected Rows: 0 + +CREATE TABLE t ( + ts TIMESTAMP TIME INDEX, + v INT, + PRIMARY KEY (v), +) ENGINE = mito; + +Affected Rows: 0 + +INSERT INTO t VALUES + ('2024-01-01 00:00:00', -11), + ('2024-01-01 00:01:00', 3), + ('2024-01-01 00:02:00', -7), + ('2024-01-01 00:03:00', 5); + +Affected Rows: 4 + +-- Expression aggregate regression. +SELECT MAX(ABS(v)), MAX(ts) FROM t; + ++---------------+---------------------+ +| max(abs(t.v)) | max(t.ts) | ++---------------+---------------------+ +| 11 | 2024-01-01T00:03:00 | ++---------------+---------------------+ + +SELECT MAX(ts), MAX(ABS(v)) FROM t; + ++---------------------+---------------+ +| max(t.ts) | max(abs(t.v)) | ++---------------------+---------------+ +| 2024-01-01T00:03:00 | 11 | ++---------------------+---------------+ + +-- Column-only multiple-MAX positive control. +SELECT MAX(v), MAX(ts) FROM t; + ++----------+---------------------+ +| max(t.v) | max(t.ts) | ++----------+---------------------+ +| 5 | 2024-01-01T00:03:00 | ++----------+---------------------+ + +-- A mixed expression/column MAX must not send an incomplete dynamic filter to DN SeqScans. +-- SQLNESS REPLACE ("metrics_per_partition":\s*.*metrics=) "metrics_per_partition": REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (=Hash.*) =REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\["DynamicFilter\s\[[^"]*\]"\] , "dyn_filters": ["DynamicFilter [ REDACTED ]"] +-- SQLNESS REPLACE metrics=REDACTED\s*\| metrics=REDACTED_| +EXPLAIN ANALYZE VERBOSE SELECT MAX(ABS(v)), MAX(ts) FROM t; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[max(abs(t.v)), max(t.ts)] metrics=REDACTED_| +|_|_|_CoalescePartitionsExec metrics=REDACTED_| +|_|_|_AggregateExec: mode=Partial, gby=[], aggr=[max(abs(t.v)), max(t.ts)] metrics=REDACTED_| +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["ts", "v"], "flat_format":REDACTED, "metrics_per_partition": REDACTED metrics=REDACTED_| +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +-- A column-only multiple-MAX query remains a dynamic-filter positive control. +-- SQLNESS REPLACE ("metrics_per_partition":\s*.*metrics=) "metrics_per_partition": REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (=Hash.*) =REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\["DynamicFilter\s\[[^"]*\]"\] , "dyn_filters": ["DynamicFilter [ REDACTED ]"] +-- SQLNESS REPLACE metrics=REDACTED\s*\| metrics=REDACTED_| +EXPLAIN ANALYZE VERBOSE SELECT MAX(v), MAX(ts) FROM t; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=REDACTED_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_AggregateExec: mode=Final, gby=[], aggr=[max(t.v), max(t.ts)] metrics=REDACTED_| +|_|_|_CoalescePartitionsExec metrics=REDACTED_| +|_|_|_AggregateExec: mode=Partial, gby=[], aggr=[max(t.v), max(t.ts)] metrics=REDACTED_| +|_|_|_SeqScan: region=REDACTED, {"partition_count":REDACTED, "projection": ["ts", "v"], "dyn_filters": ["DynamicFilter [ REDACTED ]"], "flat_format":REDACTED, "metrics_per_partition": REDACTED metrics=REDACTED_| +|_|_|_| +|_|_| Total rows: REDACTED_| ++-+-+-+ + +DROP TABLE t; + +Affected Rows: 0 + +USE public; + +Affected Rows: 0 + +DROP DATABASE max_abs_regression; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/aggregate/max_abs_regression.sql b/tests/cases/standalone/common/aggregate/max_abs_regression.sql new file mode 100644 index 00000000000..c9705f6086e --- /dev/null +++ b/tests/cases/standalone/common/aggregate/max_abs_regression.sql @@ -0,0 +1,69 @@ +-- Regression coverage for aggregate expressions with multiple MAX aggregates. +CREATE DATABASE max_abs_regression; + +USE max_abs_regression; + +CREATE TABLE t ( + ts TIMESTAMP TIME INDEX, + v INT, + PRIMARY KEY (v), +) ENGINE = mito; + +INSERT INTO t VALUES + ('2024-01-01 00:00:00', -11), + ('2024-01-01 00:01:00', 3), + ('2024-01-01 00:02:00', -7), + ('2024-01-01 00:03:00', 5); + +-- Expression aggregate regression. +SELECT MAX(ABS(v)), MAX(ts) FROM t; +SELECT MAX(ts), MAX(ABS(v)) FROM t; + +-- Column-only multiple-MAX positive control. +SELECT MAX(v), MAX(ts) FROM t; + +-- A mixed expression/column MAX must not send an incomplete dynamic filter to DN SeqScans. +-- SQLNESS REPLACE ("metrics_per_partition":\s*.*metrics=) "metrics_per_partition": REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (=Hash.*) =REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\["DynamicFilter\s\[[^"]*\]"\] , "dyn_filters": ["DynamicFilter [ REDACTED ]"] +-- SQLNESS REPLACE metrics=REDACTED\s*\| metrics=REDACTED_| +EXPLAIN ANALYZE VERBOSE SELECT MAX(ABS(v)), MAX(ts) FROM t; + +-- A column-only multiple-MAX query remains a dynamic-filter positive control. +-- SQLNESS REPLACE ("metrics_per_partition":\s*.*metrics=) "metrics_per_partition": REDACTED metrics= +-- SQLNESS REPLACE (metrics=\{.*\}) metrics=REDACTED +-- SQLNESS REPLACE (metrics=\[[^\]]*\]) metrics=REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (=Hash.*) =REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE "(file_id|time_range_start|time_range_end)":"[^"]+" "$1":"REDACTED" +-- SQLNESS REPLACE ("[a-z_]+":"[0-9\.]+(ns|us|µs|ms|s)") "DURATION": REDACTED +-- SQLNESS REPLACE "(size|flat_format)":\s*(\d+|true|false) "$1":REDACTED +-- SQLNESS REPLACE ,\s*filter=.*?metrics= metrics= +-- SQLNESS REPLACE Total\s+rows:\s+\d+ Total rows: REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED +-- SQLNESS REPLACE ,\s"dyn_filters":\s\["DynamicFilter\s\[[^"]*\]"\] , "dyn_filters": ["DynamicFilter [ REDACTED ]"] +-- SQLNESS REPLACE metrics=REDACTED\s*\| metrics=REDACTED_| +EXPLAIN ANALYZE VERBOSE SELECT MAX(v), MAX(ts) FROM t; + +DROP TABLE t; + +USE public; + +DROP DATABASE max_abs_regression; diff --git a/tests/cases/standalone/common/promql/counter_reset_precision.result b/tests/cases/standalone/common/promql/counter_reset_precision.result new file mode 100644 index 00000000000..3482f7ca9be --- /dev/null +++ b/tests/cases/standalone/common/promql/counter_reset_precision.result @@ -0,0 +1,43 @@ +CREATE TABLE counter_reset_precision ( + ts TIMESTAMP TIME INDEX, + greptime_value DOUBLE, + series STRING PRIMARY KEY +); + +Affected Rows: 0 + +INSERT INTO counter_reset_precision VALUES + (0, 10000000000000000.0, 'a'), + (30000, 1.0, 'a'), + (60000, 0.0, 'a'), + (90000, 1.0, 'a'), + (120000, 2.0, 'a'), + (150000, 3.0, 'a'), + (180000, 4.0, 'a'); + +Affected Rows: 7 + +-- The counter resets twice, at 30s from 1e16 and at 60s from 1.0, and 1e16 + 1.0 is 1e16 in +-- f64. Windows are two minutes wide and step by one sample, so every row here depends on the +-- 1.0 reset surviving. +-- +-- Summing the two corrections on their own drops the 1.0 and then cancels against the first +-- sample, so the 90s window reports no increase. Carrying that sum into the next window and +-- subtracting the 1e16 that left it reports half the increase, and the windows after that +-- subtract a reset the sum never held, so the correction turns negative and stays there for +-- the rest of the batch. +TQL EVAL (90, 180, '30s') increase(counter_reset_precision[2m]); + ++---------------------+---------------------------------------------------------+--------+ +| ts | prom_increase(ts_range,greptime_value,ts,Int64(120000)) | series | ++---------------------+---------------------------------------------------------+--------+ +| 1970-01-01T00:01:30 | 1.3333333333333333 | a | +| 1970-01-01T00:02:00 | 2.6666666666666665 | a | +| 1970-01-01T00:02:30 | 3.0 | a | +| 1970-01-01T00:03:00 | 4.0 | a | ++---------------------+---------------------------------------------------------+--------+ + +DROP TABLE counter_reset_precision; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/counter_reset_precision.sql b/tests/cases/standalone/common/promql/counter_reset_precision.sql new file mode 100644 index 00000000000..867260dc3da --- /dev/null +++ b/tests/cases/standalone/common/promql/counter_reset_precision.sql @@ -0,0 +1,27 @@ +CREATE TABLE counter_reset_precision ( + ts TIMESTAMP TIME INDEX, + greptime_value DOUBLE, + series STRING PRIMARY KEY +); + +INSERT INTO counter_reset_precision VALUES + (0, 10000000000000000.0, 'a'), + (30000, 1.0, 'a'), + (60000, 0.0, 'a'), + (90000, 1.0, 'a'), + (120000, 2.0, 'a'), + (150000, 3.0, 'a'), + (180000, 4.0, 'a'); + +-- The counter resets twice, at 30s from 1e16 and at 60s from 1.0, and 1e16 + 1.0 is 1e16 in +-- f64. Windows are two minutes wide and step by one sample, so every row here depends on the +-- 1.0 reset surviving. +-- +-- Summing the two corrections on their own drops the 1.0 and then cancels against the first +-- sample, so the 90s window reports no increase. Carrying that sum into the next window and +-- subtracting the 1e16 that left it reports half the increase, and the windows after that +-- subtract a reset the sum never held, so the correction turns negative and stays there for +-- the rest of the batch. +TQL EVAL (90, 180, '30s') increase(counter_reset_precision[2m]); + +DROP TABLE counter_reset_precision; diff --git a/tests/cases/standalone/common/promql/native_time_selection.result b/tests/cases/standalone/common/promql/native_time_selection.result new file mode 100644 index 00000000000..80d18319ad7 --- /dev/null +++ b/tests/cases/standalone/common/promql/native_time_selection.result @@ -0,0 +1,611 @@ +-- Regression coverage for instant and range selection on native microsecond and +-- nanosecond time indexes. +CREATE TABLE native_time_us ( + ts TIMESTAMP(6) TIME INDEX, + series STRING PRIMARY KEY, + val DOUBLE, +); + +Affected Rows: 0 + +INSERT INTO native_time_us VALUES + (1000001, 'future', 101), + (1000000, 'exact', 201), + (1000001, 'exact', 202), + (-299000000, 'lowerbound', 301), + (-298999999, 'lowerplus', 302), + (1000000, 'positive_lowerbound', 701), + (1000001, 'positive_lowerplus', 702), + (1000001, 'multi', 401), + (-1000000, 'offset', 501), + (0, 'offset', 502), + (1000000, 'offset', 503), + (999999, 'past', 602), + (999001, 'past', 601), + (0, 'window', 1), + (1, 'window', 2), + (2, 'window', 5), + (1000000, 'window', 3), + (1000001, 'window', 4); + +Affected Rows: 18 + +-- The native projection and exact 1ms-lookback bounds must reach the scan; +-- the 1s+tick row must not displace 201. +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED +-- SQLNESS REPLACE native_time_us.__table_id\s*=\s*UInt32\(\d+\) native_time_us.__table_id=UInt32(REDACTED) +TQL EXPLAIN (1, 1, '1s', '1ms') native_time_us{series="exact"}; + ++---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | PromInstantManipulate: range=[1000..1000], lookback=[1], interval=[1000], time index=[ts] | +| | PromSeriesDivide: tags=["series"] | +| | Sort: native_time_us.series ASC NULLS FIRST, native_time_us.ts ASC NULLS FIRST | +| | Projection: native_time_us.val, native_time_us.series, native_time_us.ts | +| | Filter: native_time_us.series = Utf8("exact") AND native_time_us.ts > TimestampMicrosecond(999000, None) AND native_time_us.ts <= TimestampMicrosecond(1000000, None) | +| | TableScan: native_time_us, partial_filters=[native_time_us.series = Utf8("exact"), native_time_us.ts > TimestampMicrosecond(999000, None), native_time_us.ts <= TimestampMicrosecond(1000000, None)] | +| | ]] | +| physical_plan | CooperativeExec | +| | MergeScanExec: REDACTED +| | | ++---------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ + +-- The actual memtable scan must use LastRow { after_merge: true } with native +-- 1ms-lookback bounds; it must select exact 1s rather than the future tick. +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +TQL ANALYZE VERBOSE (1, 1, '1s', '1ms') native_time_us{series="exact"}; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=[]_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_PromInstantManipulateExec: range=[1000..1000], lookback=[1], interval=[1000], time index=[ts] metrics=[output_rows: 1, REDACTED +|_|_|_PromSeriesDivideExec: tags=["series"] metrics=[output_rows: 1, REDACTED +|_|_|_ProjectionExec: expr=[val@2 as val, series@1 as series, ts@0 as ts] metrics=[output_rows: 1, REDACTED +|_|_|_CooperativeExec metrics=[]_| +|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "series", "val"], "filters": ["series = Dictionary(UInt32, Utf8(\"exact\"))", "ts > TimestampMicrosecond(999000, None)", "ts <= TimestampMicrosecond(1000000, None)"], "REDACTED +|_|_|_| +|_|_| Total rows: 1_| ++-+-+-+ + +-- The same-series future tick is in the memtable, while exact 1s remains selected. +TQL EVAL (1, 1, '1s', '1ms') native_time_us{series="exact"}; + ++-------+--------+---------------------+ +| val | series | ts | ++-------+--------+---------------------+ +| 201.0 | exact | 1970-01-01T00:00:01 | ++-------+--------+---------------------+ + +-- Future-only selection is empty before flushing, exercising the memtable path. +TQL EVAL (1, 1, '1s', '300s') native_time_us{series="future"}; + +++ +++ + +ADMIN FLUSH_TABLE('native_time_us'); + ++-------------------------------------+ +| ADMIN FLUSH_TABLE('native_time_us') | ++-------------------------------------+ +| 0 | ++-------------------------------------+ + +-- The exact native sample remains selected from the flushed SST. +TQL EVAL (1, 1, '1s', '1ms') native_time_us{series="exact"}; + ++-------+--------+---------------------+ +| val | series | ts | ++-------+--------+---------------------+ +| 201.0 | exact | 1970-01-01T00:00:01 | ++-------+--------+---------------------+ + +TQL EVAL (1, 1, '1s', '300s') timestamp(native_time_us{series="future"}); + +++ +++ + +TQL EVAL (1, 1, '1s', '300s') timestamp(native_time_us{series="exact"}); + ++---------------------+-------+--------+ +| ts | value | series | ++---------------------+-------+--------+ +| 1970-01-01T00:00:01 | 1.0 | exact | ++---------------------+-------+--------+ + +-- Instant lookback bounds are exclusive: these return only 302 and 702. +TQL EVAL (1, 1, '1s', '300s') native_time_us{series=~"lower.*"}; + ++-------+-----------+---------------------+ +| val | series | ts | ++-------+-----------+---------------------+ +| 302.0 | lowerplus | 1970-01-01T00:00:01 | ++-------+-----------+---------------------+ + +TQL EVAL (301, 301, '1s', '300s') native_time_us{series=~"positive_lower.*"}; + ++-------+--------------------+---------------------+ +| val | series | ts | ++-------+--------------------+---------------------+ +| 702.0 | positive_lowerplus | 1970-01-01T00:05:01 | ++-------+--------------------+---------------------+ + +-- The sub-millisecond point belongs only to the 2s evaluation step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (1, 2, '1s', '300s') native_time_us{series="multi"}; + ++-------+--------+---------------------+ +| val | series | ts | ++-------+--------+---------------------+ +| 401.0 | multi | 1970-01-01T00:00:02 | ++-------+--------+---------------------+ + +-- The latest native timestamp below 1s is retained even when inserts are unordered. +TQL EVAL (1, 1, '1s', '300s') native_time_us{series="past"}; + ++-------+--------+---------------------+ +| val | series | ts | ++-------+--------+---------------------+ +| 602.0 | past | 1970-01-01T00:00:01 | ++-------+--------+---------------------+ + +-- Offsets select native timestamps, including stored negative time. +TQL EVAL (0, 0, '1s', '300s') native_time_us{series="offset"}; + ++-------+--------+---------------------+ +| val | series | ts | ++-------+--------+---------------------+ +| 502.0 | offset | 1970-01-01T00:00:00 | ++-------+--------+---------------------+ + +TQL EVAL (0, 0, '1s', '300s') native_time_us{series="offset"} offset 1s; + ++-------+--------+---------------------+ +| val | series | ts | ++-------+--------+---------------------+ +| 501.0 | offset | 1970-01-01T00:00:00 | ++-------+--------+---------------------+ + +TQL EVAL (0, 0, '1s', '300s') native_time_us{series="offset"} offset -1s; + ++-------+--------+---------------------+ +| val | series | ts | ++-------+--------+---------------------+ +| 503.0 | offset | 1970-01-01T00:00:00 | ++-------+--------+---------------------+ + +-- [1s] at 1s excludes 0 and 1s+tick, retaining 0+tick, 0+2ticks, and 1s. +TQL EVAL (1, 1, '1s', '300s') count_over_time(native_time_us{series="window"}[1s]); + ++---------------------+------------------------------------+--------+ +| ts | prom_count_over_time(ts_range,val) | series | ++---------------------+------------------------------------+--------+ +| 1970-01-01T00:00:01 | 3.0 | window | ++---------------------+------------------------------------+--------+ + +TQL EVAL (1, 1, '1s', '300s') sum_over_time(native_time_us{series="window"}[1s]); + ++---------------------+----------------------------------+--------+ +| ts | prom_sum_over_time(ts_range,val) | series | ++---------------------+----------------------------------+--------+ +| 1970-01-01T00:00:01 | 10.0 | window | ++---------------------+----------------------------------+--------+ + +TQL EVAL (1, 1, '1s', '300s') last_over_time(native_time_us{series="window"}[1s]); + ++---------------------+-----------------------------------+--------+ +| ts | prom_last_over_time(ts_range,val) | series | ++---------------------+-----------------------------------+--------+ +| 1970-01-01T00:00:01 | 3.0 | window | ++---------------------+-----------------------------------+--------+ + +-- The inner selector consumes native time; the subquery consumes ms evaluations. +TQL EVAL (1, 1, '1s') last_over_time((native_time_us{series="exact"})[1s:1s]); + ++---------------------+-----------------------------------+--------+ +| ts | prom_last_over_time(ts_range,val) | series | ++---------------------+-----------------------------------+--------+ +| 1970-01-01T00:00:01 | 201.0 | exact | ++---------------------+-----------------------------------+--------+ + +-- Inner offsets are applied once when evaluating the subquery selector. +TQL EVAL (0, 0, '1s') last_over_time((native_time_us{series="offset"} offset 1s)[1s:1s]); + ++---------------------+-----------------------------------+--------+ +| ts | prom_last_over_time(ts_range,val) | series | ++---------------------+-----------------------------------+--------+ +| 1970-01-01T00:00:00 | 501.0 | offset | ++---------------------+-----------------------------------+--------+ + +TQL EVAL (0, 0, '1s') last_over_time((native_time_us{series="offset"} offset -1s)[1s:1s]); + ++---------------------+-----------------------------------+--------+ +| ts | prom_last_over_time(ts_range,val) | series | ++---------------------+-----------------------------------+--------+ +| 1970-01-01T00:00:00 | 503.0 | offset | ++---------------------+-----------------------------------+--------+ + +DROP TABLE native_time_us; + +Affected Rows: 0 + +CREATE TABLE native_time_ns ( + ts TIMESTAMP(9) TIME INDEX, + series STRING PRIMARY KEY, + val DOUBLE, +); + +Affected Rows: 0 + +INSERT INTO native_time_ns VALUES + (1000000001, 'future', 101), + (1000000000, 'exact', 201), + (1000000001, 'exact', 202), + (-299000000000, 'lowerbound', 301), + (-298999999999, 'lowerplus', 302), + (1000000000, 'positive_lowerbound', 701), + (1000000001, 'positive_lowerplus', 702), + (1000000001, 'multi', 401), + (-1000000000, 'offset', 501), + (0, 'offset', 502), + (1000000000, 'offset', 503), + (999999000, 'past', 602), + (999001000, 'past', 601), + (0, 'window', 1), + (1, 'window', 2), + (2, 'window', 5), + (1000000000, 'window', 3), + (1000000001, 'window', 4); + +Affected Rows: 18 + +-- The native projection and exact 1ms-lookback bounds must reach the scan; +-- the 1s+tick row must not displace 201. +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED +-- SQLNESS REPLACE native_time_ns.__table_id\s*=\s*UInt32\(\d+\) native_time_ns.__table_id=UInt32(REDACTED) +TQL EXPLAIN (1, 1, '1s', '1ms') native_time_ns{series="exact"}; + ++---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | PromInstantManipulate: range=[1000..1000], lookback=[1], interval=[1000], time index=[ts] | +| | PromSeriesDivide: tags=["series"] | +| | Sort: native_time_ns.series ASC NULLS FIRST, native_time_ns.ts ASC NULLS FIRST | +| | Projection: native_time_ns.val, native_time_ns.series, native_time_ns.ts | +| | Filter: native_time_ns.series = Utf8("exact") AND native_time_ns.ts > TimestampNanosecond(999000000, None) AND native_time_ns.ts <= TimestampNanosecond(1000000000, None) | +| | TableScan: native_time_ns, partial_filters=[native_time_ns.series = Utf8("exact"), native_time_ns.ts > TimestampNanosecond(999000000, None), native_time_ns.ts <= TimestampNanosecond(1000000000, None)] | +| | ]] | +| physical_plan | CooperativeExec | +| | MergeScanExec: REDACTED +| | | ++---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ + +-- The actual memtable scan must use LastRow { after_merge: true } with native +-- 1ms-lookback bounds; it must select exact 1s rather than the future tick. +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +TQL ANALYZE VERBOSE (1, 1, '1s', '1ms') native_time_ns{series="exact"}; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=[]_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_PromInstantManipulateExec: range=[1000..1000], lookback=[1], interval=[1000], time index=[ts] metrics=[output_rows: 1, REDACTED +|_|_|_PromSeriesDivideExec: tags=["series"] metrics=[output_rows: 1, REDACTED +|_|_|_ProjectionExec: expr=[val@2 as val, series@1 as series, ts@0 as ts] metrics=[output_rows: 1, REDACTED +|_|_|_CooperativeExec metrics=[]_| +|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "series", "val"], "filters": ["series = Dictionary(UInt32, Utf8(\"exact\"))", "ts > TimestampNanosecond(999000000, None)", "ts <= TimestampNanosecond(1000000000, None)"], "REDACTED +|_|_|_| +|_|_| Total rows: 1_| ++-+-+-+ + +-- The same-series future tick is in the memtable, while exact 1s remains selected. +TQL EVAL (1, 1, '1s', '1ms') native_time_ns{series="exact"}; + ++-------+--------+---------------------+ +| val | series | ts | ++-------+--------+---------------------+ +| 201.0 | exact | 1970-01-01T00:00:01 | ++-------+--------+---------------------+ + +-- Future-only selection is empty before flushing, exercising the memtable path. +TQL EVAL (1, 1, '1s', '300s') native_time_ns{series="future"}; + +++ +++ + +ADMIN FLUSH_TABLE('native_time_ns'); + ++-------------------------------------+ +| ADMIN FLUSH_TABLE('native_time_ns') | ++-------------------------------------+ +| 0 | ++-------------------------------------+ + +-- The exact native sample remains selected from the flushed SST. +TQL EVAL (1, 1, '1s', '1ms') native_time_ns{series="exact"}; + ++-------+--------+---------------------+ +| val | series | ts | ++-------+--------+---------------------+ +| 201.0 | exact | 1970-01-01T00:00:01 | ++-------+--------+---------------------+ + +TQL EVAL (1, 1, '1s', '300s') timestamp(native_time_ns{series="future"}); + +++ +++ + +TQL EVAL (1, 1, '1s', '300s') timestamp(native_time_ns{series="exact"}); + ++---------------------+-------+--------+ +| ts | value | series | ++---------------------+-------+--------+ +| 1970-01-01T00:00:01 | 1.0 | exact | ++---------------------+-------+--------+ + +-- Instant lookback bounds are exclusive: these return only 302 and 702. +TQL EVAL (1, 1, '1s', '300s') native_time_ns{series=~"lower.*"}; + ++-------+-----------+---------------------+ +| val | series | ts | ++-------+-----------+---------------------+ +| 302.0 | lowerplus | 1970-01-01T00:00:01 | ++-------+-----------+---------------------+ + +TQL EVAL (301, 301, '1s', '300s') native_time_ns{series=~"positive_lower.*"}; + ++-------+--------------------+---------------------+ +| val | series | ts | ++-------+--------------------+---------------------+ +| 702.0 | positive_lowerplus | 1970-01-01T00:05:01 | ++-------+--------------------+---------------------+ + +-- The sub-millisecond point belongs only to the 2s evaluation step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (1, 2, '1s', '300s') native_time_ns{series="multi"}; + ++-------+--------+---------------------+ +| val | series | ts | ++-------+--------+---------------------+ +| 401.0 | multi | 1970-01-01T00:00:02 | ++-------+--------+---------------------+ + +-- The latest native timestamp below 1s is retained even when inserts are unordered. +TQL EVAL (1, 1, '1s', '300s') native_time_ns{series="past"}; + ++-------+--------+---------------------+ +| val | series | ts | ++-------+--------+---------------------+ +| 602.0 | past | 1970-01-01T00:00:01 | ++-------+--------+---------------------+ + +-- Offsets select native timestamps, including stored negative time. +TQL EVAL (0, 0, '1s', '300s') native_time_ns{series="offset"}; + ++-------+--------+---------------------+ +| val | series | ts | ++-------+--------+---------------------+ +| 502.0 | offset | 1970-01-01T00:00:00 | ++-------+--------+---------------------+ + +TQL EVAL (0, 0, '1s', '300s') native_time_ns{series="offset"} offset 1s; + ++-------+--------+---------------------+ +| val | series | ts | ++-------+--------+---------------------+ +| 501.0 | offset | 1970-01-01T00:00:00 | ++-------+--------+---------------------+ + +TQL EVAL (0, 0, '1s', '300s') native_time_ns{series="offset"} offset -1s; + ++-------+--------+---------------------+ +| val | series | ts | ++-------+--------+---------------------+ +| 503.0 | offset | 1970-01-01T00:00:00 | ++-------+--------+---------------------+ + +-- [1s] at 1s excludes 0 and 1s+tick, retaining 0+tick, 0+2ticks, and 1s. +TQL EVAL (1, 1, '1s', '300s') count_over_time(native_time_ns{series="window"}[1s]); + ++---------------------+------------------------------------+--------+ +| ts | prom_count_over_time(ts_range,val) | series | ++---------------------+------------------------------------+--------+ +| 1970-01-01T00:00:01 | 3.0 | window | ++---------------------+------------------------------------+--------+ + +TQL EVAL (1, 1, '1s', '300s') sum_over_time(native_time_ns{series="window"}[1s]); + ++---------------------+----------------------------------+--------+ +| ts | prom_sum_over_time(ts_range,val) | series | ++---------------------+----------------------------------+--------+ +| 1970-01-01T00:00:01 | 10.0 | window | ++---------------------+----------------------------------+--------+ + +TQL EVAL (1, 1, '1s', '300s') last_over_time(native_time_ns{series="window"}[1s]); + ++---------------------+-----------------------------------+--------+ +| ts | prom_last_over_time(ts_range,val) | series | ++---------------------+-----------------------------------+--------+ +| 1970-01-01T00:00:01 | 3.0 | window | ++---------------------+-----------------------------------+--------+ + +-- The inner selector consumes native time; the subquery consumes ms evaluations. +TQL EVAL (1, 1, '1s') last_over_time((native_time_ns{series="exact"})[1s:1s]); + ++---------------------+-----------------------------------+--------+ +| ts | prom_last_over_time(ts_range,val) | series | ++---------------------+-----------------------------------+--------+ +| 1970-01-01T00:00:01 | 201.0 | exact | ++---------------------+-----------------------------------+--------+ + +-- Inner offsets are applied once when evaluating the subquery selector. +TQL EVAL (0, 0, '1s') last_over_time((native_time_ns{series="offset"} offset 1s)[1s:1s]); + ++---------------------+-----------------------------------+--------+ +| ts | prom_last_over_time(ts_range,val) | series | ++---------------------+-----------------------------------+--------+ +| 1970-01-01T00:00:00 | 501.0 | offset | ++---------------------+-----------------------------------+--------+ + +TQL EVAL (0, 0, '1s') last_over_time((native_time_ns{series="offset"} offset -1s)[1s:1s]); + ++---------------------+-----------------------------------+--------+ +| ts | prom_last_over_time(ts_range,val) | series | ++---------------------+-----------------------------------+--------+ +| 1970-01-01T00:00:00 | 503.0 | offset | ++---------------------+-----------------------------------+--------+ + +DROP TABLE native_time_ns; + +Affected Rows: 0 + +-- An unrepresentable native lower bound must not discard its representable upper bound. +-- The upper filter must reach LastRow so the 1ms-future row cannot hide the eligible row. +CREATE TABLE native_time_ns_lower_overflow ( + ts TIMESTAMP(9) TIME INDEX, + series STRING PRIMARY KEY, + val DOUBLE, +); + +Affected Rows: 0 + +INSERT INTO native_time_ns_lower_overflow VALUES + (-9223200000000000000, 'exact', 1), + (-9223199999999000000, 'exact', 2); + +Affected Rows: 2 + +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE native_time_ns_lower_overflow.__table_id\s*=\s*UInt32\(\d+\) native_time_ns_lower_overflow.__table_id=UInt32(REDACTED) +TQL EXPLAIN (0, 0, '1s', '2d') native_time_ns_lower_overflow{series="exact"} offset 106750d; + ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| plan_type | plan | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| logical_plan | MergeScan [is_placeholder=false, remote_input=[ | +| | PromInstantManipulate: range=[0..0], lookback=[172800000], interval=[1000], time index=[ts] | +| | PromSeriesNormalize: offset=[9223200000000], time index=[ts], filter NaN: [false] | +| | PromSeriesDivide: tags=["series"] | +| | Sort: native_time_ns_lower_overflow.series ASC NULLS FIRST, native_time_ns_lower_overflow.ts ASC NULLS FIRST | +| | Projection: native_time_ns_lower_overflow.val, native_time_ns_lower_overflow.series, native_time_ns_lower_overflow.ts | +| | Filter: native_time_ns_lower_overflow.series = Utf8("exact") AND native_time_ns_lower_overflow.ts <= TimestampNanosecond(-9223200000000000000, None) | +| | TableScan: native_time_ns_lower_overflow, partial_filters=[native_time_ns_lower_overflow.series = Utf8("exact"), native_time_ns_lower_overflow.ts <= TimestampNanosecond(-9223200000000000000, None)] | +| | ]] | +| physical_plan | CooperativeExec | +| | MergeScanExec: REDACTED +| | | ++---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ + +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +TQL ANALYZE VERBOSE (0, 0, '1s', '2d') native_time_ns_lower_overflow{series="exact"} offset 106750d; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_CooperativeExec metrics=[]_| +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_PromInstantManipulateExec: range=[0..0], lookback=[172800000], interval=[1000], time index=[ts] metrics=[output_rows: 1, REDACTED +|_|_|_PromSeriesNormalizeExec: offset=[9223200000000], time index=[ts], filter NaN: [false] metrics=[output_rows: 1, REDACTED +|_|_|_PromSeriesDivideExec: tags=["series"] metrics=[output_rows: 1, REDACTED +|_|_|_ProjectionExec: expr=[val@2 as val, series@1 as series, ts@0 as ts] metrics=[output_rows: 1, REDACTED +|_|_|_CooperativeExec metrics=[]_| +|_|_|_SeriesScan: region=REDACTED, {"partition_count":{"count":1, "mem_ranges":1, "files":0, "file_ranges":0}, "selector":"LastRow { after_merge: true }", "distribution":"PerSeries", "projection": ["ts", "series", "val"], "filters": ["series = Dictionary(UInt32, Utf8(\"exact\"))", "ts <= TimestampNanosecond(-9223200000000000000, None)"], "REDACTED +|_|_|_| +|_|_| Total rows: 1_| ++-+-+-+ + +-- The representable upper bound selects only the exact row. +TQL EVAL (0, 0, '1s', '2d') native_time_ns_lower_overflow{series="exact"} offset 106750d; + ++-----+--------+---------------------+ +| val | series | ts | ++-----+--------+---------------------+ +| 1.0 | exact | 1970-01-01T00:00:00 | ++-----+--------+---------------------+ + +ADMIN FLUSH_TABLE('native_time_ns_lower_overflow'); + ++----------------------------------------------------+ +| ADMIN FLUSH_TABLE('native_time_ns_lower_overflow') | ++----------------------------------------------------+ +| 0 | ++----------------------------------------------------+ + +TQL EVAL (0, 0, '1s', '2d') native_time_ns_lower_overflow{series="exact"} offset 106750d; + ++-----+--------+---------------------+ +| val | series | ts | ++-----+--------+---------------------+ +| 1.0 | exact | 1970-01-01T00:00:00 | ++-----+--------+---------------------+ + +DROP TABLE native_time_ns_lower_overflow; + +Affected Rows: 0 + +-- Second precision is promoted before applying fractional-second offsets. +CREATE TABLE native_time_sec (ts TIMESTAMP(0) TIME INDEX, val DOUBLE); + +Affected Rows: 0 + +INSERT INTO native_time_sec VALUES (0, 10), (1, 11), (2, 12); + +Affected Rows: 3 + +TQL EVAL (1, 1, '1s', '1s') native_time_sec offset 500ms; + ++------+---------------------+ +| val | ts | ++------+---------------------+ +| 10.0 | 1970-01-01T00:00:01 | ++------+---------------------+ + +TQL EVAL (1, 1, '1s', '1s') native_time_sec offset -500ms; + ++------+---------------------+ +| val | ts | ++------+---------------------+ +| 11.0 | 1970-01-01T00:00:01 | ++------+---------------------+ + +DROP TABLE native_time_sec; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/native_time_selection.sql b/tests/cases/standalone/common/promql/native_time_selection.sql new file mode 100644 index 00000000000..ca795b0b3ad --- /dev/null +++ b/tests/cases/standalone/common/promql/native_time_selection.sql @@ -0,0 +1,222 @@ +-- Regression coverage for instant and range selection on native microsecond and +-- nanosecond time indexes. + +CREATE TABLE native_time_us ( + ts TIMESTAMP(6) TIME INDEX, + series STRING PRIMARY KEY, + val DOUBLE, +); + +INSERT INTO native_time_us VALUES + (1000001, 'future', 101), + (1000000, 'exact', 201), + (1000001, 'exact', 202), + (-299000000, 'lowerbound', 301), + (-298999999, 'lowerplus', 302), + (1000000, 'positive_lowerbound', 701), + (1000001, 'positive_lowerplus', 702), + (1000001, 'multi', 401), + (-1000000, 'offset', 501), + (0, 'offset', 502), + (1000000, 'offset', 503), + (999999, 'past', 602), + (999001, 'past', 601), + (0, 'window', 1), + (1, 'window', 2), + (2, 'window', 5), + (1000000, 'window', 3), + (1000001, 'window', 4); + +-- The native projection and exact 1ms-lookback bounds must reach the scan; +-- the 1s+tick row must not displace 201. +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED +-- SQLNESS REPLACE native_time_us.__table_id\s*=\s*UInt32\(\d+\) native_time_us.__table_id=UInt32(REDACTED) +TQL EXPLAIN (1, 1, '1s', '1ms') native_time_us{series="exact"}; + +-- The actual memtable scan must use LastRow { after_merge: true } with native +-- 1ms-lookback bounds; it must select exact 1s rather than the future tick. +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +TQL ANALYZE VERBOSE (1, 1, '1s', '1ms') native_time_us{series="exact"}; + +-- The same-series future tick is in the memtable, while exact 1s remains selected. +TQL EVAL (1, 1, '1s', '1ms') native_time_us{series="exact"}; + +-- Future-only selection is empty before flushing, exercising the memtable path. +TQL EVAL (1, 1, '1s', '300s') native_time_us{series="future"}; + +ADMIN FLUSH_TABLE('native_time_us'); + +-- The exact native sample remains selected from the flushed SST. +TQL EVAL (1, 1, '1s', '1ms') native_time_us{series="exact"}; +TQL EVAL (1, 1, '1s', '300s') timestamp(native_time_us{series="future"}); +TQL EVAL (1, 1, '1s', '300s') timestamp(native_time_us{series="exact"}); + +-- Instant lookback bounds are exclusive: these return only 302 and 702. +TQL EVAL (1, 1, '1s', '300s') native_time_us{series=~"lower.*"}; +TQL EVAL (301, 301, '1s', '300s') native_time_us{series=~"positive_lower.*"}; + +-- The sub-millisecond point belongs only to the 2s evaluation step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (1, 2, '1s', '300s') native_time_us{series="multi"}; + +-- The latest native timestamp below 1s is retained even when inserts are unordered. +TQL EVAL (1, 1, '1s', '300s') native_time_us{series="past"}; + +-- Offsets select native timestamps, including stored negative time. +TQL EVAL (0, 0, '1s', '300s') native_time_us{series="offset"}; +TQL EVAL (0, 0, '1s', '300s') native_time_us{series="offset"} offset 1s; +TQL EVAL (0, 0, '1s', '300s') native_time_us{series="offset"} offset -1s; + +-- [1s] at 1s excludes 0 and 1s+tick, retaining 0+tick, 0+2ticks, and 1s. +TQL EVAL (1, 1, '1s', '300s') count_over_time(native_time_us{series="window"}[1s]); +TQL EVAL (1, 1, '1s', '300s') sum_over_time(native_time_us{series="window"}[1s]); +TQL EVAL (1, 1, '1s', '300s') last_over_time(native_time_us{series="window"}[1s]); + +-- The inner selector consumes native time; the subquery consumes ms evaluations. +TQL EVAL (1, 1, '1s') last_over_time((native_time_us{series="exact"})[1s:1s]); + +-- Inner offsets are applied once when evaluating the subquery selector. +TQL EVAL (0, 0, '1s') last_over_time((native_time_us{series="offset"} offset 1s)[1s:1s]); +TQL EVAL (0, 0, '1s') last_over_time((native_time_us{series="offset"} offset -1s)[1s:1s]); + +DROP TABLE native_time_us; + +CREATE TABLE native_time_ns ( + ts TIMESTAMP(9) TIME INDEX, + series STRING PRIMARY KEY, + val DOUBLE, +); + +INSERT INTO native_time_ns VALUES + (1000000001, 'future', 101), + (1000000000, 'exact', 201), + (1000000001, 'exact', 202), + (-299000000000, 'lowerbound', 301), + (-298999999999, 'lowerplus', 302), + (1000000000, 'positive_lowerbound', 701), + (1000000001, 'positive_lowerplus', 702), + (1000000001, 'multi', 401), + (-1000000000, 'offset', 501), + (0, 'offset', 502), + (1000000000, 'offset', 503), + (999999000, 'past', 602), + (999001000, 'past', 601), + (0, 'window', 1), + (1, 'window', 2), + (2, 'window', 5), + (1000000000, 'window', 3), + (1000000001, 'window', 4); + +-- The native projection and exact 1ms-lookback bounds must reach the scan; +-- the 1s+tick row must not displace 201. +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (RepartitionExec:.*) RepartitionExec: REDACTED +-- SQLNESS REPLACE native_time_ns.__table_id\s*=\s*UInt32\(\d+\) native_time_ns.__table_id=UInt32(REDACTED) +TQL EXPLAIN (1, 1, '1s', '1ms') native_time_ns{series="exact"}; + +-- The actual memtable scan must use LastRow { after_merge: true } with native +-- 1ms-lookback bounds; it must select exact 1s rather than the future tick. +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +TQL ANALYZE VERBOSE (1, 1, '1s', '1ms') native_time_ns{series="exact"}; + +-- The same-series future tick is in the memtable, while exact 1s remains selected. +TQL EVAL (1, 1, '1s', '1ms') native_time_ns{series="exact"}; + +-- Future-only selection is empty before flushing, exercising the memtable path. +TQL EVAL (1, 1, '1s', '300s') native_time_ns{series="future"}; + +ADMIN FLUSH_TABLE('native_time_ns'); + +-- The exact native sample remains selected from the flushed SST. +TQL EVAL (1, 1, '1s', '1ms') native_time_ns{series="exact"}; +TQL EVAL (1, 1, '1s', '300s') timestamp(native_time_ns{series="future"}); +TQL EVAL (1, 1, '1s', '300s') timestamp(native_time_ns{series="exact"}); + +-- Instant lookback bounds are exclusive: these return only 302 and 702. +TQL EVAL (1, 1, '1s', '300s') native_time_ns{series=~"lower.*"}; +TQL EVAL (301, 301, '1s', '300s') native_time_ns{series=~"positive_lower.*"}; + +-- The sub-millisecond point belongs only to the 2s evaluation step. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (1, 2, '1s', '300s') native_time_ns{series="multi"}; + +-- The latest native timestamp below 1s is retained even when inserts are unordered. +TQL EVAL (1, 1, '1s', '300s') native_time_ns{series="past"}; + +-- Offsets select native timestamps, including stored negative time. +TQL EVAL (0, 0, '1s', '300s') native_time_ns{series="offset"}; +TQL EVAL (0, 0, '1s', '300s') native_time_ns{series="offset"} offset 1s; +TQL EVAL (0, 0, '1s', '300s') native_time_ns{series="offset"} offset -1s; + +-- [1s] at 1s excludes 0 and 1s+tick, retaining 0+tick, 0+2ticks, and 1s. +TQL EVAL (1, 1, '1s', '300s') count_over_time(native_time_ns{series="window"}[1s]); +TQL EVAL (1, 1, '1s', '300s') sum_over_time(native_time_ns{series="window"}[1s]); +TQL EVAL (1, 1, '1s', '300s') last_over_time(native_time_ns{series="window"}[1s]); + +-- The inner selector consumes native time; the subquery consumes ms evaluations. +TQL EVAL (1, 1, '1s') last_over_time((native_time_ns{series="exact"})[1s:1s]); + +-- Inner offsets are applied once when evaluating the subquery selector. +TQL EVAL (0, 0, '1s') last_over_time((native_time_ns{series="offset"} offset 1s)[1s:1s]); +TQL EVAL (0, 0, '1s') last_over_time((native_time_ns{series="offset"} offset -1s)[1s:1s]); + +DROP TABLE native_time_ns; + +-- An unrepresentable native lower bound must not discard its representable upper bound. +-- The upper filter must reach LastRow so the 1ms-future row cannot hide the eligible row. +CREATE TABLE native_time_ns_lower_overflow ( + ts TIMESTAMP(9) TIME INDEX, + series STRING PRIMARY KEY, + val DOUBLE, +); +INSERT INTO native_time_ns_lower_overflow VALUES + (-9223200000000000000, 'exact', 1), + (-9223199999999000000, 'exact', 2); + +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE native_time_ns_lower_overflow.__table_id\s*=\s*UInt32\(\d+\) native_time_ns_lower_overflow.__table_id=UInt32(REDACTED) +TQL EXPLAIN (0, 0, '1s', '2d') native_time_ns_lower_overflow{series="exact"} offset 106750d; + +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (Hash.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +-- SQLNESS REPLACE (flat_format.*) REDACTED +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +TQL ANALYZE VERBOSE (0, 0, '1s', '2d') native_time_ns_lower_overflow{series="exact"} offset 106750d; + +-- The representable upper bound selects only the exact row. +TQL EVAL (0, 0, '1s', '2d') native_time_ns_lower_overflow{series="exact"} offset 106750d; +ADMIN FLUSH_TABLE('native_time_ns_lower_overflow'); +TQL EVAL (0, 0, '1s', '2d') native_time_ns_lower_overflow{series="exact"} offset 106750d; +DROP TABLE native_time_ns_lower_overflow; + +-- Second precision is promoted before applying fractional-second offsets. +CREATE TABLE native_time_sec (ts TIMESTAMP(0) TIME INDEX, val DOUBLE); +INSERT INTO native_time_sec VALUES (0, 10), (1, 11), (2, 12); +TQL EVAL (1, 1, '1s', '1s') native_time_sec offset 500ms; +TQL EVAL (1, 1, '1s', '1s') native_time_sec offset -500ms; +DROP TABLE native_time_sec; diff --git a/tests/cases/standalone/common/promql/null_samples.result b/tests/cases/standalone/common/promql/null_samples.result new file mode 100644 index 00000000000..8e0f159fe19 --- /dev/null +++ b/tests/cases/standalone/common/promql/null_samples.result @@ -0,0 +1,252 @@ +-- NULL fields are missing samples, not zero-valued samples. +CREATE TABLE null_samples ( + ts TIMESTAMP(3) TIME INDEX, + host STRING PRIMARY KEY, + val DOUBLE, +); + +Affected Rows: 0 + +INSERT INTO null_samples VALUES + (0, 'a', 1.0), + (1000, 'a', NULL), + (2000, 'a', NULL), + (3000, 'a', 4.0), + (0, 'b', NULL), + (1000, 'b', NULL), + (2000, 'b', NULL), + (3000, 'b', NULL); + +Affected Rows: 8 + +-- At t=2 the trailing NULLs must not hide 1; at t=3 count must be 2. +-- At t=7 the left-open window is empty. Valid results must disappear. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (2, 7, '1s') count_over_time(null_samples{host="a"}[4s]); + ++---------------------+------------------------------------+------+ +| ts | prom_count_over_time(ts_range,val) | host | ++---------------------+------------------------------------+------+ +| 1970-01-01T00:00:02 | 1.0 | a | +| 1970-01-01T00:00:03 | 2.0 | a | +| 1970-01-01T00:00:04 | 1.0 | a | +| 1970-01-01T00:00:05 | 1.0 | a | +| 1970-01-01T00:00:06 | 1.0 | a | ++---------------------+------------------------------------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (2, 7, '1s') last_over_time(null_samples{host="a"}[4s]); + ++---------------------+-----------------------------------+------+ +| ts | prom_last_over_time(ts_range,val) | host | ++---------------------+-----------------------------------+------+ +| 1970-01-01T00:00:02 | 1.0 | a | +| 1970-01-01T00:00:03 | 4.0 | a | +| 1970-01-01T00:00:04 | 4.0 | a | +| 1970-01-01T00:00:05 | 4.0 | a | +| 1970-01-01T00:00:06 | 4.0 | a | ++---------------------+-----------------------------------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (2, 7, '1s') present_over_time(null_samples{host="a"}[4s]); + ++---------------------+--------------------------------------+------+ +| ts | prom_present_over_time(ts_range,val) | host | ++---------------------+--------------------------------------+------+ +| 1970-01-01T00:00:02 | 1.0 | a | +| 1970-01-01T00:00:03 | 1.0 | a | +| 1970-01-01T00:00:04 | 1.0 | a | +| 1970-01-01T00:00:05 | 1.0 | a | +| 1970-01-01T00:00:06 | 1.0 | a | ++---------------------+--------------------------------------+------+ + +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (2, 7, '1s') absent_over_time(null_samples{host="a"}[4s]); + ++---------------------+-------------------------------------+------+ +| ts | prom_absent_over_time(ts_range,val) | host | ++---------------------+-------------------------------------+------+ +| 1970-01-01T00:00:07 | 1.0 | a | ++---------------------+-------------------------------------+------+ + +-- All-NULL windows have no samples: only absent_over_time returns 1. +TQL EVAL (3, 3, '1s') count_over_time(null_samples{host="b"}[4s]); + +++ +++ + +TQL EVAL (3, 3, '1s') last_over_time(null_samples{host="b"}[4s]); + +++ +++ + +TQL EVAL (3, 3, '1s') present_over_time(null_samples{host="b"}[4s]); + +++ +++ + +TQL EVAL (3, 3, '1s') absent_over_time(null_samples{host="b"}[4s]); + ++---------------------+-------------------------------------+------+ +| ts | prom_absent_over_time(ts_range,val) | host | ++---------------------+-------------------------------------+------+ +| 1970-01-01T00:00:03 | 1.0 | b | ++---------------------+-------------------------------------+------+ + +-- Every function below sees the same two samples, 1.0 at 0s and 4.0 at 3s. +TQL EVAL (3, 3, '1s') rate(null_samples{host="a"}[4s]); + ++---------------------+----------------------------------------+------+ +| ts | prom_rate(ts_range,val,ts,Int64(4000)) | host | ++---------------------+----------------------------------------+------+ +| 1970-01-01T00:00:03 | 1.0 | a | ++---------------------+----------------------------------------+------+ + +TQL EVAL (3, 3, '1s') increase(null_samples{host="a"}[4s]); + ++---------------------+--------------------------------------------+------+ +| ts | prom_increase(ts_range,val,ts,Int64(4000)) | host | ++---------------------+--------------------------------------------+------+ +| 1970-01-01T00:00:03 | 4.0 | a | ++---------------------+--------------------------------------------+------+ + +TQL EVAL (3, 3, '1s') delta(null_samples{host="a"}[4s]); + ++---------------------+-----------------------------------------+------+ +| ts | prom_delta(ts_range,val,ts,Int64(4000)) | host | ++---------------------+-----------------------------------------+------+ +| 1970-01-01T00:00:03 | 4.0 | a | ++---------------------+-----------------------------------------+------+ + +TQL EVAL (3, 3, '1s') idelta(null_samples{host="a"}[4s]); + ++---------------------+---------------------------+------+ +| ts | prom_idelta(ts_range,val) | host | ++---------------------+---------------------------+------+ +| 1970-01-01T00:00:03 | 3.0 | a | ++---------------------+---------------------------+------+ + +TQL EVAL (3, 3, '1s') irate(null_samples{host="a"}[4s]); + ++---------------------+--------------------------+------+ +| ts | prom_irate(ts_range,val) | host | ++---------------------+--------------------------+------+ +| 1970-01-01T00:00:03 | 1.0 | a | ++---------------------+--------------------------+------+ + +TQL EVAL (3, 3, '1s') changes(null_samples{host="a"}[4s]); + ++---------------------+----------------------------+------+ +| ts | prom_changes(ts_range,val) | host | ++---------------------+----------------------------+------+ +| 1970-01-01T00:00:03 | 1.0 | a | ++---------------------+----------------------------+------+ + +TQL EVAL (3, 3, '1s') resets(null_samples{host="a"}[4s]); + ++---------------------+---------------------------+------+ +| ts | prom_resets(ts_range,val) | host | ++---------------------+---------------------------+------+ +| 1970-01-01T00:00:03 | 0.0 | a | ++---------------------+---------------------------+------+ + +TQL EVAL (3, 3, '1s') avg_over_time(null_samples{host="a"}[4s]); + ++---------------------+----------------------------------+------+ +| ts | prom_avg_over_time(ts_range,val) | host | ++---------------------+----------------------------------+------+ +| 1970-01-01T00:00:03 | 2.5 | a | ++---------------------+----------------------------------+------+ + +TQL EVAL (3, 3, '1s') stddev_over_time(null_samples{host="a"}[4s]); + ++---------------------+-------------------------------------+------+ +| ts | prom_stddev_over_time(ts_range,val) | host | ++---------------------+-------------------------------------+------+ +| 1970-01-01T00:00:03 | 1.5 | a | ++---------------------+-------------------------------------+------+ + +TQL EVAL (3, 3, '1s') stdvar_over_time(null_samples{host="a"}[4s]); + ++---------------------+-------------------------------------+------+ +| ts | prom_stdvar_over_time(ts_range,val) | host | ++---------------------+-------------------------------------+------+ +| 1970-01-01T00:00:03 | 2.25 | a | ++---------------------+-------------------------------------+------+ + +TQL EVAL (3, 3, '1s') quantile_over_time(0.5, null_samples{host="a"}[4s]); + ++---------------------+----------------------------------------------------+------+ +| ts | prom_quantile_over_time(ts_range,val,Float64(0.5)) | host | ++---------------------+----------------------------------------------------+------+ +| 1970-01-01T00:00:03 | 2.5 | a | ++---------------------+----------------------------------------------------+------+ + +-- A window whose only slot is NULL holds no sample, like a window with no row. +TQL EVAL (1, 1, '1s') rate(null_samples{host="a"}[1s]); + +++ +++ + +-- Prometheus returns an empty vector for a range without samples, not NaN. +TQL EVAL (1, 1, '1s') quantile_over_time(0.5, null_samples{host="a"}[1s]); + +++ +++ + +-- `b` never has a sample, so it must not reach the aggregation. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '1s') avg by (host) (rate(null_samples[4s])); + ++------+---------------------+---------------------------------------------+ +| host | ts | avg(prom_rate(ts_range,val,ts,Int64(4000))) | ++------+---------------------+---------------------------------------------+ +| a | 1970-01-01T00:00:03 | 1.0 | ++------+---------------------+---------------------------------------------+ + +DROP TABLE null_samples; + +Affected Rows: 0 + +CREATE TABLE multi_field ( + ts TIMESTAMP(3) TIME INDEX, + host STRING PRIMARY KEY, + f1 DOUBLE, + f2 DOUBLE +); + +Affected Rows: 0 + +INSERT INTO multi_field VALUES + (0, 'a', 1.0, 10.0), + (1000, 'a', 2.0, NULL), + (2000, 'a', 3.0, NULL), + (3000, 'a', 4.0, 40.0); + +Affected Rows: 4 + +-- f1 has two samples in this window and f2 only one, so f1 keeps its result while f2 is NULL. +-- Dropping the whole row would take f1's samples with it. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (1, 1, '1s') rate(multi_field[4s]); + ++---------------------+---------------------------------------+---------------------------------------+------+ +| ts | prom_rate(ts_range,f1,ts,Int64(4000)) | prom_rate(ts_range,f2,ts,Int64(4000)) | host | ++---------------------+---------------------------------------+---------------------------------------+------+ +| 1970-01-01T00:00:01 | 0.375 | | a | ++---------------------+---------------------------------------+---------------------------------------+------+ + +-- A selector emits the same shape: the row stays, the field without a sample is NULL. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (1, 1, '1s') multi_field; + ++---------------------+------+-----+----+ +| ts | host | f1 | f2 | ++---------------------+------+-----+----+ +| 1970-01-01T00:00:01 | a | 2.0 | | ++---------------------+------+-----+----+ + +DROP TABLE multi_field; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/promql/null_samples.sql b/tests/cases/standalone/common/promql/null_samples.sql new file mode 100644 index 00000000000..d0a02158a44 --- /dev/null +++ b/tests/cases/standalone/common/promql/null_samples.sql @@ -0,0 +1,92 @@ +-- NULL fields are missing samples, not zero-valued samples. +CREATE TABLE null_samples ( + ts TIMESTAMP(3) TIME INDEX, + host STRING PRIMARY KEY, + val DOUBLE, +); + +INSERT INTO null_samples VALUES + (0, 'a', 1.0), + (1000, 'a', NULL), + (2000, 'a', NULL), + (3000, 'a', 4.0), + (0, 'b', NULL), + (1000, 'b', NULL), + (2000, 'b', NULL), + (3000, 'b', NULL); + +-- At t=2 the trailing NULLs must not hide 1; at t=3 count must be 2. +-- At t=7 the left-open window is empty. Valid results must disappear. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (2, 7, '1s') count_over_time(null_samples{host="a"}[4s]); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (2, 7, '1s') last_over_time(null_samples{host="a"}[4s]); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (2, 7, '1s') present_over_time(null_samples{host="a"}[4s]); +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (2, 7, '1s') absent_over_time(null_samples{host="a"}[4s]); + +-- All-NULL windows have no samples: only absent_over_time returns 1. +TQL EVAL (3, 3, '1s') count_over_time(null_samples{host="b"}[4s]); +TQL EVAL (3, 3, '1s') last_over_time(null_samples{host="b"}[4s]); +TQL EVAL (3, 3, '1s') present_over_time(null_samples{host="b"}[4s]); +TQL EVAL (3, 3, '1s') absent_over_time(null_samples{host="b"}[4s]); + +-- Every function below sees the same two samples, 1.0 at 0s and 4.0 at 3s. +TQL EVAL (3, 3, '1s') rate(null_samples{host="a"}[4s]); + +TQL EVAL (3, 3, '1s') increase(null_samples{host="a"}[4s]); + +TQL EVAL (3, 3, '1s') delta(null_samples{host="a"}[4s]); + +TQL EVAL (3, 3, '1s') idelta(null_samples{host="a"}[4s]); + +TQL EVAL (3, 3, '1s') irate(null_samples{host="a"}[4s]); + +TQL EVAL (3, 3, '1s') changes(null_samples{host="a"}[4s]); + +TQL EVAL (3, 3, '1s') resets(null_samples{host="a"}[4s]); + +TQL EVAL (3, 3, '1s') avg_over_time(null_samples{host="a"}[4s]); + +TQL EVAL (3, 3, '1s') stddev_over_time(null_samples{host="a"}[4s]); + +TQL EVAL (3, 3, '1s') stdvar_over_time(null_samples{host="a"}[4s]); + +TQL EVAL (3, 3, '1s') quantile_over_time(0.5, null_samples{host="a"}[4s]); + +-- A window whose only slot is NULL holds no sample, like a window with no row. +TQL EVAL (1, 1, '1s') rate(null_samples{host="a"}[1s]); + +-- Prometheus returns an empty vector for a range without samples, not NaN. +TQL EVAL (1, 1, '1s') quantile_over_time(0.5, null_samples{host="a"}[1s]); + +-- `b` never has a sample, so it must not reach the aggregation. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (0, 15, '1s') avg by (host) (rate(null_samples[4s])); + +DROP TABLE null_samples; + +CREATE TABLE multi_field ( + ts TIMESTAMP(3) TIME INDEX, + host STRING PRIMARY KEY, + f1 DOUBLE, + f2 DOUBLE +); + +INSERT INTO multi_field VALUES + (0, 'a', 1.0, 10.0), + (1000, 'a', 2.0, NULL), + (2000, 'a', 3.0, NULL), + (3000, 'a', 4.0, 40.0); + +-- f1 has two samples in this window and f2 only one, so f1 keeps its result while f2 is NULL. +-- Dropping the whole row would take f1's samples with it. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (1, 1, '1s') rate(multi_field[4s]); + +-- A selector emits the same shape: the row stays, the field without a sample is NULL. +-- SQLNESS SORT_RESULT 3 1 +TQL EVAL (1, 1, '1s') multi_field; + +DROP TABLE multi_field; diff --git a/tests/cases/standalone/common/promql/precisions.result b/tests/cases/standalone/common/promql/precisions.result index e57f5b04ee4..4026a358322 100644 --- a/tests/cases/standalone/common/promql/precisions.result +++ b/tests/cases/standalone/common/promql/precisions.result @@ -133,10 +133,8 @@ TQL EVAL (0, 15, '5s') avg_over_time(host_sec{host="host1"}[5s]) + avg_over_time -- Verify that PromQL time predicates on non-millisecond time indexes are -- pushed into the scan as native timestamp range filters. --- Original instant selector filter is built on the millisecond alias: --- host = "host1" AND ts_ms >= -299999ms AND ts_ms <= 10000ms --- After pushing through `CAST(raw_ts AS Timestamp(ms)) AS ts` and applying --- DataFusion cast preimage, it becomes a native half-open range on raw_ts. +-- Instant selection compares raw timestamps before millisecond output conversion: +-- host = "host1" AND ts_us > -300000000us AND ts_us <= 10000000us. -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED @@ -152,17 +150,17 @@ TQL EXPLAIN (0, 10, '5s') host_micro{host="host1"}; | | PromInstantManipulate: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] | | | PromSeriesDivide: tags=["host"] | | | Sort: host_micro.host ASC NULLS FIRST, host_micro.ts ASC NULLS FIRST | -| | Projection: host_micro.val, host_micro.host, CAST(host_micro.ts AS Timestamp(ms)) AS ts | -| | Filter: host_micro.host = Utf8("host1") AND host_micro.ts >= TimestampMicrosecond(-299999999, None) AND host_micro.ts < TimestampMicrosecond(10001000, None) | -| | TableScan: host_micro, partial_filters=[host_micro.host = Utf8("host1"), host_micro.ts >= TimestampMicrosecond(-299999999, None), host_micro.ts < TimestampMicrosecond(10001000, None)] | +| | Projection: host_micro.val, host_micro.host, host_micro.ts | +| | Filter: host_micro.host = Utf8("host1") AND host_micro.ts > TimestampMicrosecond(-300000000, None) AND host_micro.ts <= TimestampMicrosecond(10000000, None) | +| | TableScan: host_micro, partial_filters=[host_micro.host = Utf8("host1"), host_micro.ts > TimestampMicrosecond(-300000000, None), host_micro.ts <= TimestampMicrosecond(10000000, None)] | | | ]] | | physical_plan | CooperativeExec | | | MergeScanExec: REDACTED | | | +---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ --- The same instant-selector cast-preimage path should work for nanosecond indexes. --- Expected native bounds: ts_ns >= -299999999999ns AND ts_ns < 10001000000ns. +-- The same exclusive-lower, inclusive-upper window applies to nanosecond indexes. +-- Expected native bounds: ts_ns > -300000000000ns AND ts_ns <= 10000000000ns. -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED @@ -177,9 +175,9 @@ TQL EXPLAIN (0, 10, '5s') host_nano{host="host1"}; | | PromInstantManipulate: range=[0..10000], lookback=[300000], interval=[5000], time index=[ts] | | | PromSeriesDivide: tags=["host"] | | | Sort: host_nano.host ASC NULLS FIRST, host_nano.ts ASC NULLS FIRST | -| | Projection: host_nano.val, host_nano.host, CAST(host_nano.ts AS Timestamp(ms)) AS ts | -| | Filter: host_nano.host = Utf8("host1") AND host_nano.ts >= TimestampNanosecond(-299999999999, None) AND host_nano.ts < TimestampNanosecond(10001000000, None) | -| | TableScan: host_nano, partial_filters=[host_nano.host = Utf8("host1"), host_nano.ts >= TimestampNanosecond(-299999999999, None), host_nano.ts < TimestampNanosecond(10001000000, None)] | +| | Projection: host_nano.val, host_nano.host, host_nano.ts | +| | Filter: host_nano.host = Utf8("host1") AND host_nano.ts > TimestampNanosecond(-300000000000, None) AND host_nano.ts <= TimestampNanosecond(10000000000, None) | +| | TableScan: host_nano, partial_filters=[host_nano.host = Utf8("host1"), host_nano.ts > TimestampNanosecond(-300000000000, None), host_nano.ts <= TimestampNanosecond(10000000000, None)] | | | ]] | | physical_plan | CooperativeExec | | | MergeScanExec: REDACTED @@ -187,8 +185,8 @@ TQL EXPLAIN (0, 10, '5s') host_nano{host="host1"}; +---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ -- Range selectors use their range window instead of the default lookback. --- Original range selector filter for [5s]: --- host = "host1" AND ts_ms >= -4999ms AND ts_ms <= 10000ms +-- Native range selector filter for [5s]: +-- host = "host1" AND ts_us > -5000000us AND ts_us <= 10000000us -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED @@ -206,17 +204,17 @@ TQL EXPLAIN (0, 10, '5s') avg_over_time(host_micro{host="host1"}[5s]); | | PromSeriesNormalize: offset=[0], time index=[ts], filter NaN: [true] | | | PromSeriesDivide: tags=["host"] | | | Sort: host_micro.host ASC NULLS FIRST, host_micro.ts ASC NULLS FIRST | -| | Projection: host_micro.val, host_micro.host, CAST(host_micro.ts AS Timestamp(ms)) AS ts | -| | Filter: host_micro.host = Utf8("host1") AND host_micro.ts >= TimestampMicrosecond(-4999999, None) AND host_micro.ts < TimestampMicrosecond(10001000, None) | -| | TableScan: host_micro, partial_filters=[host_micro.host = Utf8("host1"), host_micro.ts >= TimestampMicrosecond(-4999999, None), host_micro.ts < TimestampMicrosecond(10001000, None)] | +| | Projection: host_micro.val, host_micro.host, host_micro.ts | +| | Filter: host_micro.host = Utf8("host1") AND host_micro.ts > TimestampMicrosecond(-5000000, None) AND host_micro.ts <= TimestampMicrosecond(10000000, None) | +| | TableScan: host_micro, partial_filters=[host_micro.host = Utf8("host1"), host_micro.ts > TimestampMicrosecond(-5000000, None), host_micro.ts <= TimestampMicrosecond(10000000, None)] | | | ]] | | physical_plan | CooperativeExec | | | MergeScanExec: REDACTED | | | +---------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ --- The same range-selector cast-preimage path should work for nanosecond indexes. --- Expected native bounds: ts_ns >= -4999999999ns AND ts_ns < 10001000000ns. +-- Range membership also retains nanosecond precision. +-- Expected native bounds: ts_ns > -5000000000ns AND ts_ns <= 10000000000ns. -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED @@ -234,9 +232,9 @@ TQL EXPLAIN (0, 10, '5s') avg_over_time(host_nano{host="host1"}[5s]); | | PromSeriesNormalize: offset=[0], time index=[ts], filter NaN: [true] | | | PromSeriesDivide: tags=["host"] | | | Sort: host_nano.host ASC NULLS FIRST, host_nano.ts ASC NULLS FIRST | -| | Projection: host_nano.val, host_nano.host, CAST(host_nano.ts AS Timestamp(ms)) AS ts | -| | Filter: host_nano.host = Utf8("host1") AND host_nano.ts >= TimestampNanosecond(-4999999999, None) AND host_nano.ts < TimestampNanosecond(10001000000, None) | -| | TableScan: host_nano, partial_filters=[host_nano.host = Utf8("host1"), host_nano.ts >= TimestampNanosecond(-4999999999, None), host_nano.ts < TimestampNanosecond(10001000000, None)] | +| | Projection: host_nano.val, host_nano.host, host_nano.ts | +| | Filter: host_nano.host = Utf8("host1") AND host_nano.ts > TimestampNanosecond(-5000000000, None) AND host_nano.ts <= TimestampNanosecond(10000000000, None) | +| | TableScan: host_nano, partial_filters=[host_nano.host = Utf8("host1"), host_nano.ts > TimestampNanosecond(-5000000000, None), host_nano.ts <= TimestampNanosecond(10000000000, None)] | | | ]] | | physical_plan | CooperativeExec | | | MergeScanExec: REDACTED diff --git a/tests/cases/standalone/common/promql/precisions.sql b/tests/cases/standalone/common/promql/precisions.sql index 01e6cae9fec..8f7fec37a86 100644 --- a/tests/cases/standalone/common/promql/precisions.sql +++ b/tests/cases/standalone/common/promql/precisions.sql @@ -68,10 +68,8 @@ TQL EVAL (0, 15, '5s') avg_over_time(host_sec{host="host1"}[5s]) + avg_over_time -- Verify that PromQL time predicates on non-millisecond time indexes are -- pushed into the scan as native timestamp range filters. --- Original instant selector filter is built on the millisecond alias: --- host = "host1" AND ts_ms >= -299999ms AND ts_ms <= 10000ms --- After pushing through `CAST(raw_ts AS Timestamp(ms)) AS ts` and applying --- DataFusion cast preimage, it becomes a native half-open range on raw_ts. +-- Instant selection compares raw timestamps before millisecond output conversion: +-- host = "host1" AND ts_us > -300000000us AND ts_us <= 10000000us. -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED @@ -80,8 +78,8 @@ TQL EVAL (0, 15, '5s') avg_over_time(host_sec{host="host1"}[5s]) + avg_over_time -- SQLNESS REPLACE host_nano.__table_id\s*=\s*UInt32\(\d+\) host_nano.__table_id=UInt32(REDACTED) TQL EXPLAIN (0, 10, '5s') host_micro{host="host1"}; --- The same instant-selector cast-preimage path should work for nanosecond indexes. --- Expected native bounds: ts_ns >= -299999999999ns AND ts_ns < 10001000000ns. +-- The same exclusive-lower, inclusive-upper window applies to nanosecond indexes. +-- Expected native bounds: ts_ns > -300000000000ns AND ts_ns <= 10000000000ns. -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED @@ -90,8 +88,8 @@ TQL EXPLAIN (0, 10, '5s') host_micro{host="host1"}; TQL EXPLAIN (0, 10, '5s') host_nano{host="host1"}; -- Range selectors use their range window instead of the default lookback. --- Original range selector filter for [5s]: --- host = "host1" AND ts_ms >= -4999ms AND ts_ms <= 10000ms +-- Native range selector filter for [5s]: +-- host = "host1" AND ts_us > -5000000us AND ts_us <= 10000000us -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED @@ -99,8 +97,8 @@ TQL EXPLAIN (0, 10, '5s') host_nano{host="host1"}; -- SQLNESS REPLACE host_micro.__table_id\s*=\s*UInt32\(\d+\) host_micro.__table_id=UInt32(REDACTED) TQL EXPLAIN (0, 10, '5s') avg_over_time(host_micro{host="host1"}[5s]); --- The same range-selector cast-preimage path should work for nanosecond indexes. --- Expected native bounds: ts_ns >= -4999999999ns AND ts_ns < 10001000000ns. +-- Range membership also retains nanosecond precision. +-- Expected native bounds: ts_ns > -5000000000ns AND ts_ns <= 10000000000ns. -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED -- SQLNESS REPLACE (peers.*) REDACTED -- SQLNESS REPLACE (Hash.*) REDACTED diff --git a/tests/cases/standalone/common/promql/set_operation.result b/tests/cases/standalone/common/promql/set_operation.result index 50a0b6cf3b6..2c56567f08d 100644 --- a/tests/cases/standalone/common/promql/set_operation.result +++ b/tests/cases/standalone/common/promql/set_operation.result @@ -816,12 +816,12 @@ tql eval(1000, 2000, '300s') sum by (src, src_pod, src_namespace, src_node, dest +---------------------+-----------+---------+----------+---------+----------+---------------+----------+---------+----------------------------------------------------------------------------------------------+ | greptime_timestamp | az | cloud | dest | region | src | src_namespace | src_node | src_pod | sum(prom_increase(greptime_timestamp_range,greptime_value,greptime_timestamp,Int64(900000))) | +---------------------+-----------+---------+----------+---------+----------+---------------+----------+---------+----------------------------------------------------------------------------------------------+ -| 1970-01-01T00:21:40 | us-west-6 | cloud-1 | 10.0.0.2 | us-west | 10.0.0.1 | namespace-1 | node-1 | pod-1 | 2500.0 | +| 1970-01-01T00:21:40 | us-west-6 | cloud-1 | 10.0.0.2 | us-west | 10.0.0.1 | namespace-1 | node-1 | pod-1 | 2000.0 | | 1970-01-01T00:21:40 | us-west-6 | cloud-1 | 10.0.0.3 | us-west | 10.0.0.1 | namespace-1 | node-2 | pod-2 | 2000.0 | -| 1970-01-01T00:21:40 | us-west-6 | cloud-2 | 10.0.0.5 | us-west | 10.0.0.4 | namespace-2 | node-3 | pod-3 | 2300.0 | -| 1970-01-01T00:26:40 | us-west-6 | cloud-1 | 10.0.0.2 | us-west | 10.0.0.1 | namespace-1 | node-1 | pod-1 | 2500.0 | +| 1970-01-01T00:21:40 | us-west-6 | cloud-2 | 10.0.0.5 | us-west | 10.0.0.4 | namespace-2 | node-3 | pod-3 | 2000.0 | +| 1970-01-01T00:26:40 | us-west-6 | cloud-1 | 10.0.0.2 | us-west | 10.0.0.1 | namespace-1 | node-1 | pod-1 | 2000.0 | | 1970-01-01T00:26:40 | us-west-6 | cloud-1 | 10.0.0.3 | us-west | 10.0.0.1 | namespace-1 | node-2 | pod-2 | 2000.0 | -| 1970-01-01T00:26:40 | us-west-6 | cloud-2 | 10.0.0.5 | us-west | 10.0.0.4 | namespace-2 | node-3 | pod-3 | 2300.0 | +| 1970-01-01T00:26:40 | us-west-6 | cloud-2 | 10.0.0.5 | us-west | 10.0.0.4 | namespace-2 | node-3 | pod-3 | 2000.0 | +---------------------+-----------+---------+----------+---------+----------+---------------+----------+---------+----------------------------------------------------------------------------------------------+ DROP TABLE node_network_transmit_bytes_total; diff --git a/tests/cases/standalone/common/session_skip_wal.result b/tests/cases/standalone/common/session_skip_wal.result new file mode 100644 index 00000000000..4343e3cd334 --- /dev/null +++ b/tests/cases/standalone/common/session_skip_wal.result @@ -0,0 +1,224 @@ +CREATE TABLE session_skip_wal(host STRING, ts TIMESTAMP TIME INDEX); + +Affected Rows: 0 + +-- MYSQL: SET persists across statements on the same connection. +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = true; + +affected_rows: 0 + +-- Invalid SET must leave the enabled policy unchanged. +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = 'invalid'; + +Failed to execute query, err: MySqlError { ERROR 1235 (42000): (Unsupported): Not supported: Invalid skip_wal value "invalid": expected true or false } + +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = 1; + +Failed to execute query, err: MySqlError { ERROR 1235 (42000): (Unsupported): Not supported: SET skip_wal requires true or false } + +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = NULL; + +Failed to execute query, err: MySqlError { ERROR 1235 (42000): (Unsupported): Not supported: SET skip_wal requires true or false } + +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = false, true; + +Failed to execute query, err: MySqlError { ERROR 1235 (42000): (Unsupported): Not supported: SET skip_wal requires exactly one boolean value } + +-- SQLNESS PROTOCOL MYSQL +INSERT INTO session_skip_wal VALUES ('skipped', 1000); + +affected_rows: 1 + +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM session_skip_wal ORDER BY ts; + ++---------+---------------------+ +| host | ts | ++---------+---------------------+ +| skipped | 1970-01-01 00:00:01 | ++---------+---------------------+ + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM session_skip_wal ORDER BY ts; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = false; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = 'true'; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = 'false'; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +INSERT INTO session_skip_wal VALUES ('persisted', 2000); + +affected_rows: 1 + +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM session_skip_wal ORDER BY ts; + ++-----------+---------------------+ +| host | ts | ++-----------+---------------------+ +| persisted | 1970-01-01 00:00:02 | ++-----------+---------------------+ + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM session_skip_wal ORDER BY ts; + ++-----------+---------------------+ +| host | ts | ++-----------+---------------------+ +| persisted | 1970-01-01 00:00:02 | ++-----------+---------------------+ + +TRUNCATE TABLE session_skip_wal; + +Affected Rows: 0 + +-- POSTGRES: SET persists across statements on the same connection. +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = true; + +Affected Rows: 0 + +-- Invalid SET must leave the enabled policy unchanged. +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = 'invalid'; + +Failed to execute query, encountered: Error { kind: Db, cause: Some(DbError { severity: "ERROR", parsed_severity: None, code: SqlState(E0A000), message: "Not supported: Invalid skip_wal value \"invalid\": expected true or false", detail: None, hint: None, position: None, where_: None, schema: None, table: None, column: None, datatype: None, constraint: None, file: None, line: None, routine: None }) } + +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = 1; + +Failed to execute query, encountered: Error { kind: Db, cause: Some(DbError { severity: "ERROR", parsed_severity: None, code: SqlState(E0A000), message: "Not supported: SET skip_wal requires true or false", detail: None, hint: None, position: None, where_: None, schema: None, table: None, column: None, datatype: None, constraint: None, file: None, line: None, routine: None }) } + +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = NULL; + +Failed to execute query, encountered: Error { kind: Db, cause: Some(DbError { severity: "ERROR", parsed_severity: None, code: SqlState(E0A000), message: "Not supported: SET skip_wal requires true or false", detail: None, hint: None, position: None, where_: None, schema: None, table: None, column: None, datatype: None, constraint: None, file: None, line: None, routine: None }) } + +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = false, true; + +Failed to execute query, encountered: Error { kind: Db, cause: Some(DbError { severity: "ERROR", parsed_severity: None, code: SqlState(E0A000), message: "Not supported: SET skip_wal requires exactly one boolean value", detail: None, hint: None, position: None, where_: None, schema: None, table: None, column: None, datatype: None, constraint: None, file: None, line: None, routine: None }) } + +-- SQLNESS PROTOCOL POSTGRES +INSERT INTO session_skip_wal VALUES ('skipped', 1000); + +Affected Rows: 1 + +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM session_skip_wal ORDER BY ts; + ++---------+----------------------------+ +| host | ts | ++---------+----------------------------+ +| skipped | 1970-01-01 00:00:01.000000 | ++---------+----------------------------+ + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM session_skip_wal ORDER BY ts; + ++------+----+ +| host | ts | ++------+----+ ++------+----+ + +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = false; + +Affected Rows: 0 + +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = 'true'; + +Affected Rows: 0 + +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = 'false'; + +Affected Rows: 0 + +-- SQLNESS PROTOCOL POSTGRES +INSERT INTO session_skip_wal VALUES ('persisted', 2000); + +Affected Rows: 1 + +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM session_skip_wal ORDER BY ts; + ++-----------+----------------------------+ +| host | ts | ++-----------+----------------------------+ +| persisted | 1970-01-01 00:00:02.000000 | ++-----------+----------------------------+ + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM session_skip_wal ORDER BY ts; + ++-----------+----------------------------+ +| host | ts | ++-----------+----------------------------+ +| persisted | 1970-01-01 00:00:02.000000 | ++-----------+----------------------------+ + +TRUNCATE TABLE session_skip_wal; + +Affected Rows: 0 + +DROP TABLE session_skip_wal; + +Affected Rows: 0 + +-- The default gRPC protocol has request-scoped contexts, not a persistent SQL session. +SET skip_wal = true; + +Affected Rows: 0 + +SET skip_wal = false; + +Affected Rows: 0 + +SET skip_wal = 'true'; + +Affected Rows: 0 + +SET skip_wal = 'false'; + +Affected Rows: 0 + +SET skip_wal = 'invalid'; + +Error: 1001(Unsupported), Not supported: Invalid skip_wal value "invalid": expected true or false + +SET skip_wal = 1; + +Error: 1001(Unsupported), Not supported: SET skip_wal requires true or false + +SET skip_wal = NULL; + +Error: 1001(Unsupported), Not supported: SET skip_wal requires true or false + +SET skip_wal = false, true; + +Error: 1001(Unsupported), Not supported: SET skip_wal requires exactly one boolean value + diff --git a/tests/cases/standalone/common/session_skip_wal.sql b/tests/cases/standalone/common/session_skip_wal.sql new file mode 100644 index 00000000000..8be063253bc --- /dev/null +++ b/tests/cases/standalone/common/session_skip_wal.sql @@ -0,0 +1,118 @@ +CREATE TABLE session_skip_wal(host STRING, ts TIMESTAMP TIME INDEX); + +-- MYSQL: SET persists across statements on the same connection. +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = true; + +-- Invalid SET must leave the enabled policy unchanged. + +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = 'invalid'; + +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = 1; + +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = NULL; + +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = false, true; + +-- SQLNESS PROTOCOL MYSQL +INSERT INTO session_skip_wal VALUES ('skipped', 1000); + +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM session_skip_wal ORDER BY ts; + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM session_skip_wal ORDER BY ts; + +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = false; + +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = 'true'; + +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = 'false'; + +-- SQLNESS PROTOCOL MYSQL +INSERT INTO session_skip_wal VALUES ('persisted', 2000); + +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM session_skip_wal ORDER BY ts; + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM session_skip_wal ORDER BY ts; + +TRUNCATE TABLE session_skip_wal; + +-- POSTGRES: SET persists across statements on the same connection. +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = true; + +-- Invalid SET must leave the enabled policy unchanged. + +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = 'invalid'; + +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = 1; + +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = NULL; + +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = false, true; + +-- SQLNESS PROTOCOL POSTGRES +INSERT INTO session_skip_wal VALUES ('skipped', 1000); + +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM session_skip_wal ORDER BY ts; + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM session_skip_wal ORDER BY ts; + +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = false; + +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = 'true'; + +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = 'false'; + +-- SQLNESS PROTOCOL POSTGRES +INSERT INTO session_skip_wal VALUES ('persisted', 2000); + +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM session_skip_wal ORDER BY ts; + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM session_skip_wal ORDER BY ts; + +TRUNCATE TABLE session_skip_wal; + +DROP TABLE session_skip_wal; + +-- The default gRPC protocol has request-scoped contexts, not a persistent SQL session. +SET skip_wal = true; + +SET skip_wal = false; + +SET skip_wal = 'true'; + +SET skip_wal = 'false'; + +SET skip_wal = 'invalid'; + +SET skip_wal = 1; + +SET skip_wal = NULL; + +SET skip_wal = false, true; diff --git a/tests/cases/standalone/common/show/show_create.result b/tests/cases/standalone/common/show/show_create.result index 18879ed95e8..ba2c39386d6 100644 --- a/tests/cases/standalone/common/show/show_create.result +++ b/tests/cases/standalone/common/show/show_create.result @@ -180,6 +180,310 @@ DROP TABLE alter_hint_table; Affected Rows: 0 +CREATE TABLE partition_num_hint_table ( + host STRING PRIMARY KEY, + ts TIMESTAMP TIME INDEX +) +PARTITION ON COLUMNS (host) ( + host < 'm', + host >= 'm' +) +WITH ('repartition.partition.num.hint' = ' 8 '); + +Affected Rows: 0 + +SHOW CREATE TABLE partition_num_hint_table; + ++--------------------------+---------------------------------------------------------+ +| Table | Create Table | ++--------------------------+---------------------------------------------------------+ +| partition_num_hint_table | CREATE TABLE IF NOT EXISTS "partition_num_hint_table" ( | +| | "host" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | PARTITION ON COLUMNS ("host") ( | +| | host < 'm', | +| | host >= 'm' | +| | ) | +| | ENGINE=mito | +| | WITH( | +| | 'repartition.partition.num.hint' = '8' | +| | ) | ++--------------------------+---------------------------------------------------------+ + +ALTER TABLE partition_num_hint_table SET 'repartition.partition.num.hint' = 16; + +Affected Rows: 0 + +SHOW CREATE TABLE partition_num_hint_table; + ++--------------------------+---------------------------------------------------------+ +| Table | Create Table | ++--------------------------+---------------------------------------------------------+ +| partition_num_hint_table | CREATE TABLE IF NOT EXISTS "partition_num_hint_table" ( | +| | "host" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | PARTITION ON COLUMNS ("host") ( | +| | host < 'm', | +| | host >= 'm' | +| | ) | +| | ENGINE=mito | +| | WITH( | +| | 'repartition.partition.num.hint' = '16' | +| | ) | ++--------------------------+---------------------------------------------------------+ + +ALTER TABLE partition_num_hint_table SET 'repartition.partition.num.hint' = '0'; + +Error: 1004(InvalidArguments), Invalid alter table(partition_num_hint_table) request: repartition.partition.num.hint expects a positive integer within u32 range, got `0` + +ALTER TABLE partition_num_hint_table UNSET 'repartition.partition.num.hint'; + +Affected Rows: 0 + +ALTER TABLE partition_num_hint_table UNSET 'repartition.partition.num.hint'; + +Affected Rows: 0 + +SHOW CREATE TABLE partition_num_hint_table; + ++--------------------------+---------------------------------------------------------+ +| Table | Create Table | ++--------------------------+---------------------------------------------------------+ +| partition_num_hint_table | CREATE TABLE IF NOT EXISTS "partition_num_hint_table" ( | +| | "host" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | PARTITION ON COLUMNS ("host") ( | +| | host < 'm', | +| | host >= 'm' | +| | ) | +| | ENGINE=mito | +| | | ++--------------------------+---------------------------------------------------------+ + +DROP TABLE partition_num_hint_table; + +Affected Rows: 0 + +CREATE TABLE partition_num_hint_table ( + host STRING PRIMARY KEY, + ts TIMESTAMP TIME INDEX +) +WITH ('repartition.partition.num.hint' = '10'); + +Affected Rows: 0 + +SHOW CREATE TABLE partition_num_hint_table; + ++--------------------------+---------------------------------------------------------+ +| Table | Create Table | ++--------------------------+---------------------------------------------------------+ +| partition_num_hint_table | CREATE TABLE IF NOT EXISTS "partition_num_hint_table" ( | +| | "host" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | 'repartition.partition.num.hint' = '10' | +| | ) | ++--------------------------+---------------------------------------------------------+ + +DROP TABLE partition_num_hint_table; + +Affected Rows: 0 + +CREATE TABLE partition_num_hint_table ( + host STRING PRIMARY KEY, + ts TIMESTAMP TIME INDEX +) +WITH ('repartition.partition.num.hint' = 10); + +Affected Rows: 0 + +SHOW CREATE TABLE partition_num_hint_table; + ++--------------------------+---------------------------------------------------------+ +| Table | Create Table | ++--------------------------+---------------------------------------------------------+ +| partition_num_hint_table | CREATE TABLE IF NOT EXISTS "partition_num_hint_table" ( | +| | "host" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | 'repartition.partition.num.hint' = '10' | +| | ) | ++--------------------------+---------------------------------------------------------+ + +DROP TABLE partition_num_hint_table; + +Affected Rows: 0 + +CREATE TABLE combined_hint_table ( + host STRING PRIMARY KEY, + `service` STRING, + ts TIMESTAMP TIME INDEX +); + +Affected Rows: 0 + +ALTER TABLE combined_hint_table SET + 'repartition.column.hint' = 'host', + 'repartition.partition.num.hint' = 10; + +Affected Rows: 0 + +SHOW CREATE TABLE combined_hint_table; + ++---------------------+----------------------------------------------------+ +| Table | Create Table | ++---------------------+----------------------------------------------------+ +| combined_hint_table | CREATE TABLE IF NOT EXISTS "combined_hint_table" ( | +| | "host" STRING NULL, | +| | "service" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | 'repartition.column.hint' = 'host', | +| | 'repartition.partition.num.hint' = '10' | +| | ) | ++---------------------+----------------------------------------------------+ + +ALTER TABLE combined_hint_table SET + 'repartition.column.hint' = 'service', + 'repartition.partition.num.hint' = 0; + +Error: 1004(InvalidArguments), Invalid alter table(combined_hint_table) request: repartition.partition.num.hint expects a positive integer within u32 range, got `0` + +SHOW CREATE TABLE combined_hint_table; + ++---------------------+----------------------------------------------------+ +| Table | Create Table | ++---------------------+----------------------------------------------------+ +| combined_hint_table | CREATE TABLE IF NOT EXISTS "combined_hint_table" ( | +| | "host" STRING NULL, | +| | "service" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | 'repartition.column.hint' = 'host', | +| | 'repartition.partition.num.hint' = '10' | +| | ) | ++---------------------+----------------------------------------------------+ + +ALTER TABLE combined_hint_table SET + 'repartition.partition.num.hint' = 20, + 'repartition.column.hint' = 'missing'; + +Error: 4002(TableColumnNotFound), Column missing not exists in table combined_hint_table + +SHOW CREATE TABLE combined_hint_table; + ++---------------------+----------------------------------------------------+ +| Table | Create Table | ++---------------------+----------------------------------------------------+ +| combined_hint_table | CREATE TABLE IF NOT EXISTS "combined_hint_table" ( | +| | "host" STRING NULL, | +| | "service" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | 'repartition.column.hint' = 'host', | +| | 'repartition.partition.num.hint' = '10' | +| | ) | ++---------------------+----------------------------------------------------+ + +ALTER TABLE combined_hint_table SET + 'repartition.partition.num.hint' = 20, + 'repartition.column.hint' = 'service'; + +Affected Rows: 0 + +SHOW CREATE TABLE combined_hint_table; + ++---------------------+----------------------------------------------------+ +| Table | Create Table | ++---------------------+----------------------------------------------------+ +| combined_hint_table | CREATE TABLE IF NOT EXISTS "combined_hint_table" ( | +| | "host" STRING NULL, | +| | "service" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | WITH( | +| | 'repartition.column.hint' = 'service', | +| | 'repartition.partition.num.hint' = '20' | +| | ) | ++---------------------+----------------------------------------------------+ + +ALTER TABLE combined_hint_table SET + 'repartition.partition.num.hint' = 30, + 'repartition.partition.num.hint' = 40; + +Error: 1004(InvalidArguments), Invalid table option request: duplicate repartition hint keys + +ALTER TABLE combined_hint_table SET + 'repartition.partition.num.hint' = 30, + 'ttl' = '7d'; + +Error: 1004(InvalidArguments), Invalid table option request: repartition hints must be altered separately from other table options + +ALTER TABLE combined_hint_table UNSET + 'repartition.column.hint', + 'repartition.partition.num.hint'; + +Affected Rows: 0 + +SHOW CREATE TABLE combined_hint_table; + ++---------------------+----------------------------------------------------+ +| Table | Create Table | ++---------------------+----------------------------------------------------+ +| combined_hint_table | CREATE TABLE IF NOT EXISTS "combined_hint_table" ( | +| | "host" STRING NULL, | +| | "service" STRING NULL, | +| | "ts" TIMESTAMP(3) NOT NULL, | +| | TIME INDEX ("ts"), | +| | PRIMARY KEY ("host") | +| | ) | +| | | +| | ENGINE=mito | +| | | ++---------------------+----------------------------------------------------+ + +DROP TABLE combined_hint_table; + +Affected Rows: 0 + CREATE TABLE not_supported_table_storage_option ( `id` INT UNSIGNED, host STRING, diff --git a/tests/cases/standalone/common/show/show_create.sql b/tests/cases/standalone/common/show/show_create.sql index fd8b1c16ea2..5e86273091a 100644 --- a/tests/cases/standalone/common/show/show_create.sql +++ b/tests/cases/standalone/common/show/show_create.sql @@ -61,6 +61,98 @@ SHOW CREATE TABLE alter_hint_table; DROP TABLE alter_hint_table; +CREATE TABLE partition_num_hint_table ( + host STRING PRIMARY KEY, + ts TIMESTAMP TIME INDEX +) +PARTITION ON COLUMNS (host) ( + host < 'm', + host >= 'm' +) +WITH ('repartition.partition.num.hint' = ' 8 '); + +SHOW CREATE TABLE partition_num_hint_table; + +ALTER TABLE partition_num_hint_table SET 'repartition.partition.num.hint' = 16; + +SHOW CREATE TABLE partition_num_hint_table; + +ALTER TABLE partition_num_hint_table SET 'repartition.partition.num.hint' = '0'; + +ALTER TABLE partition_num_hint_table UNSET 'repartition.partition.num.hint'; + +ALTER TABLE partition_num_hint_table UNSET 'repartition.partition.num.hint'; + +SHOW CREATE TABLE partition_num_hint_table; + +DROP TABLE partition_num_hint_table; + +CREATE TABLE partition_num_hint_table ( + host STRING PRIMARY KEY, + ts TIMESTAMP TIME INDEX +) +WITH ('repartition.partition.num.hint' = '10'); + +SHOW CREATE TABLE partition_num_hint_table; + +DROP TABLE partition_num_hint_table; + +CREATE TABLE partition_num_hint_table ( + host STRING PRIMARY KEY, + ts TIMESTAMP TIME INDEX +) +WITH ('repartition.partition.num.hint' = 10); + +SHOW CREATE TABLE partition_num_hint_table; + +DROP TABLE partition_num_hint_table; + +CREATE TABLE combined_hint_table ( + host STRING PRIMARY KEY, + `service` STRING, + ts TIMESTAMP TIME INDEX +); + +ALTER TABLE combined_hint_table SET + 'repartition.column.hint' = 'host', + 'repartition.partition.num.hint' = 10; + +SHOW CREATE TABLE combined_hint_table; + +ALTER TABLE combined_hint_table SET + 'repartition.column.hint' = 'service', + 'repartition.partition.num.hint' = 0; + +SHOW CREATE TABLE combined_hint_table; + +ALTER TABLE combined_hint_table SET + 'repartition.partition.num.hint' = 20, + 'repartition.column.hint' = 'missing'; + +SHOW CREATE TABLE combined_hint_table; + +ALTER TABLE combined_hint_table SET + 'repartition.partition.num.hint' = 20, + 'repartition.column.hint' = 'service'; + +SHOW CREATE TABLE combined_hint_table; + +ALTER TABLE combined_hint_table SET + 'repartition.partition.num.hint' = 30, + 'repartition.partition.num.hint' = 40; + +ALTER TABLE combined_hint_table SET + 'repartition.partition.num.hint' = 30, + 'ttl' = '7d'; + +ALTER TABLE combined_hint_table UNSET + 'repartition.column.hint', + 'repartition.partition.num.hint'; + +SHOW CREATE TABLE combined_hint_table; + +DROP TABLE combined_hint_table; + CREATE TABLE not_supported_table_storage_option ( `id` INT UNSIGNED, host STRING, diff --git a/tests/cases/standalone/common/tql-explain-analyze/explain.result b/tests/cases/standalone/common/tql-explain-analyze/explain.result index 2d6eab49da3..6b03513a116 100644 --- a/tests/cases/standalone/common/tql-explain-analyze/explain.result +++ b/tests/cases/standalone/common/tql-explain-analyze/explain.result @@ -391,9 +391,9 @@ TQL EXPLAIN (0, 10, '5s') test_nano; | | PromInstantManipulate: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] | | | PromSeriesDivide: tags=["k"] | | | Sort: test_nano.k ASC NULLS FIRST, test_nano.j ASC NULLS FIRST | -| | Projection: test_nano.i, test_nano.k, CAST(test_nano.j AS Timestamp(ms)) AS j | -| | Filter: test_nano.j >= TimestampNanosecond(-299999999999, None) AND test_nano.j < TimestampNanosecond(10001000000, None) | -| | TableScan: test_nano, partial_filters=[test_nano.j >= TimestampNanosecond(-299999999999, None), test_nano.j < TimestampNanosecond(10001000000, None)] | +| | Projection: test_nano.i, test_nano.k, test_nano.j | +| | Filter: test_nano.j > TimestampNanosecond(-300000000000, None) AND test_nano.j <= TimestampNanosecond(10000000000, None) | +| | TableScan: test_nano, partial_filters=[test_nano.j > TimestampNanosecond(-300000000000, None), test_nano.j <= TimestampNanosecond(10000000000, None)] | | | ]] | | physical_plan | CooperativeExec | | | MergeScanExec: REDACTED @@ -415,8 +415,8 @@ TQL EXPLAIN VERBOSE (0, 10, '5s') test_nano; | initial_logical_plan_| PromInstantManipulate: range=[0..10000], lookback=[300000], interval=[5000], time index=[j]_| |_|_PromSeriesDivide: tags=["k"]_| |_|_Sort: test_nano.k ASC NULLS FIRST, test_nano.j ASC NULLS FIRST_| -|_|_Filter: test_nano.j >= TimestampMillisecond(-299999, None) AND test_nano.j <= TimestampMillisecond(10000, None)_| -|_|_Projection: test_nano.i, test_nano.k, CAST(test_nano.j AS Timestamp(ms)) AS j_| +|_|_Filter: test_nano.j > TimestampNanosecond(-300000000000, None) AND test_nano.j <= TimestampNanosecond(10000000000, None)_| +|_|_Projection: test_nano.i, test_nano.k, test_nano.j_| |_|_TableScan: test_nano_| | logical_plan after apply_function_rewrites_| SAME TEXT AS ABOVE_| | logical_plan after count_wildcard_to_time_index_rule_| SAME TEXT AS ABOVE_| @@ -430,9 +430,9 @@ TQL EXPLAIN VERBOSE (0, 10, '5s') test_nano; |_| PromInstantManipulate: range=[0..10000], lookback=[300000], interval=[5000], time index=[j]_| |_|_PromSeriesDivide: tags=["k"]_| |_|_Sort: test_nano.k ASC NULLS FIRST, test_nano.j ASC NULLS FIRST_| -|_|_Projection: test_nano.i, test_nano.k, CAST(test_nano.j AS Timestamp(ms)) AS j_| -|_|_Filter: test_nano.j >= TimestampNanosecond(-299999999999, None) AND test_nano.j < TimestampNanosecond(10001000000, None)_| -|_|_TableScan: test_nano, partial_filters=[test_nano.j >= TimestampNanosecond(-299999999999, None), test_nano.j < TimestampNanosecond(10001000000, None)] | +|_|_Projection: test_nano.i, test_nano.k, test_nano.j_| +|_|_Filter: test_nano.j > TimestampNanosecond(-300000000000, None) AND test_nano.j <= TimestampNanosecond(10000000000, None)_| +|_|_TableScan: test_nano, partial_filters=[test_nano.j > TimestampNanosecond(-300000000000, None), test_nano.j <= TimestampNanosecond(10000000000, None)] | |_| ]]_| | logical_plan after JsonSchemaConcretizeRule_| SAME TEXT AS ABOVE_| | logical_plan after FixStateUdafOrderingAnalyzer_| SAME TEXT AS ABOVE_| @@ -464,9 +464,9 @@ TQL EXPLAIN VERBOSE (0, 10, '5s') test_nano; |_| PromInstantManipulate: range=[0..10000], lookback=[300000], interval=[5000], time index=[j]_| |_|_PromSeriesDivide: tags=["k"]_| |_|_Sort: test_nano.k ASC NULLS FIRST, test_nano.j ASC NULLS FIRST_| -|_|_Projection: test_nano.i, test_nano.k, CAST(test_nano.j AS Timestamp(ms)) AS j_| -|_|_Filter: test_nano.j >= TimestampNanosecond(-299999999999, None) AND test_nano.j < TimestampNanosecond(10001000000, None)_| -|_|_TableScan: test_nano, partial_filters=[test_nano.j >= TimestampNanosecond(-299999999999, None), test_nano.j < TimestampNanosecond(10001000000, None)] | +|_|_Projection: test_nano.i, test_nano.k, test_nano.j_| +|_|_Filter: test_nano.j > TimestampNanosecond(-300000000000, None) AND test_nano.j <= TimestampNanosecond(10000000000, None)_| +|_|_TableScan: test_nano, partial_filters=[test_nano.j > TimestampNanosecond(-300000000000, None), test_nano.j <= TimestampNanosecond(10000000000, None)] | |_| ]]_| | logical_plan after ScanHintRule_| SAME TEXT AS ABOVE_| | logical_plan after JsonTypeConcretizeRule_| SAME TEXT AS ABOVE_| @@ -500,9 +500,9 @@ TQL EXPLAIN VERBOSE (0, 10, '5s') test_nano; |_| PromInstantManipulate: range=[0..10000], lookback=[300000], interval=[5000], time index=[j]_| |_|_PromSeriesDivide: tags=["k"]_| |_|_Sort: test_nano.k ASC NULLS FIRST, test_nano.j ASC NULLS FIRST_| -|_|_Projection: test_nano.i, test_nano.k, CAST(test_nano.j AS Timestamp(ms)) AS j_| -|_|_Filter: test_nano.j >= TimestampNanosecond(-299999999999, None) AND test_nano.j < TimestampNanosecond(10001000000, None)_| -|_|_TableScan: test_nano, partial_filters=[test_nano.j >= TimestampNanosecond(-299999999999, None), test_nano.j < TimestampNanosecond(10001000000, None)] | +|_|_Projection: test_nano.i, test_nano.k, test_nano.j_| +|_|_Filter: test_nano.j > TimestampNanosecond(-300000000000, None) AND test_nano.j <= TimestampNanosecond(10000000000, None)_| +|_|_TableScan: test_nano, partial_filters=[test_nano.j > TimestampNanosecond(-300000000000, None), test_nano.j <= TimestampNanosecond(10000000000, None)] | |_| ]]_| | initial_physical_plan_| MergeScanExec: REDACTED |_|_| diff --git a/tests/cases/standalone/common/tql/general_table.result b/tests/cases/standalone/common/tql/general_table.result index 6172f40aaad..61d414c3339 100644 --- a/tests/cases/standalone/common/tql/general_table.result +++ b/tests/cases/standalone/common/tql/general_table.result @@ -42,7 +42,7 @@ TQL analyze (0, 10, '1s') sum by(job) (irate(cpu_usage{job="fire"}[5s])) / 1e9; |_|_|_PromRangeManipulateExec: req range=[0..10000], interval=[1000], eval range=[5000], time index=[ts] REDACTED |_|_|_PromSeriesNormalizeExec: offset=[0], time index=[ts], filter NaN: [true] REDACTED |_|_|_PromSeriesDivideExec: tags=["job"] REDACTED -|_|_|_ProjectionExec: expr=[value@1 as value, job@0 as job, CAST(ts@2 AS Timestamp(ms)) as ts] REDACTED +|_|_|_ProjectionExec: expr=[value@1 as value, job@0 as job, ts@2 as ts] REDACTED |_|_|_ScanExec: REDACTED |_|_|_| |_|_| Total rows: 0_| diff --git a/tests/cases/standalone/common/types/json/json2_alter.result b/tests/cases/standalone/common/types/json/json2_alter.result new file mode 100644 index 00000000000..5662907d1d7 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_alter.result @@ -0,0 +1,51 @@ +CREATE TABLE application_logs ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true' +); + +Affected Rows: 0 + +ALTER TABLE application_logs + MODIFY COLUMN attrs JSON2; + +Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type + +ALTER TABLE application_logs + MODIFY COLUMN attrs JSON2 (); + +Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type + +ALTER TABLE application_logs + MODIFY COLUMN attrs JSON2 ( + max_auto_expanded_paths = 2000 + ); + +Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type + +ALTER TABLE application_logs + MODIFY COLUMN attrs JSON2 ( + trace_id STRING, + user.id STRING NOT NULL, + user.name STRING DEFAULT 'anonymous', + request_id STRING INVERTED INDEX + ); + +Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type + +ALTER TABLE application_logs + MODIFY COLUMN attrs JSON2 ( + max_auto_expanded_paths = 2000, + trace_id STRING, + user.id STRING NOT NULL, + user.name STRING DEFAULT 'anonymous', + request_id STRING INVERTED INDEX + ); + +Error: 1001(Unsupported), Not supported: ALTER TABLE MODIFY COLUMN to JSON2 type + +DROP TABLE application_logs; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_alter.sql b/tests/cases/standalone/common/types/json/json2_alter.sql new file mode 100644 index 00000000000..2c17138b86d --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_alter.sql @@ -0,0 +1,36 @@ +CREATE TABLE application_logs ( + ts TIMESTAMP TIME INDEX, + attrs JSON2 +) WITH ( + 'append_mode' = 'true' +); + +ALTER TABLE application_logs + MODIFY COLUMN attrs JSON2; + +ALTER TABLE application_logs + MODIFY COLUMN attrs JSON2 (); + +ALTER TABLE application_logs + MODIFY COLUMN attrs JSON2 ( + max_auto_expanded_paths = 2000 + ); + +ALTER TABLE application_logs + MODIFY COLUMN attrs JSON2 ( + trace_id STRING, + user.id STRING NOT NULL, + user.name STRING DEFAULT 'anonymous', + request_id STRING INVERTED INDEX + ); + +ALTER TABLE application_logs + MODIFY COLUMN attrs JSON2 ( + max_auto_expanded_paths = 2000, + trace_id STRING, + user.id STRING NOT NULL, + user.name STRING DEFAULT 'anonymous', + request_id STRING INVERTED INDEX + ); + +DROP TABLE application_logs; diff --git a/tests/cases/standalone/common/types/json/json2_swcs_minmax_stats.result b/tests/cases/standalone/common/types/json/json2_swcs_minmax_stats.result new file mode 100644 index 00000000000..77d48225039 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_swcs_minmax_stats.result @@ -0,0 +1,142 @@ +-- Regression test: min-max pruning must look up row group statistics by +-- parquet leaf column index. On flat-format tables with a JSON2 column, the +-- JSON2 struct expands to multiple leaf columns, so columns after it were +-- previously misaligned with leaves of the JSON2 struct when reading row +-- group statistics. Here `event_time` is the 8th logical column (2 tags + 5 +-- fields), so its statistics were read from the 8th parquet leaf, which is +-- the small-integer `ns_edge` path of the payload. SWCS compaction reads +-- inputs with a time window predicate, so min-max pruning compared the window +-- against the tiny `ns_edge` statistics, dropped every row group and silently +-- lost all rows. The same misaligned statistics also affected plain queries +-- with time-range predicates. +create table json2_swcs_minmax_stats ( + workspace_id string not null, + session_id string not null, + seq bigint not null, + event_time timestamp(9) not null time index, + entry_kind string not null, + payload json2 not null, + schema_version int not null, + is_error boolean not null, + primary key (workspace_id, session_id) +) with ( + 'append_mode' = 'true', + 'sst_format' = 'flat', + 'ttl' = 'forever' +); + +Affected Rows: 0 + +-- 9 rows in the 2026-09-10 window and 2 rows in the 2026-09-11 window. +insert into json2_swcs_minmax_stats values + ('ws', 's1', 1, 1788998400000000001, 'synthetic', '{"case":"swcs-repro","nested":[1,true],"ns_edge":0,"ordinal":1,"synthetic":true}', 1, false), + ('ws', 's2', 2, 1788998400000000002, 'synthetic', '{"case":"swcs-repro","nested":[2,false],"ns_edge":1,"ordinal":2,"synthetic":true}', 1, false), + ('ws', 's3', 3, 1788998400000000003, 'synthetic', '{"case":"swcs-repro","nested":[3,null],"ns_edge":2,"ordinal":3,"synthetic":true}', 1, false), + ('ws', 's1', 4, 1788998400000000004, 'synthetic', '{"case":"swcs-repro","nested":[4,true],"ns_edge":3,"ordinal":4,"synthetic":true}', 1, false), + ('ws', 's2', 5, 1788998400000000005, 'synthetic', '{"case":"swcs-repro","nested":[5,false],"ns_edge":4,"ordinal":5,"synthetic":true}', 1, false), + ('ws', 's3', 6, 1788998400000000006, 'synthetic', '{"case":"swcs-repro","nested":[6,null],"ns_edge":5,"ordinal":6,"synthetic":true}', 1, false), + ('ws', 's1', 7, 1788998400000000007, 'synthetic', '{"case":"swcs-repro","nested":[7,true],"ns_edge":6,"ordinal":7,"synthetic":true}', 1, false), + ('ws', 's2', 8, 1788998400000000008, 'synthetic', '{"case":"swcs-repro","nested":[8,false],"ns_edge":7,"ordinal":8,"synthetic":true}', 1, false), + ('ws', 's3', 9, 1788998400000000009, 'synthetic', '{"case":"swcs-repro","nested":[9,null],"ns_edge":8,"ordinal":9,"synthetic":true}', 1, false), + ('ws', 's1', 10, 1789084800000000001, 'synthetic', '{"case":"swcs-repro","nested":[10,true],"ns_edge":9,"ordinal":10,"synthetic":true}', 1, false), + ('ws', 's2', 11, 1789084800000000002, 'synthetic', '{"case":"swcs-repro","nested":[11,false],"ns_edge":10,"ordinal":11,"synthetic":true}', 1, false); + +Affected Rows: 11 + +select count(*) from json2_swcs_minmax_stats; + ++----------+ +| count(*) | ++----------+ +| 11 | ++----------+ + +admin flush_table('json2_swcs_minmax_stats'); + ++----------------------------------------------+ +| ADMIN flush_table('json2_swcs_minmax_stats') | ++----------------------------------------------+ +| 0 | ++----------------------------------------------+ + +-- Check pruning on flush-written SSTs before any manual compaction. +select count(*) from json2_swcs_minmax_stats +where event_time >= 1788998400000000000 and event_time < 1789084800000000000; + ++----------+ +| count(*) | ++----------+ +| 9 | ++----------+ + +-- SWCS compaction reads the SST with a time window predicate. Before the +-- fix, min-max pruning compared the window against the small `ns_edge` +-- statistics and dropped every row group. +admin compact_table('json2_swcs_minmax_stats', 'swcs', '86400'); + ++-----------------------------------------------------------------+ +| ADMIN compact_table('json2_swcs_minmax_stats', 'swcs', '86400') | ++-----------------------------------------------------------------+ +| 0 | ++-----------------------------------------------------------------+ + +select count(*) from json2_swcs_minmax_stats; + ++----------+ +| count(*) | ++----------+ +| 11 | ++----------+ + +-- Compact once more to also exercise reading compaction-written SSTs. +admin compact_table('json2_swcs_minmax_stats', 'swcs', '86400'); + ++-----------------------------------------------------------------+ +| ADMIN compact_table('json2_swcs_minmax_stats', 'swcs', '86400') | ++-----------------------------------------------------------------+ +| 0 | ++-----------------------------------------------------------------+ + +select count(*) from json2_swcs_minmax_stats; + ++----------+ +| count(*) | ++----------+ +| 11 | ++----------+ + +-- Time-range predicates use the same statistics and must not prune the rows +-- either. +select count(*) from json2_swcs_minmax_stats +where event_time >= 1788998400000000000 and event_time < 1789084800000000000; + ++----------+ +| count(*) | ++----------+ +| 9 | ++----------+ + +select workspace_id, session_id, seq, event_time, entry_kind, payload, schema_version, is_error +from json2_swcs_minmax_stats +order by seq; + ++--------------+------------+-----+-------------------------------+------------+--------------------------------------------------------------------------------------+----------------+----------+ +| workspace_id | session_id | seq | event_time | entry_kind | payload | schema_version | is_error | ++--------------+------------+-----+-------------------------------+------------+--------------------------------------------------------------------------------------+----------------+----------+ +| ws | s1 | 1 | 2026-09-10T00:00:00.000000001 | synthetic | {"case":"swcs-repro","nested":[1,true],"ns_edge":0,"ordinal":1,"synthetic":true} | 1 | false | +| ws | s2 | 2 | 2026-09-10T00:00:00.000000002 | synthetic | {"case":"swcs-repro","nested":[2,false],"ns_edge":1,"ordinal":2,"synthetic":true} | 1 | false | +| ws | s3 | 3 | 2026-09-10T00:00:00.000000003 | synthetic | {"case":"swcs-repro","nested":[3,null],"ns_edge":2,"ordinal":3,"synthetic":true} | 1 | false | +| ws | s1 | 4 | 2026-09-10T00:00:00.000000004 | synthetic | {"case":"swcs-repro","nested":[4,true],"ns_edge":3,"ordinal":4,"synthetic":true} | 1 | false | +| ws | s2 | 5 | 2026-09-10T00:00:00.000000005 | synthetic | {"case":"swcs-repro","nested":[5,false],"ns_edge":4,"ordinal":5,"synthetic":true} | 1 | false | +| ws | s3 | 6 | 2026-09-10T00:00:00.000000006 | synthetic | {"case":"swcs-repro","nested":[6,null],"ns_edge":5,"ordinal":6,"synthetic":true} | 1 | false | +| ws | s1 | 7 | 2026-09-10T00:00:00.000000007 | synthetic | {"case":"swcs-repro","nested":[7,true],"ns_edge":6,"ordinal":7,"synthetic":true} | 1 | false | +| ws | s2 | 8 | 2026-09-10T00:00:00.000000008 | synthetic | {"case":"swcs-repro","nested":[8,false],"ns_edge":7,"ordinal":8,"synthetic":true} | 1 | false | +| ws | s3 | 9 | 2026-09-10T00:00:00.000000009 | synthetic | {"case":"swcs-repro","nested":[9,null],"ns_edge":8,"ordinal":9,"synthetic":true} | 1 | false | +| ws | s1 | 10 | 2026-09-11T00:00:00.000000001 | synthetic | {"case":"swcs-repro","nested":[10,true],"ns_edge":9,"ordinal":10,"synthetic":true} | 1 | false | +| ws | s2 | 11 | 2026-09-11T00:00:00.000000002 | synthetic | {"case":"swcs-repro","nested":[11,false],"ns_edge":10,"ordinal":11,"synthetic":true} | 1 | false | ++--------------+------------+-----+-------------------------------+------------+--------------------------------------------------------------------------------------+----------------+----------+ + +drop table json2_swcs_minmax_stats; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/common/types/json/json2_swcs_minmax_stats.sql b/tests/cases/standalone/common/types/json/json2_swcs_minmax_stats.sql new file mode 100644 index 00000000000..7bfea334420 --- /dev/null +++ b/tests/cases/standalone/common/types/json/json2_swcs_minmax_stats.sql @@ -0,0 +1,72 @@ +-- Regression test: min-max pruning must look up row group statistics by +-- parquet leaf column index. On flat-format tables with a JSON2 column, the +-- JSON2 struct expands to multiple leaf columns, so columns after it were +-- previously misaligned with leaves of the JSON2 struct when reading row +-- group statistics. Here `event_time` is the 8th logical column (2 tags + 5 +-- fields), so its statistics were read from the 8th parquet leaf, which is +-- the small-integer `ns_edge` path of the payload. SWCS compaction reads +-- inputs with a time window predicate, so min-max pruning compared the window +-- against the tiny `ns_edge` statistics, dropped every row group and silently +-- lost all rows. The same misaligned statistics also affected plain queries +-- with time-range predicates. + +create table json2_swcs_minmax_stats ( + workspace_id string not null, + session_id string not null, + seq bigint not null, + event_time timestamp(9) not null time index, + entry_kind string not null, + payload json2 not null, + schema_version int not null, + is_error boolean not null, + primary key (workspace_id, session_id) +) with ( + 'append_mode' = 'true', + 'sst_format' = 'flat', + 'ttl' = 'forever' +); + +-- 9 rows in the 2026-09-10 window and 2 rows in the 2026-09-11 window. +insert into json2_swcs_minmax_stats values + ('ws', 's1', 1, 1788998400000000001, 'synthetic', '{"case":"swcs-repro","nested":[1,true],"ns_edge":0,"ordinal":1,"synthetic":true}', 1, false), + ('ws', 's2', 2, 1788998400000000002, 'synthetic', '{"case":"swcs-repro","nested":[2,false],"ns_edge":1,"ordinal":2,"synthetic":true}', 1, false), + ('ws', 's3', 3, 1788998400000000003, 'synthetic', '{"case":"swcs-repro","nested":[3,null],"ns_edge":2,"ordinal":3,"synthetic":true}', 1, false), + ('ws', 's1', 4, 1788998400000000004, 'synthetic', '{"case":"swcs-repro","nested":[4,true],"ns_edge":3,"ordinal":4,"synthetic":true}', 1, false), + ('ws', 's2', 5, 1788998400000000005, 'synthetic', '{"case":"swcs-repro","nested":[5,false],"ns_edge":4,"ordinal":5,"synthetic":true}', 1, false), + ('ws', 's3', 6, 1788998400000000006, 'synthetic', '{"case":"swcs-repro","nested":[6,null],"ns_edge":5,"ordinal":6,"synthetic":true}', 1, false), + ('ws', 's1', 7, 1788998400000000007, 'synthetic', '{"case":"swcs-repro","nested":[7,true],"ns_edge":6,"ordinal":7,"synthetic":true}', 1, false), + ('ws', 's2', 8, 1788998400000000008, 'synthetic', '{"case":"swcs-repro","nested":[8,false],"ns_edge":7,"ordinal":8,"synthetic":true}', 1, false), + ('ws', 's3', 9, 1788998400000000009, 'synthetic', '{"case":"swcs-repro","nested":[9,null],"ns_edge":8,"ordinal":9,"synthetic":true}', 1, false), + ('ws', 's1', 10, 1789084800000000001, 'synthetic', '{"case":"swcs-repro","nested":[10,true],"ns_edge":9,"ordinal":10,"synthetic":true}', 1, false), + ('ws', 's2', 11, 1789084800000000002, 'synthetic', '{"case":"swcs-repro","nested":[11,false],"ns_edge":10,"ordinal":11,"synthetic":true}', 1, false); + +select count(*) from json2_swcs_minmax_stats; + +admin flush_table('json2_swcs_minmax_stats'); + +-- Check pruning on flush-written SSTs before any manual compaction. +select count(*) from json2_swcs_minmax_stats +where event_time >= 1788998400000000000 and event_time < 1789084800000000000; + +-- SWCS compaction reads the SST with a time window predicate. Before the +-- fix, min-max pruning compared the window against the small `ns_edge` +-- statistics and dropped every row group. +admin compact_table('json2_swcs_minmax_stats', 'swcs', '86400'); + +select count(*) from json2_swcs_minmax_stats; + +-- Compact once more to also exercise reading compaction-written SSTs. +admin compact_table('json2_swcs_minmax_stats', 'swcs', '86400'); + +select count(*) from json2_swcs_minmax_stats; + +-- Time-range predicates use the same statistics and must not prune the rows +-- either. +select count(*) from json2_swcs_minmax_stats +where event_time >= 1788998400000000000 and event_time < 1789084800000000000; + +select workspace_id, session_id, seq, event_time, entry_kind, payload, schema_version, is_error +from json2_swcs_minmax_stats +order by seq; + +drop table json2_swcs_minmax_stats; diff --git a/tests/cases/standalone/copy/copy_skip_wal.result b/tests/cases/standalone/copy/copy_skip_wal.result new file mode 100644 index 00000000000..9f7a2b71cea --- /dev/null +++ b/tests/cases/standalone/copy/copy_skip_wal.result @@ -0,0 +1,274 @@ +CREATE TABLE copy_skip_wal(host STRING, ts TIMESTAMP TIME INDEX); + +Affected Rows: 0 + +INSERT INTO copy_skip_wal VALUES ('host1', 1000), ('host2', 2000); + +Affected Rows: 2 + +COPY copy_skip_wal TO '${SQLNESS_HOME}/copy_skip_wal/copy_skip_wal.csv' WITH (FORMAT='csv'); + +Affected Rows: 2 + +TRUNCATE TABLE copy_skip_wal; + +Affected Rows: 0 + +-- MYSQL: COPY TABLE inherits the connection WAL policy. +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = true; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +COPY copy_skip_wal FROM '${SQLNESS_HOME}/copy_skip_wal/copy_skip_wal.csv' WITH (FORMAT='csv'); + +affected_rows: 2 + +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM copy_skip_wal ORDER BY ts; + ++-------+---------------------+ +| host | ts | ++-------+---------------------+ +| host1 | 1970-01-01 00:00:01 | +| host2 | 1970-01-01 00:00:02 | ++-------+---------------------+ + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM copy_skip_wal ORDER BY ts; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = false; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +COPY copy_skip_wal FROM '${SQLNESS_HOME}/copy_skip_wal/copy_skip_wal.csv' WITH (FORMAT='csv'); + +affected_rows: 2 + +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM copy_skip_wal ORDER BY ts; + ++-------+---------------------+ +| host | ts | ++-------+---------------------+ +| host1 | 1970-01-01 00:00:01 | +| host2 | 1970-01-01 00:00:02 | ++-------+---------------------+ + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM copy_skip_wal ORDER BY ts; + ++-------+---------------------+ +| host | ts | ++-------+---------------------+ +| host1 | 1970-01-01 00:00:01 | +| host2 | 1970-01-01 00:00:02 | ++-------+---------------------+ + +TRUNCATE TABLE copy_skip_wal; + +Affected Rows: 0 + +-- MYSQL: COPY DATABASE inherits the connection WAL policy. +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = true; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +COPY DATABASE public FROM '${SQLNESS_HOME}/copy_skip_wal/' WITH (FORMAT='csv'); + +affected_rows: 2 + +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM copy_skip_wal ORDER BY ts; + ++-------+---------------------+ +| host | ts | ++-------+---------------------+ +| host1 | 1970-01-01 00:00:01 | +| host2 | 1970-01-01 00:00:02 | ++-------+---------------------+ + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM copy_skip_wal ORDER BY ts; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = false; + +affected_rows: 0 + +-- SQLNESS PROTOCOL MYSQL +COPY DATABASE public FROM '${SQLNESS_HOME}/copy_skip_wal/' WITH (FORMAT='csv'); + +affected_rows: 2 + +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM copy_skip_wal ORDER BY ts; + ++-------+---------------------+ +| host | ts | ++-------+---------------------+ +| host1 | 1970-01-01 00:00:01 | +| host2 | 1970-01-01 00:00:02 | ++-------+---------------------+ + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM copy_skip_wal ORDER BY ts; + ++-------+---------------------+ +| host | ts | ++-------+---------------------+ +| host1 | 1970-01-01 00:00:01 | +| host2 | 1970-01-01 00:00:02 | ++-------+---------------------+ + +TRUNCATE TABLE copy_skip_wal; + +Affected Rows: 0 + +-- POSTGRES: COPY TABLE inherits the connection WAL policy. +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = true; + +Affected Rows: 0 + +-- SQLNESS PROTOCOL POSTGRES +COPY copy_skip_wal FROM '${SQLNESS_HOME}/copy_skip_wal/copy_skip_wal.csv' WITH (FORMAT='csv'); + +Affected Rows: 2 + +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM copy_skip_wal ORDER BY ts; + ++-------+----------------------------+ +| host | ts | ++-------+----------------------------+ +| host1 | 1970-01-01 00:00:01.000000 | +| host2 | 1970-01-01 00:00:02.000000 | ++-------+----------------------------+ + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM copy_skip_wal ORDER BY ts; + ++------+----+ +| host | ts | ++------+----+ ++------+----+ + +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = false; + +Affected Rows: 0 + +-- SQLNESS PROTOCOL POSTGRES +COPY copy_skip_wal FROM '${SQLNESS_HOME}/copy_skip_wal/copy_skip_wal.csv' WITH (FORMAT='csv'); + +Affected Rows: 2 + +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM copy_skip_wal ORDER BY ts; + ++-------+----------------------------+ +| host | ts | ++-------+----------------------------+ +| host1 | 1970-01-01 00:00:01.000000 | +| host2 | 1970-01-01 00:00:02.000000 | ++-------+----------------------------+ + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM copy_skip_wal ORDER BY ts; + ++-------+----------------------------+ +| host | ts | ++-------+----------------------------+ +| host1 | 1970-01-01 00:00:01.000000 | +| host2 | 1970-01-01 00:00:02.000000 | ++-------+----------------------------+ + +TRUNCATE TABLE copy_skip_wal; + +Affected Rows: 0 + +-- POSTGRES: COPY DATABASE inherits the connection WAL policy. +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = true; + +Affected Rows: 0 + +-- SQLNESS PROTOCOL POSTGRES +COPY DATABASE public FROM '${SQLNESS_HOME}/copy_skip_wal/' WITH (FORMAT='csv'); + +Affected Rows: 2 + +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM copy_skip_wal ORDER BY ts; + ++-------+----------------------------+ +| host | ts | ++-------+----------------------------+ +| host1 | 1970-01-01 00:00:01.000000 | +| host2 | 1970-01-01 00:00:02.000000 | ++-------+----------------------------+ + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM copy_skip_wal ORDER BY ts; + ++------+----+ +| host | ts | ++------+----+ ++------+----+ + +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = false; + +Affected Rows: 0 + +-- SQLNESS PROTOCOL POSTGRES +COPY DATABASE public FROM '${SQLNESS_HOME}/copy_skip_wal/' WITH (FORMAT='csv'); + +Affected Rows: 2 + +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM copy_skip_wal ORDER BY ts; + ++-------+----------------------------+ +| host | ts | ++-------+----------------------------+ +| host1 | 1970-01-01 00:00:01.000000 | +| host2 | 1970-01-01 00:00:02.000000 | ++-------+----------------------------+ + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM copy_skip_wal ORDER BY ts; + ++-------+----------------------------+ +| host | ts | ++-------+----------------------------+ +| host1 | 1970-01-01 00:00:01.000000 | +| host2 | 1970-01-01 00:00:02.000000 | ++-------+----------------------------+ + +TRUNCATE TABLE copy_skip_wal; + +Affected Rows: 0 + +DROP TABLE copy_skip_wal; + +Affected Rows: 0 + diff --git a/tests/cases/standalone/copy/copy_skip_wal.sql b/tests/cases/standalone/copy/copy_skip_wal.sql new file mode 100644 index 00000000000..7ead655de58 --- /dev/null +++ b/tests/cases/standalone/copy/copy_skip_wal.sql @@ -0,0 +1,125 @@ +CREATE TABLE copy_skip_wal(host STRING, ts TIMESTAMP TIME INDEX); + +INSERT INTO copy_skip_wal VALUES ('host1', 1000), ('host2', 2000); + +COPY copy_skip_wal TO '${SQLNESS_HOME}/copy_skip_wal/copy_skip_wal.csv' WITH (FORMAT='csv'); + +TRUNCATE TABLE copy_skip_wal; + +-- MYSQL: COPY TABLE inherits the connection WAL policy. +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = true; + +-- SQLNESS PROTOCOL MYSQL +COPY copy_skip_wal FROM '${SQLNESS_HOME}/copy_skip_wal/copy_skip_wal.csv' WITH (FORMAT='csv'); + +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM copy_skip_wal ORDER BY ts; + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM copy_skip_wal ORDER BY ts; + +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = false; + +-- SQLNESS PROTOCOL MYSQL +COPY copy_skip_wal FROM '${SQLNESS_HOME}/copy_skip_wal/copy_skip_wal.csv' WITH (FORMAT='csv'); + +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM copy_skip_wal ORDER BY ts; + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM copy_skip_wal ORDER BY ts; + +TRUNCATE TABLE copy_skip_wal; + +-- MYSQL: COPY DATABASE inherits the connection WAL policy. +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = true; + +-- SQLNESS PROTOCOL MYSQL +COPY DATABASE public FROM '${SQLNESS_HOME}/copy_skip_wal/' WITH (FORMAT='csv'); + +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM copy_skip_wal ORDER BY ts; + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM copy_skip_wal ORDER BY ts; + +-- SQLNESS PROTOCOL MYSQL +SET skip_wal = false; + +-- SQLNESS PROTOCOL MYSQL +COPY DATABASE public FROM '${SQLNESS_HOME}/copy_skip_wal/' WITH (FORMAT='csv'); + +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM copy_skip_wal ORDER BY ts; + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL MYSQL +SELECT * FROM copy_skip_wal ORDER BY ts; + +TRUNCATE TABLE copy_skip_wal; + +-- POSTGRES: COPY TABLE inherits the connection WAL policy. +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = true; + +-- SQLNESS PROTOCOL POSTGRES +COPY copy_skip_wal FROM '${SQLNESS_HOME}/copy_skip_wal/copy_skip_wal.csv' WITH (FORMAT='csv'); + +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM copy_skip_wal ORDER BY ts; + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM copy_skip_wal ORDER BY ts; + +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = false; + +-- SQLNESS PROTOCOL POSTGRES +COPY copy_skip_wal FROM '${SQLNESS_HOME}/copy_skip_wal/copy_skip_wal.csv' WITH (FORMAT='csv'); + +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM copy_skip_wal ORDER BY ts; + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM copy_skip_wal ORDER BY ts; + +TRUNCATE TABLE copy_skip_wal; + +-- POSTGRES: COPY DATABASE inherits the connection WAL policy. +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = true; + +-- SQLNESS PROTOCOL POSTGRES +COPY DATABASE public FROM '${SQLNESS_HOME}/copy_skip_wal/' WITH (FORMAT='csv'); + +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM copy_skip_wal ORDER BY ts; + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM copy_skip_wal ORDER BY ts; + +-- SQLNESS PROTOCOL POSTGRES +SET skip_wal = false; + +-- SQLNESS PROTOCOL POSTGRES +COPY DATABASE public FROM '${SQLNESS_HOME}/copy_skip_wal/' WITH (FORMAT='csv'); + +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM copy_skip_wal ORDER BY ts; + +-- SQLNESS ARG restart=true +-- SQLNESS PROTOCOL POSTGRES +SELECT * FROM copy_skip_wal ORDER BY ts; + +TRUNCATE TABLE copy_skip_wal; + +DROP TABLE copy_skip_wal; diff --git a/tests/compatibility/cases/downgrade_compatibility/case.toml b/tests/compatibility/cases/downgrade_compatibility/case.toml index ba5ee5065ed..ab0496e523d 100644 --- a/tests/compatibility/cases/downgrade_compatibility/case.toml +++ b/tests/compatibility/cases/downgrade_compatibility/case.toml @@ -1,8 +1,8 @@ name = "downgrade_compatibility" -reason = "Verify v1.1.4 can reopen a table whose region WAL options were written by the current binary, and can read all rows of an append-only table flushed and compacted with preserve_row_sequence enabled." +reason = "Verify v1.1.4 can reopen tables whose region WAL options or byte-stream-split (BSS) float SSTs were written by the current binary, and can read all rows of an append-only table flushed and compacted with preserve_row_sequence enabled." introduced_by = "fix: preserve legacy region WAL options format; preserve_row_sequence" topologies = ["distributed", "standalone"] from_range = [">=v1.2.0"] to_range = ["=v1.1.4"] -features = ["table", "wal", "downgrade", "append", "preserve_row_sequence"] +features = ["table", "wal", "downgrade", "append", "preserve_row_sequence", "sst", "float", "byte_stream_split"] owner = "metasrv" diff --git a/tests/compatibility/cases/downgrade_compatibility/setup.sql b/tests/compatibility/cases/downgrade_compatibility/setup.sql index a9415570d31..267466f9dfe 100644 --- a/tests/compatibility/cases/downgrade_compatibility/setup.sql +++ b/tests/compatibility/cases/downgrade_compatibility/setup.sql @@ -32,3 +32,18 @@ INSERT INTO t_preserve_sequence_downgrade VALUES ADMIN FLUSH_TABLE('t_preserve_sequence_downgrade'); ADMIN COMPACT_TABLE('t_preserve_sequence_downgrade'); + +CREATE TABLE t_sst_float_field_encoding_downgrade( + ts TIMESTAMP TIME INDEX, + host STRING PRIMARY KEY, + f FLOAT, + d DOUBLE +) ENGINE=mito +WITH('experimental_sst_float_field_encoding'='byte_stream_split'); + +INSERT INTO t_sst_float_field_encoding_downgrade VALUES +('2024-02-10 00:00:00+0000', 'host_a', 1.25, 10.5), +('2024-02-10 00:01:00+0000', 'host_b', -2.5, 20.25), +('2024-02-10 00:02:00+0000', 'host_c', NULL, NULL); + +ADMIN FLUSH_TABLE('t_sst_float_field_encoding_downgrade'); diff --git a/tests/compatibility/cases/downgrade_compatibility/verify.result b/tests/compatibility/cases/downgrade_compatibility/verify.result index 3c052392927..c350288d19d 100644 --- a/tests/compatibility/cases/downgrade_compatibility/verify.result +++ b/tests/compatibility/cases/downgrade_compatibility/verify.result @@ -17,3 +17,15 @@ SELECT ts, host, val FROM t_preserve_sequence_downgrade ORDER BY ts, host; | 2024-02-09T00:02:00 | host_a | 3 | | 2024-02-09T00:03:00 | host_c | 4 | +---------------------+--------+-----+ + +SELECT ts, host, f, d +FROM t_sst_float_field_encoding_downgrade +ORDER BY ts, host; + ++---------------------+--------+------+-------+ +| ts | host | f | d | ++---------------------+--------+------+-------+ +| 2024-02-10T00:00:00 | host_a | 1.25 | 10.5 | +| 2024-02-10T00:01:00 | host_b | -2.5 | 20.25 | +| 2024-02-10T00:02:00 | host_c | | | ++---------------------+--------+------+-------+ diff --git a/tests/compatibility/cases/downgrade_compatibility/verify.sql b/tests/compatibility/cases/downgrade_compatibility/verify.sql index 764968bf1b8..6b73c886117 100644 --- a/tests/compatibility/cases/downgrade_compatibility/verify.sql +++ b/tests/compatibility/cases/downgrade_compatibility/verify.sql @@ -1,3 +1,7 @@ SELECT ts, host, val FROM t_downgrade_compatibility ORDER BY ts, host; SELECT ts, host, val FROM t_preserve_sequence_downgrade ORDER BY ts, host; + +SELECT ts, host, f, d +FROM t_sst_float_field_encoding_downgrade +ORDER BY ts, host; diff --git a/tests/compatibility/cases/sst_float_field_encoding_upgrade/case.toml b/tests/compatibility/cases/sst_float_field_encoding_upgrade/case.toml new file mode 100644 index 00000000000..4f008775bb2 --- /dev/null +++ b/tests/compatibility/cases/sst_float_field_encoding_upgrade/case.toml @@ -0,0 +1,8 @@ +name = "sst_float_field_encoding_upgrade" +reason = "Verify newer binaries can read default-encoded Mito SSTs containing FLOAT and DOUBLE fields written by older binaries." +introduced_by = "experimental_sst_float_field_encoding" +topologies = ["distributed", "standalone"] +from_range = ["*"] +to_range = ["*"] +features = ["table", "sst", "float", "upgrade"] +owner = "storage" diff --git a/tests/compatibility/cases/sst_float_field_encoding_upgrade/setup.sql b/tests/compatibility/cases/sst_float_field_encoding_upgrade/setup.sql new file mode 100644 index 00000000000..a84b4339dfc --- /dev/null +++ b/tests/compatibility/cases/sst_float_field_encoding_upgrade/setup.sql @@ -0,0 +1,13 @@ +CREATE TABLE t_sst_float_field_encoding_upgrade( + ts TIMESTAMP TIME INDEX, + host STRING PRIMARY KEY, + f FLOAT, + d DOUBLE +) ENGINE=mito; + +INSERT INTO t_sst_float_field_encoding_upgrade VALUES +('2024-02-10 00:00:00+0000', 'host_a', 1.25, 10.5), +('2024-02-10 00:01:00+0000', 'host_b', -2.5, 20.25), +('2024-02-10 00:02:00+0000', 'host_c', NULL, NULL); + +ADMIN FLUSH_TABLE('t_sst_float_field_encoding_upgrade'); diff --git a/tests/compatibility/cases/sst_float_field_encoding_upgrade/verify.result b/tests/compatibility/cases/sst_float_field_encoding_upgrade/verify.result new file mode 100644 index 00000000000..d6ae25d9747 --- /dev/null +++ b/tests/compatibility/cases/sst_float_field_encoding_upgrade/verify.result @@ -0,0 +1,11 @@ +SELECT ts, host, f, d +FROM t_sst_float_field_encoding_upgrade +ORDER BY ts, host; + ++---------------------+--------+------+-------+ +| ts | host | f | d | ++---------------------+--------+------+-------+ +| 2024-02-10T00:00:00 | host_a | 1.25 | 10.5 | +| 2024-02-10T00:01:00 | host_b | -2.5 | 20.25 | +| 2024-02-10T00:02:00 | host_c | | | ++---------------------+--------+------+-------+ diff --git a/tests/compatibility/cases/sst_float_field_encoding_upgrade/verify.sql b/tests/compatibility/cases/sst_float_field_encoding_upgrade/verify.sql new file mode 100644 index 00000000000..81822cb0ac8 --- /dev/null +++ b/tests/compatibility/cases/sst_float_field_encoding_upgrade/verify.sql @@ -0,0 +1,3 @@ +SELECT ts, host, f, d +FROM t_sst_float_field_encoding_upgrade +ORDER BY ts, host; diff --git a/tests/perf/README.md b/tests/perf/README.md index d518fec996c..0361f598a3b 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -86,15 +86,19 @@ sample values without changing label cardinality or the runner lifecycle: [scenario.remote_write.value] pattern = "quantized_signal" # linear, constant, modulo, unique, seeded_random, # run_length, quantized_signal, - # signal_with_sporadic_stalls, mixed_signal_repeated + # signal_with_sporadic_stalls, mixed_signal_repeated, + # bounded_mixed base = 0.0 step = 0.125 -cardinality = 4096 # buckets for modulo/seeded_random/quantized_signal/run_length -seed = 12345 # deterministic seeded_random input -run_length = 8 # adjacent samples per bucket for run_length/quantized_signal +cardinality = 4096 # buckets for modulo/seeded_random/quantized_signal/run_length; + # baseline buckets for bounded_mixed +seed = 12345 # deterministic seeded_random input and bounded_mixed series hash +run_length = 8 # adjacent samples per bucket for run_length/quantized_signal; + # interpolation interval for bounded_mixed stall_every = 100 # interval for signal_with_sporadic_stalls stall_length = 16 # held samples inside each stall interval -mixed_every = 5 # every Nth sample becomes the repeated base value +mixed_every = 5 # every Nth sample becomes the repeated base value; + # one fractional bounded_mixed series per N series ``` The default `linear` pattern preserves the helper's historical formula. Use @@ -103,12 +107,29 @@ data shapes, `run_length` for run-heavy low-cardinality series, `quantized_signa for signal-like values collapsed into a finite bucket set, `signal_with_sporadic_stalls` for mostly continuous signals with periodic flat spots, and `mixed_signal_repeated` for signal-plus-periodic-default mixtures. `unique` or high-cardinality buckets -still work for broad sample-value distributions. This is a generic sample-value -control for query/ingestion cases; it does not inspect or assert storage -encoding, Parquet footers, or storage policy choices. For chunked remote-write -ingestion, the runner passes the sample offset and total sample count to the -helper so non-linear value patterns use a stable global/per-series ordinal across -chunks. +still work for broad sample-value distributions. + +`bounded_mixed` is synthetic rather than an empirical workload. For series `s`, +`h = splitmix64(s ^ seed)` selects a span from `[10, 1_000, 100_000, 10_000_000]` +using its top two bits, multiplied by `1 + ((h >> 32) % 10)`. Its baseline is +`(h % max(cardinality, 1)) * span / 10`. Ranges can overlap. At local sample +`l = sample_offset + sample_idx`, it linearly interpolates hash-derived anchors +in `[-span + 1, span - 1]`, changing anchors every `max(run_length, 1)` samples. +It applies `base + step * signal`, then rounds to an integer. Every +`max(mixed_every, 1)`-th series adds a hash-derived fraction from `1/1000` to +`999/1000`. The 95/5 finite, nonintegral guarantee applies to the case's bounded +parameters, not arbitrary huge or nonfinite base/step inputs. With `base = 0` +and `step = 1`, output lies within `baseline ± span`, allowing one unit for +rounding and the fractional addition. The explicit case uses 1,000 equal-length +series and `mixed_every = 20`, giving exactly 95% integer-valued and 5% +fractional Float64 samples. Its ranges and temporal shape are synthetic design +parameters, not measured properties from the survey. + +This is a generic sample-value control for query/ingestion cases; it does not +inspect or assert storage encoding, Parquet footers, or storage policy choices. +For chunked remote-write ingestion, the runner passes the sample offset and total +sample count to the helper so non-linear value patterns use a stable global/per-series +ordinal across chunks. Case schema, value distribution defaults, storage defaults, and read-bench defaults are owned by Rust. The outer CI driver calls @@ -162,12 +183,15 @@ measures region scan cost, and query measurements still exercise the SQL/TQL frontend path. Treat all performance conclusions as release-only; debug builds are suitable only for command wiring and correctness checks. -`prepare-remote` creates the configured database if needed. The outer driver writes a per-target -frontend config enabling `[prom_store]` with metric engine storage and a non-zero -`pending_rows_flush_interval`, and validates that the logical metric table reaches -`series_count * samples_per_series` rows before trusting the query measurements. -Use `--fixture-generator /path/to/query_perf_fixture` to provide the Rust helper -to the outer driver. +`prepare-remote` creates the configured database if needed. A remote-write case +can provide `base_setup_sql` and `candidate_setup_sql` lists; each target runs its +own complete statements in order after database creation and before ingestion. +The outer driver writes a per-target frontend config enabling `[prom_store]` with +metric engine storage and a non-zero `pending_rows_flush_interval`, and validates +that the logical metric table reaches `series_count * samples_per_series` rows +before trusting the query measurements. Use +`--fixture-generator /path/to/query_perf_fixture` to provide the Rust helper to +the outer driver. Large manual remote-write cases can set `sample_chunk_size` to split ingestion by time. For each chunk, `prepare-remote` invokes `query_perf_fixture prom-remote-write` with the @@ -204,9 +228,16 @@ case for issue #7913. It writes 8192 series × 20160 samples through remote-writ in 1440-sample daily time chunks, flushing after each chunk before running 1d/7d/14d TQL selectors. It is not included in the default `all` case set because ingestion cost dominates routine CI validation. Commenting `/query-regression heavy` runs -only this case; `/query-regression` runs the seven routine default cases. Manual +only this case; `/query-regression` runs the nine routine default cases, including +`sst_float_bss`, `promql_instant_last_row_9034`, and +`mito_prefilter_all_match`. Manual workflow dispatch accepts the `heavy` token to select this case. +The routine default set also includes +`tests/perf/query_cases/mito_prefilter_all_match/case.toml`. It covers the +Mito prefilter all-match path, a mixed integer/float filter, and integer and +float selective controls over a 262144-row direct SST fixture. + ## OTLP trace load scenario `scenario.kind = "otlp_trace_load"` runs a bounded native `otelgen` process @@ -366,6 +397,91 @@ uv run --no-project python .github/scripts/query-regression-run.py \ --work-dir /tmp/query-regression-work ``` +### SST float BSS comparison + +`tests/perf/query_cases/sst_float_bss/case.toml` is included in the routine +`all` default case group and compares a default empty physical metric table with +a candidate byte-stream-split (BSS) physical table. It writes 1,000 series × 4,320 samples +(4,320,000 rows) using synthetic `bounded_mixed` values: 95% integral series and +5% nonintegral series, with bounded per-series fluctuation rather than globally +unique values. This deliberately matches a 95/5 design; it is not an empirical +claim about any production population. A separate survey sample found 94.79% +integral values, but that observation does not make its values globally unique. +Current mixed-data evidence, including SST inspection, exact-bit row verification, +warmed endpoint SQL, and both warm-reader projections, is recorded in +[`query_cases/sst_float_bss/RESULTS.md`](query_cases/sst_float_bss/RESULTS.md). +Timing observations were collected while other builds saturated the shared host; +they are not performance acceptance evidence. Before rerunning timings, ensure +the host is idle (not merely this agent), record load throughout, and avoid +concurrent builds. CPU affinity alone does not isolate memory or I/O contention. +The historical unique-integer workload is preserved only in Git history and does +not apply to this case. +Both targets must use the exact same release `greptime` binary; only the +per-target table setup SQL differs. Run the existing driver from a checkout +containing that binary, with absolute paths and fresh data directories for every +run: + +```bash +REPO="$(pwd -P)" +RELEASE_GREPTIME="/absolute/path/to/release/greptime" +FIXTURE_GENERATOR="/absolute/path/to/release/query_perf_fixture" +RUNNER="/absolute/path/to/release/query_regression_runner" +WORK_DIR="/absolute/path/to/fresh/sst-float-bss-run-1" +cd "$REPO" +uv run --no-project python "$REPO/.github/scripts/query-regression-run.py" \ + --cases "$REPO/tests/perf/query_cases/sst_float_bss/case.toml" \ + --base-src "$REPO" \ + --candidate-src . \ + --base-bin "$RELEASE_GREPTIME" \ + --candidate-bin "$RELEASE_GREPTIME" \ + --fixture-generator "$FIXTURE_GENERATOR" \ + --runner "$RUNNER" \ + --work-dir "$WORK_DIR" \ + --summary-script "$REPO/.github/scripts/query-regression-summary.py" +``` + +Repeat the command three times with a different fresh absolute `WORK_DIR` each +run. In each `query-regression-report.json`, compare base and candidate +`targets[].storage_inspection.summary.summary.total_file_size`, and each query's +`targets[].measurements[].latency_ms_median`. Storage percentage is +`(candidate_total_file_size - base_total_file_size) / base_total_file_size * 100`; +query latency percentage is +`(candidate_latency_ms_median - base_latency_ms_median) / base_latency_ms_median * 100`. +The configured three warmups occur after the initial query validation and before +that query's 15 measured endpoint requests, so query latency is a warmed +frontend/cache measurement. The case restricts storage inspection to +`data/greptime/public`, excluding `greptime_private` SSTs. After the datanodes +stop, it also runs seven iterations of value-only `parquetbench` for all +inspected SST files and sequential `scanbench` over the corresponding regions, +with parallelism one. The three explicit flushes yielded six SSTs in the observed +mixed runs—three 1,105,920-row files and three 334,080-row files—because storage split each +flush; this is an observation, not a guaranteed file layout. Their per-run +output and aggregate `parquetbench_median_average_ms` and +`scanbench_median_average_ms` are under +`targets[].read_bench`; they are quiescent local-file read/scan diagnostics, not +warmed frontend-query latency measurements. Bench averages include iteration one; +no OS cache is dropped, and the driver runs base before candidate. + +Historical [PR #8548](https://github.com/GreptimeTeam/greptimedb/pull/8548) +reported storage savings alongside warm-read slowdowns for a different mixed +counter/gauge study. It used value-only and all-column projections, discarded the +first iteration, and alternated target order. This case reuses that reader +measurement approach, not its dataset or results. For the supplementary warm +comparison, reuse recorded commands against stopped data directories with eight +iterations, discard iteration one, and alternate target order for eight rounds. +For parquetbench, sum all per-file warm medians before taking the outer median; +for scanbench, take the outer median of whole-region warm medians. Keep +post-flush and post-compaction measurements separate. + +The case's `-5.0` storage target and `25` query-latency guardrail are experimental +acceptance targets, not observed-benefit claims; do not relax them if a run +fails. Its three periodic flushes and shared high TWCS trigger avoid the normal +four-file compaction trigger from confounding the layout. Check footer encodings +(BSS on candidate and no BSS on base) and data equality from the artifacts rather +than through a new harness framework. + +For a focused manual reproduction of the Mito prefilter all-match optimization, use the existing lifecycle command above with `--cases tests/perf/query_cases/mito_prefilter_all_match/case.toml`. Its four count probes require explicit result inspection rather than automatic validation: expect `262144`, `131072`, `16896`, and `16896` in query order. The default lifecycle retains caches, so the optimization remains exercised but its timing includes warm-cache effects; use an explicitly configured cold environment when a cold comparison is required. + The Rust runner subcommands are also useful for focused diagnostics: ```bash @@ -393,12 +509,13 @@ parquetbench/scanbench` as the read-bench tool against each target's data direct The workflow runs when an allowlisted repository admin comments `/query-regression` on a non-draft PR. It does not rerun on pushes, -ready-for-review, or reopen events. `/query-regression` runs the seven routine -default cases, including `promql_instant_last_row_9034`; -`/query-regression heavy` runs only the high-cardinality remote-write #7913 case. -PR runs build base/candidate once and use `--allow-large-fixture`. Manual -`workflow_dispatch` runs can pass `all`, `heavy`, one case path, or a -comma/whitespace-separated list of case paths, and can override refs. +ready-for-review, or reopen events. `/query-regression` runs the nine routine +default cases, including `sst_float_bss`, `promql_instant_last_row_9034`, and +`mito_prefilter_all_match`; `/query-regression heavy` runs only the +high-cardinality remote-write #7913 case. PR runs build base/candidate once and +use `--allow-large-fixture`. Manual `workflow_dispatch` runs can pass `all`, +`heavy`, one case path, or a comma/whitespace-separated list of case paths, and +can override refs. Comment admission is two workflows. `slash-command-dispatch.yml` uses [peter-evans/slash-command-dispatch](https://github.com/peter-evans/slash-command-dispatch) diff --git a/tests/perf/query_cases/mito_prefilter_all_match/case.toml b/tests/perf/query_cases/mito_prefilter_all_match/case.toml new file mode 100644 index 00000000000..9db69eb446a --- /dev/null +++ b/tests/perf/query_cases/mito_prefilter_all_match/case.toml @@ -0,0 +1,213 @@ +# Performance regression coverage for the Mito prefilter all-match optimization. +# +# The count probes are deliberately not automatic assertions: inspect their +# results manually. Each should return 262144, 131072, 16896, and 16896 +# respectively. + +[case] +name = "mito_prefilter_all_match" +description = "Mito prefilter all-match regression coverage with integer and float selective controls" + +[scenario] +kind = "direct_readable_sst" +seed = 9402 + +[[scenario.tables]] +database = "public" +name = "mito_prefilter_all_match" +engine = "mito" +append_mode = true +sst_format = "flat" +primary_key = ["tag_a", "tag_b"] +time_index = "ts" + +[[scenario.tables.columns]] +name = "tag_a" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 1, prefix = "tag_a" } + +[[scenario.tables.columns]] +name = "tag_b" +type = "STRING" +semantic = "tag" +distribution = { kind = "cardinality", values = 1, prefix = "tag_b" } + +[[scenario.tables.columns]] +name = "field_01" +type = "UINT64" +semantic = "field" + +[[scenario.tables.columns]] +name = "field_02" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 2.0 } + +[[scenario.tables.columns]] +name = "field_03" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 3.0 } + +[[scenario.tables.columns]] +name = "field_04" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 4.0 } + +[[scenario.tables.columns]] +name = "field_05" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 5.0 } + +[[scenario.tables.columns]] +name = "field_06" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 6.0 } + +[[scenario.tables.columns]] +name = "field_07" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 7.0 } + +[[scenario.tables.columns]] +name = "field_08" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 8.0 } + +[[scenario.tables.columns]] +name = "field_09" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 9.0 } + +[[scenario.tables.columns]] +name = "field_10" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 10.0 } + +[[scenario.tables.columns]] +name = "field_11" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 11.0 } + +[[scenario.tables.columns]] +name = "field_12" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 12.0 } + +[[scenario.tables.columns]] +name = "field_13" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 13.0 } + +[[scenario.tables.columns]] +name = "field_14" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 14.0 } + +[[scenario.tables.columns]] +name = "field_15" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 15.0 } + +[[scenario.tables.columns]] +name = "field_16" +type = "DOUBLE" +semantic = "field" +distribution = { kind = "deterministic_wave", min = 0.0, max = 16.0 } + +[[scenario.tables.columns]] +name = "ts" +type = "TIMESTAMP(9)" +semantic = "timestamp" + +[scenario.layout] +regions = 1 +sst_count = 1 +rows_per_sst = 262144 +row_group_size = 262144 +series_count = 1 +start_unix_nanos = 1704067200000000000 +step_nanos = 1000000000 +time_range_layout = "non_overlapping_per_sst" +series_layout = "timestamp_major" + +[[scenario.queries]] +name = "all_match_field_01" +kind = "sql" +query = "EXPLAIN ANALYZE VERBOSE SELECT * FROM mito_prefilter_all_match WHERE field_01 >= 0" +warmup = 2 +iterations = 10 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 30 + +[[scenario.queries]] +name = "all_match_field_01_and_field_02" +kind = "sql" +query = "EXPLAIN ANALYZE VERBOSE SELECT * FROM mito_prefilter_all_match WHERE field_01 >= 0 AND field_02 > 1.98" +warmup = 2 +iterations = 10 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 30 + +[[scenario.queries]] +name = "selective_control_field_01" +kind = "sql" +query = "EXPLAIN ANALYZE VERBOSE SELECT * FROM mito_prefilter_all_match WHERE field_01 >= 131072" +warmup = 2 +iterations = 10 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 30 + +[[scenario.queries]] +name = "selective_control_field_02" +kind = "sql" +query = "EXPLAIN ANALYZE VERBOSE SELECT * FROM mito_prefilter_all_match WHERE field_02 > 1.98" +warmup = 2 +iterations = 10 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 30 + +[[scenario.queries]] +name = "count_all_match_field_01" +kind = "sql" +query = "SELECT count(*) FROM mito_prefilter_all_match WHERE field_01 >= 0" +warmup = 0 +iterations = 1 + +[[scenario.queries]] +name = "count_selective_control_field_01" +kind = "sql" +query = "SELECT count(*) FROM mito_prefilter_all_match WHERE field_01 >= 131072" +warmup = 0 +iterations = 1 + +[[scenario.queries]] +name = "count_all_match_field_01_and_field_02" +kind = "sql" +query = "SELECT count(*) FROM mito_prefilter_all_match WHERE field_01 >= 0 AND field_02 > 1.98" +warmup = 0 +iterations = 1 + +[[scenario.queries]] +name = "count_selective_control_field_02" +kind = "sql" +query = "SELECT count(*) FROM mito_prefilter_all_match WHERE field_02 > 1.98" +warmup = 0 +iterations = 1 diff --git a/tests/perf/query_cases/promql_simple_range_allocation/case.toml b/tests/perf/query_cases/promql_simple_range_allocation/case.toml new file mode 100644 index 00000000000..95d00ad069e --- /dev/null +++ b/tests/perf/query_cases/promql_simple_range_allocation/case.toml @@ -0,0 +1,143 @@ +# Bounded PromQL simple-range allocation regression coverage. +# +# 128 series × 780 samples at a 15-second cadence = 99,840 rows. Linear +# values are `base + series_idx % 97 + local_ordinal`; the host0000 HTTP +# control therefore has last values 240, 480, and 720 at its hourly evaluations. + +[case] +name = "promql_simple_range_allocation" +description = "Bounded PromQL simple-range UDF allocation regression coverage" + +[scenario] +kind = "prom_remote_write_then_query" + +[scenario.remote_write] +database = "public" +metric = "promql_simple_range_allocation" +physical_table = "greptime_physical_table" +series_count = 128 +samples_per_series = 780 +sample_chunk_size = 195 +flush_every_sample_chunks = 1 +start_unix_millis = 1704067200000 +step_millis = 15000 +chunk_series_count = 128 +timeout_seconds = 600 +visibility_timeout_seconds = 300 + +[scenario.remote_write.value] +pattern = "linear" +base = 0 +step = 1 + +[scenario.remote_write.prom_store] +pending_rows_flush_interval = "1s" +max_batch_rows = 100000 + +[[scenario.queries]] +name = "selector_2h15s_control" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704070800, 1704078000, '15s') promql_simple_range_allocation{host=~'host.*'}" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +[[scenario.queries]] +name = "count_over_time_1m" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704070800, 1704078000, '15s') count_over_time(promql_simple_range_allocation{host=~'host.*'}[1m])" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +[[scenario.queries]] +name = "count_over_time_1h" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704070800, 1704078000, '15s') count_over_time(promql_simple_range_allocation{host=~'host.*'}[1h])" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +[[scenario.queries]] +name = "last_over_time_1m" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704070800, 1704078000, '15s') last_over_time(promql_simple_range_allocation{host=~'host.*'}[1m])" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +[[scenario.queries]] +name = "last_over_time_1h" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704070800, 1704078000, '15s') last_over_time(promql_simple_range_allocation{host=~'host.*'}[1h])" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +[[scenario.queries]] +name = "present_over_time_1m" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704070800, 1704078000, '15s') present_over_time(promql_simple_range_allocation{host=~'host.*'}[1m])" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +[[scenario.queries]] +name = "present_over_time_1h" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704070800, 1704078000, '15s') present_over_time(promql_simple_range_allocation{host=~'host.*'}[1h])" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +# This matcher is intentionally nonempty in the input but produces no output. +# ANALYZE remains valid while exercising the absent-over-time allocation path. +[[scenario.queries]] +name = "absent_over_time_1m_empty" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704070800, 1704078000, '15s') absent_over_time(promql_simple_range_allocation{host=~'host.*'}[1m])" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +[[scenario.queries]] +name = "absent_over_time_1h_empty" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704070800, 1704078000, '15s') absent_over_time(promql_simple_range_allocation{host=~'host.*'}[1h])" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +# Inspect the retained response manually: the runner does not assert numeric values. +# For host0000, last_over_time at the hourly evaluations should return 240, 480, +# and 720 respectively. +[[scenario.queries]] +name = "prom_http_last_over_time_1h_numeric_control" +kind = "prom_http" +query = "last_over_time(promql_simple_range_allocation{host=\"host0000\"}[1h])" +start = "1704070800" +end = "1704078000" +step = "1h" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 diff --git a/tests/perf/query_cases/promql_smoothing_copy/case.toml b/tests/perf/query_cases/promql_smoothing_copy/case.toml new file mode 100644 index 00000000000..187794dd1c4 --- /dev/null +++ b/tests/perf/query_cases/promql_smoothing_copy/case.toml @@ -0,0 +1,115 @@ +# Bounded nightly PromQL double-exponential-smoothing regression coverage. +# +# 128 series × 780 samples at 15-second cadence = 99,840 rows. Remote write +# ingests four 195-sample chunks and flushes each chunk before querying. + +[case] +name = "promql_smoothing_copy" +description = "Bounded nightly PromQL double-exponential-smoothing copy-removal regression" + +[scenario] +kind = "prom_remote_write_then_query" + +[scenario.remote_write] +database = "public" +metric = "promql_smoothing_copy" +physical_table = "greptime_physical_table" +series_count = 128 +samples_per_series = 780 +sample_chunk_size = 195 +flush_every_sample_chunks = 1 +start_unix_millis = 1_704_067_200_000 +step_millis = 15_000 +chunk_series_count = 128 +timeout_seconds = 180 +visibility_timeout_seconds = 120 + +[scenario.remote_write.value] +pattern = "linear" +base = 0 +step = 1 + +[scenario.remote_write.prom_store] +pending_rows_flush_interval = "1s" +max_batch_rows = 100000 + +# Selector-only scan control for unrelated PromQL/metric-engine variance. +[[scenario.queries]] +name = "selector_2h_15s_control" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704070800, 1704078000, '15s') promql_smoothing_copy{host=~'host.*'}" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +[[scenario.queries]] +name = "double_exponential_smoothing_1m_15s" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704070800, 1704078000, '15s') double_exponential_smoothing(promql_smoothing_copy{host=~'host.*'}[1m], 0.5, 0.1)" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +[[scenario.queries]] +name = "double_exponential_smoothing_5m_15s" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704070800, 1704078000, '15s') double_exponential_smoothing(promql_smoothing_copy{host=~'host.*'}[5m], 0.5, 0.1)" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +# Primary copy-removal stress case: 481 highly-overlapping 240-sample windows. +[[scenario.queries]] +name = "double_exponential_smoothing_1h_15s" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704070800, 1704078000, '15s') double_exponential_smoothing(promql_smoothing_copy{host=~'host.*'}[1h], 0.5, 0.1)" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +[[scenario.queries]] +name = "double_exponential_smoothing_1h_5m_overlap" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704070800, 1704078000, '5m') double_exponential_smoothing(promql_smoothing_copy{host=~'host.*'}[1h], 0.5, 0.1)" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +# Three disjoint 240-sample windows provide a non-overlapping control. +[[scenario.queries]] +name = "double_exponential_smoothing_1h_1h_nonoverlap" +kind = "tql" +query = "TQL ANALYZE VERBOSE (1704070800, 1704078000, '1h') double_exponential_smoothing(promql_smoothing_copy{host=~'host.*'}[1h], 0.5, 0.1)" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 + +# Inspect retained responses manually; the runner does not assert numeric values. +# The three left-open hourly windows contain host0000 values 1..240, 241..480, +# and 481..720. Initial trend is 1 and remains unchanged on the first iteration; +# subsequent trend updates preserve it. With sf=0.5 and tf=0.1, the expected +# outputs are 240, 480, and 720. +[[scenario.queries]] +name = "double_exponential_smoothing_1h_host0000_numeric_control" +kind = "prom_http" +query = "double_exponential_smoothing(promql_smoothing_copy{host=\"host0000\"}[1h], 0.5, 0.1)" +start = "1704070800" +end = "1704078000" +step = "1h" +warmup = 3 +iterations = 9 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 10 diff --git a/tests/perf/query_cases/sst_float_bss/RESULTS.md b/tests/perf/query_cases/sst_float_bss/RESULTS.md new file mode 100644 index 00000000000..ba932d0e42a --- /dev/null +++ b/tests/perf/query_cases/sst_float_bss/RESULTS.md @@ -0,0 +1,145 @@ +# SST float BSS: mixed-data post-flush evidence + +This report replaces the earlier unique-integer experiment as the primary evidence +for `sst_float_bss`. The detailed historical numbers remain in Git history; they +apply only to that superseded workload and must not be used for the current mixed +case. + +> **Timing measurements are confounded.** After these runs, the user reported +> that other agents' builds saturated all 20 logical CPUs during measurement. +> No host-wide idle/load gate was enforced. The latency tables below are retained +> as raw observations only: neither improvements nor threshold failures establish +> BSS performance. SST byte counts and cross-target stored-row checks remain +> valid for the inspected dataset. Stable-host timing verification is pending. + +## Case and data shape + +- Generator/case revision: `149b49de5f3`. +- Base and candidate used the same release `greptime` binary. The only physical + table difference was `experimental_sst_float_field_encoding`: `default` versus + `byte_stream_split`; base used dictionary encoding and candidate used BSS + without dictionaries. This is not a dictionary-policy-held-constant comparison. +- Each of three fresh post-flush runs wrote 1,000 series × 4,320 samples at + 60-second intervals: 4,320,000 rows per target. Three explicit flushes produced + six SSTs per target: three × 1,105,920 rows and three × 334,080 rows, with 45 + `greptime_value` row groups total. This is observed layout, not a layout + guarantee. +- `bounded_mixed` is a synthetic, uncalibrated design: a series hash selects its + baseline and span band, and 60-sample anchor interpolation produces bounded + fluctuation. It does not represent a measured production distribution. +- Exactly 950 equal-length series are integer-only and 50 are fractional-only, + yielding 4,104,000 integer rows and 216,000 fractional rows per target. The + equal-length 1,000-series design deliberately makes both the series and row + shares 95%/5%. +- The 95% design is not an empirical claim. A [local preview survey project](http://192.168.50.85:8765/view/prometheus-metrics-186-project-preliminary) + observed 94.79% integral **samples**; it does not establish a global rate, + global uniqueness, or a 95% series distribution. + +## Initial SST inspection and value verification + +Inspection is limited to `data/greptime/public`; bytes exclude WAL, indexes, +metadata, and internal/private-table storage. Each of the three runs had the +same totals: + +| Measure | Default | BSS | BSS change | +| --- | ---: | ---: | ---: | +| Selected SST bytes | 20,443,131 | 8,884,743 | -56.5392% | +| `greptime_value` compressed bytes | 20,201,188 | 8,643,964 | -57.2106% | +| SST files / rows / value row groups | 6 / 4,320,000 / 45 | 6 / 4,320,000 / 45 | — | + +Footer inspection found `PLAIN`/`RLE`/`RLE_DICTIONARY` for the default value +column and `RLE`/`BYTE_STREAM_SPLIT` for BSS (both ZSTD level 1). The configured +storage threshold remains a 5% reduction; this repeated SST-size result passes +that threshold. It is post-flush evidence only, not a compaction result or a +general float-compression claim. + +For every base/candidate group in all three runs, rows were compared by the same +primary key and timestamp and their Float64 bits were exact matches: 950 +integer-only and 50 fractional-only series, 4,104,000 and 216,000 rows +respectively. The six canonical sorted-row SHA-256 results are identical: + +``` +48f87d7d01bc8a0ce9a9cdb3033f322e787784fcae61f1417ac25b05c4a40375 +``` + +This verifies exact cross-target stored rows, not independent reconstruction of +the generator output. All 1,000 series had both rises and falls, and every one +of the 36 inspected SSTs contained both integer-valued and fractional values; +these were not stored as two disjoint file sets. + +## Warmed endpoint SQL + +Each query used three warmups followed by 15 measured endpoint requests. Values +are median/p95 milliseconds; delta is `(BSS median - default median) / default +median × 100%`. The guardrail remains a maximum 25% candidate median regression and +was not changed for these runs. + +| Run | Query | Default median / p95 | BSS median / p95 | Delta | Guardrail | +| --- | --- | ---: | ---: | ---: | --- | +| mixed-1 | `sum_all_values` | 27.13 / 37.38 | 25.45 / 38.99 | -6.21% | pass | +| mixed-1 | `hourly_sum_values` | 118.78 / 136.63 | 130.93 / 185.26 | +10.23% | pass | +| mixed-2 | `sum_all_values` | 24.02 / 35.02 | 36.34 / 56.01 | +51.31% | **fail** | +| mixed-2 | `hourly_sum_values` | 95.58 / 120.40 | 85.45 / 148.03 | -10.60% | pass | +| mixed-3 | `sum_all_values` | 38.25 / 55.41 | 19.68 / 29.49 | -48.55% | pass | +| mixed-3 | `hourly_sum_values` | 56.77 / 148.29 | 121.87 / 143.16 | +114.66% | **fail** | + +The two recorded threshold failures remain in the artifacts, but the saturated +shared host prevents attributing them to BSS. Passing entries likewise do not +establish performance acceptance. No CPU isolation or cache flush was used. + +## Stopped-directory warm readers + +For each of the three datasets, eight outer rounds alternate base/candidate order. +Each command runs eight iterations, discards the first, and takes the median of +iterations 2–8. Parquetbench reads every SST: sum the six per-file medians, then +take the median of the eight round totals. Scanbench scans the whole region and +reports the outer median of its eight warm medians. These are distinct metrics, +not interchangeable whole-query latencies. Deltas are ratios of separately +aggregated medians. Value projection selects `greptime_value`; all-column +parquetbench returns five physical columns, while scanbench uses `{}` (all +region columns). Projection order is fixed; target order alternates. + +| Dataset | Projection | Reader | Base ms | BSS ms | Change | +| --- | --- | --- | ---: | ---: | ---: | +| mixed-1 | value | parquetbench | 95.957 | 68.442 | -28.67% | +| mixed-1 | value | scanbench | 144.293 | 115.421 | -20.01% | +| mixed-1 | allcolumns | parquetbench | 176.620 | 103.761 | -41.25% | +| mixed-1 | allcolumns | scanbench | 200.881 | 202.660 | +0.89% | +| mixed-2 | value | parquetbench | 112.151 | 74.275 | -33.77% | +| mixed-2 | value | scanbench | 109.824 | 78.189 | -28.80% | +| mixed-2 | allcolumns | parquetbench | 141.944 | 102.681 | -27.66% | +| mixed-2 | allcolumns | scanbench | 211.414 | 163.261 | -22.78% | +| mixed-3 | value | parquetbench | 121.827 | 67.565 | -44.54% | +| mixed-3 | value | scanbench | 100.968 | 79.878 | -20.89% | +| mixed-3 | allcolumns | parquetbench | 113.984 | 94.361 | -17.22% | +| mixed-3 | allcolumns | scanbench | 193.101 | 164.175 | -14.98% | + +All 672 commands / 5,376 raw iterations returned the expected rows; timing, +projection configurations, parquet schemas and order metadata were checked. +Run 1 all-column scanbench differed by **+0.89%**, negligible relative to the +observed variation, not evidence of a regression. Its round medians ranged from +167.197–663.465 ms (base) and 133.339–557.352 ms (BSS). There is substantial +variability; these measurements do not establish a stable gain on every read +path. The reported concurrent CPU saturation affects apparent improvements as +well as slowdowns; no causal performance conclusion is drawn from these runs. + +## Scope and recorded checks + +This case compares default and BSS post-flush data only. It adds no benchmark +framework and no public artifact upload. Historical unique-integer results and +other mixed counter/gauge studies have different data and methods; they are +context only and do not apply to this workload. + +Recorded checks for this change set: 37 targeted Rust tests, 20 Python tooling +tests, and 15 plans passed. Full-workspace tests were not run. + +Local artifacts are under +`/mnt/nvme_rust/rust-targets/metric-bss-perf/experiments/mixed-{1,2,3}/`, +including `value-verification.json`, `verification.txt`, and each run's +`sst_float_bss/query-regression-report.json`. + +Warm artifacts are in sibling `warm-mixed-{1,2,3}/` directories; +`mixed-warm-comparison.json` contains the table and per-round ranges. +`mixed-file-composition.json` records integer/fractional counts for each SST. +The shared release `greptime` SHA-256 is +`2ce0ab1670cd4bdb0f60c87af3e1f494ee66c3c6ab83bce943a073638514eb2b`. diff --git a/tests/perf/query_cases/sst_float_bss/case.toml b/tests/perf/query_cases/sst_float_bss/case.toml new file mode 100644 index 00000000000..272e0eb4bc4 --- /dev/null +++ b/tests/perf/query_cases/sst_float_bss/case.toml @@ -0,0 +1,83 @@ +# Default all-group default-versus-BSS storage and warmed-query comparison. + +[case] +name = "sst_float_bss" +description = "Compare default and byte-stream-split floating-point SST bytes and warmed value-scanning query latency with bounded 95/5 per-series DOUBLE values" + +[scenario] +kind = "prom_remote_write_then_query" + +[scenario.remote_write] +database = "public" +metric = "sst_float_bss" +physical_table = "sst_float_bss_physical" +series_count = 1000 +samples_per_series = 4320 +sample_chunk_size = 1440 +flush_every_sample_chunks = 1 +start_unix_millis = 1_704_067_200_000 # 2024-01-01T00:00:00Z +step_millis = 60000 +chunk_series_count = 128 +timeout_seconds = 600 +visibility_timeout_seconds = 300 + +# Omit the default encoding from base setup because older binaries do not recognize it. +# The high shared TWCS trigger avoids compaction changing the three-flush layout. +base_setup_sql = [ + "CREATE TABLE sst_float_bss_physical (greptime_timestamp TIMESTAMP TIME INDEX, greptime_value DOUBLE) ENGINE=metric WITH ('physical_metric_table'='', 'compaction.twcs.trigger_file_num'='100')", +] +candidate_setup_sql = [ + "CREATE TABLE sst_float_bss_physical (greptime_timestamp TIMESTAMP TIME INDEX, greptime_value DOUBLE) ENGINE=metric WITH ('physical_metric_table'='', 'experimental_sst_float_field_encoding'='byte_stream_split', 'compaction.twcs.trigger_file_num'='100')", +] + +[scenario.remote_write.value] +# Synthetic, not empirical: 95% integral and 5% nonintegral series vary +# smoothly within deterministic per-series ranges; series are not a global monotonic counter. +pattern = "bounded_mixed" +base = 0 +step = 1 +cardinality = 1000 +seed = 42 +mixed_every = 20 +run_length = 60 + +[scenario.remote_write.prom_store] +pending_rows_flush_interval = "1s" +max_batch_rows = 100000 + +[scenario.remote_write.storage] +root_suffix = "data/greptime/public" +column = "greptime_value" +min_files = 1 +min_files_with_column = 1 +max_candidate_total_file_size_regression_pct = -5.0 + +[scenario.remote_write.read_bench] +enabled = true +parquetbench = true +scanbench = true +iterations = 7 +projection = ["greptime_value"] +parquet_reader = "direct" +scan_scanner = "seq" +parallelism = 1 + +[[scenario.queries]] +name = "sum_all_values" +kind = "sql" +query = "SELECT sum(greptime_value) FROM sst_float_bss" +warmup = 3 +iterations = 15 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 25 + +[[scenario.queries]] +name = "hourly_sum_values" +kind = "sql" +query = "SELECT date_bin(INTERVAL '1 hour', greptime_timestamp) AS time_window, sum(greptime_value) FROM sst_float_bss WHERE greptime_timestamp >= TIMESTAMP '2024-01-01 00:00:00' AND greptime_timestamp < TIMESTAMP '2024-01-04 00:00:00' GROUP BY time_window ORDER BY time_window" +warmup = 3 +iterations = 15 + +[scenario.queries.thresholds] +max_candidate_latency_regression_pct = 25 diff --git a/tests/perf/test_query_regression_case_selection.py b/tests/perf/test_query_regression_case_selection.py index 84b3cd44fa7..04a542e23ad 100644 --- a/tests/perf/test_query_regression_case_selection.py +++ b/tests/perf/test_query_regression_case_selection.py @@ -16,6 +16,7 @@ """Coverage for query-regression case group selection.""" import importlib.util +import re import sys import unittest from pathlib import Path @@ -35,15 +36,21 @@ class QueryRegressionCaseSelectionTest(unittest.TestCase): runner.split_cases(["all"]), [ "tests/perf/query_cases/smoke_direct_sst/case.toml", + "tests/perf/query_cases/sst_float_bss/case.toml", "tests/perf/query_cases/prom_remote_write_seeded_random/case.toml", "tests/perf/query_cases/prom_remote_write_run_heavy/case.toml", "tests/perf/query_cases/prom_remote_write_mixed_every/case.toml", "tests/perf/query_cases/prom_remote_write_integer_counter/case.toml", "tests/perf/query_cases/promql_range_boundary/case.toml", "tests/perf/query_cases/promql_instant_last_row_9034/case.toml", + "tests/perf/query_cases/mito_prefilter_all_match/case.toml", ], ) + def test_implicit_selection_has_nine_routine_cases(self) -> None: + self.assertEqual(len(runner.DEFAULT_CASES), 9) + self.assertEqual(runner.split_cases([]), runner.DEFAULT_CASES) + def test_heavy_selects_only_remote_write_7913(self) -> None: self.assertEqual( runner.split_cases(["heavy"]), @@ -54,6 +61,23 @@ class QueryRegressionCaseSelectionTest(unittest.TestCase): case = "tests/perf/query_cases/sql_topk_order_by/case.toml" self.assertEqual(runner.split_cases([case]), [case]) + def test_sst_float_bss_omits_base_encoding_for_older_binaries(self) -> None: + case_path = Path(__file__).parent / "query_cases/sst_float_bss/case.toml" + setup_sql = dict(re.findall( + r'(base|candidate)_setup_sql = \[\s*"([^"\n]+)"', + case_path.read_text(), + )) + base_setup = setup_sql["base"] + candidate_setup = setup_sql["candidate"] + candidate_option = "'experimental_sst_float_field_encoding'='byte_stream_split'" + + self.assertNotIn("experimental_sst_float_field_encoding", base_setup) + self.assertIn(candidate_option, candidate_setup) + self.assertEqual( + base_setup, + candidate_setup.replace(f", {candidate_option}", ""), + ) + if __name__ == "__main__": unittest.main()